code: def test_3d_tensor(self):
      sn = spectral_norm.SpectralNorm()
      input_3d = (4.0 * jnp.eye(8, 8))[None, :, :]
      sn(input_3d)
      with self.assertRaisesRegex(ValueError, "Input is 3D but"):
          sn(input_3d, error_on_non_matrix=True)
  
file_path: dm-haiku/haiku/_src/spectral_norm_test.py
imports: from absl.testing import absltest
  from haiku._src import basic
  from haiku._src import spectral_norm
  from haiku._src import test_utils
  from haiku._src import transform
  import jax.numpy
  import jax.random
  import numpy
label_desc: Use image node_id 3 for calling the SpectralNormTest obj's underlying member method code with example usage: obj.test_3d_tensor() without return types