code: def test_multihead_sow_attention_weights(self):
      class Model(nnx.Module):
          attention_kwargs: dict
  
          def __init__(self, attention_kwargs, rng):
              self.attention_layers = [
                  nnx.MultiHeadAttention(**attention_kwargs, rngs=rng)
                  for i in range(3)
              ]
  
          def __call__(self, x, sow_weights=False):
              x = self.attention_layers[0](x, sow_weights=sow_weights)
              x = self.attention_layers[1](x)
              x = self.attention_layers[2](x, sow_weights=sow_weights)
              return x
  
      rng = nnx.Rngs(0)
      x = jnp.ones((4, 6, 8))
  
      module = Model(
          dict(
              in_features=8,
              num_heads=8,
              kernel_init=nnx.initializers.ones_init(),
              bias_init=nnx.initializers.zeros_init(),
              deterministic=False,
          ),
          rng,
      )
  
      with nnx.flags(decode=False):
          _ = module(x, True)
      intermediates = module.pop(nnx.Intermediate)
      assert intermediates["attention_layers/0/attention_weights"][
          0
      ].shape == (
          4,
          8,
          6,
          6,
      )
      assert "attention_layers/1/attention_weights" not in intermediates
      assert intermediates["attention_layers/2/attention_weights"][
          0
      ].shape == (
          4,
          8,
          6,
          6,
      )
  
      with nnx.flags(decode=False):
          _ = module(x)
      intermediates = module.pop(nnx.Intermediate)
      assert not intermediates
  
imports: import jax, jax.numpy
  from flax.experimental import nnx
label_desc: Use image node_id 2 for calling the TestMultiHeadAttention obj's underlying member method code with example usage: obj.test_multihead_sow_attention_weights() and returns: x