From 61382083f398673371a9b36ce085617e6f9777d7 Mon Sep 17 00:00:00 2001
From: Jake VanderPlas <vanderplas@google.com>
Date: Mon, 11 Nov 2024 09:37:42 -0800
Subject: [PATCH] No public description

PiperOrigin-RevId: 695375318
---
 vmoe/initialization/rules.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/vmoe/initialization/rules.py b/vmoe/initialization/rules.py
index 3e1825c..cdb6373 100644
--- 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):