code: def test_invalid_axis(self, axis):
      with self.assertRaisesRegex(
          ValueError,
          "`axis` should be an int, slice or iterable of ints.",
      ):
          rms_norm.RMSNorm(axis)
  
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 6 for calling the RMSNormTest obj's underlying member method code with example usage: obj.test_invalid_axis(axis) without return types