From 4c3c3765616a16122ac0cf98d747a284d3def1f2 Mon Sep 17 00:00:00 2001 From: Hyun Noh Date: Mon, 10 Aug 2026 16:13:57 -0700 Subject: [PATCH] Fix weight dictionary lookup to handle nn.share_scope. PiperOrigin-RevId: 962424478 --- gemma/gm/nn/gemma4/_layers.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) 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