Allow init_from_checkpoint to accept a list of pairs, so as to enable initialization of multiple variables in the graph from the same variable in the checkpoint.

nshazeer · Mesh TensorFlow Team · commit 1f6929e23436 · 2021-03-10T09:08:18.000-08:00
PiperOrigin-RevId: 361920642
diff --git a/mesh_tensorflow/transformer/transformer.py b/mesh_tensorflow/transformer/transformer.py
@@ -45,6 +45,7 @@
 
 import json
 import math
+import re
 
 import gin
 import mesh_tensorflow as mtf
