code: def test_param_axis_sets_param_shape(self, param_axis, param_shape):
      ln = rms_norm.RMSNorm(axis=-1, param_axis=param_axis)
      ln(jnp.ones([3, 4, 5, 6]))
      self.assertEqual(
          ln.params_dict()["rms_norm/scale"].shape, param_shape
      )
  
imports: from absl.testing import absltest
  from absl.testing import parameterized
  from haiku._src import initializers
  from haiku._src import rms_norm
  from haiku._src import test_utils
  from haiku._src import transform
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 9 for calling the RMSNormTest obj's underlying member method code with example usage: obj.test_param_axis_sets_param_shape(param_axis, param_shape) without return types