No public description

PiperOrigin-RevId: 633764452
google-research · May 15, 2024 · 3cbb975 · 3cbb975
1 parent da5922d
commit 3cbb975
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/vmoe/train/trainer.py b/vmoe/train/trainer.py
@@ -634,7 +634,7 @@ def mix(x):
     a = alpha.reshape(alpha.shape + (1,) * (x.ndim - batch_ndim))
     return sum(a[i] * jnp.roll(x, -i, axis=roll_axis) for i in range(shape[-1]))
   arrays = list(map(mix, arrays))
-  return jax.tree_unflatten(treedef, arrays)
+  return jax.tree.unflatten(treedef, arrays)
 
 
 def override_base_config(