No public description

PiperOrigin-RevId: 695375318
google-research · Nov 11, 2024 · 6138208 · 6138208
1 parent 2826f11
commit 6138208
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/vmoe/initialization/rules.py b/vmoe/initialization/rules.py
@@ -263,7 +263,7 @@ class ReshapeTransformation(Transformation):
   shape: Tuple[int, ...] = flax.struct.field(pytree_node=False)
 
   def __call__(self) -> Array:
-    return jnp.reshape(self.array, newshape=self.shape)
+    return jnp.reshape(self.array, shape=self.shape)
 
 
 class SqueezeTransformation(Transformation):