code: def test_basic(self):
      module = nnx.MultiHeadAttention(
          num_heads=2,
          in_features=3,
          qkv_features=6,
          out_features=6,
          rngs=nnx.Rngs(0),
      )
      y = module(jnp.ones((1, 7, 3)), decode=False)
      assert y.shape == (1, 7, 6)
  
imports: import jax, jax.numpy
  from flax.experimental import nnx
label_desc: Use image node_id 1 for calling the TestMultiHeadAttention obj's underlying member method code with example usage: obj.test_basic() without return types