diff --git a/gemma/gm/nn/gemma4/_layers.py b/gemma/gm/nn/gemma4/_layers.py
index 59e77988..10c876ba 100644
--- a/gemma/gm/nn/gemma4/_layers.py
+++ b/gemma/gm/nn/gemma4/_layers.py
@@ -37,11 +37,10 @@ def __call__(self, eqn: str, x: jax.Array) -> jax.Array:
         self.dtype if self.dtype is not None else None,
     )
     # Workaround for behavior with nn.share_scope in parent modules:
-    # self.param might return a dict {'w': tensor} instead of the bare tensor.
-    # The key appears to be the default class weight_name 'w'.
-    if isinstance(w, dict):
-      if 'w' in w:
-        w = w['w']
+    # self.param might return a dict {self.weight_name: tensor} instead of the
+    # bare tensor.
+    if isinstance(w, dict) and self.weight_name in w:
+      w = w[self.weight_name]
     if self.w_scale is not None:
       w *= self.w_scale
     return jnp.einsum(eqn, x, w)
@@ -64,9 +63,8 @@ def __call__(self, eqn: str, x: jax.Array) -> jax.Array:
         self.shape,
         self.dtype if self.dtype is not None else None,
     )
-    if isinstance(w, dict):
-      if 'w' in w:
-        w = w['w']
+    if isinstance(w, dict) and self.weight_name in w:
+      w = w[self.weight_name]
     if self.w_scale is not None:
       w *= self.w_scale