code: def test_spectral_norm(getkey):
      def λ1():
          u, v = state.get(spectral.uv_index)
          σ = jnp.einsum("i,ij,j->", u, spectral.layer.weight, v)
          _, s, _ = jnp.linalg.svd(
              spectral.layer.weight / σ
          )  # pyright: ignore
          return s[0]
  
      x = jrandom.normal(getkey(), (5,))
      spectral = eqx.nn.SpectralNorm(
          eqx.nn.Linear(5, 6, key=getkey()), "weight", key=getkey()
      )
      state = eqx.nn.State(spectral)
      for _ in range(100):
          _, state = spectral(x, state)
      assert jnp.allclose(λ1(), 1)
  
      # "gradient descent"
      spectral = eqx.tree_at(
          lambda s: s.layer.weight, spectral, spectral.layer.weight + 1
      )
      assert not jnp.allclose(λ1(), 1)
      for _ in range(100):
          _, state = spectral(x, state)
      assert jnp.allclose(λ1(), 1)
  
      # Test not updated at inference time
      spectral = eqx.tree_at(
          lambda s: s.layer.weight, spectral, spectral.layer.weight + 1
      )
      spectral = eqx.nn.inference_mode(spectral, value=True)
      assert not jnp.allclose(λ1(), 1)
      for _ in range(100):
          _, state = spectral(x, state)
      assert not jnp.allclose(λ1(), 1)
  
      # Test >2 dimensional input
  
      x = jrandom.normal(getkey(), (5, 8, 8, 8))
      conv = eqx.nn.Conv3d(5, 4, 3, key=getkey())
      spectral = eqx.nn.SpectralNorm(conv, "weight", key=getkey())
      state = eqx.nn.State(spectral)
      out, _ = spectral(x, state)
      assert out.shape == (4, 6, 6, 6)
  
imports: import warnings
  from typing import Union
  import equinox
  import jax
  import jax.nn
  import jax.numpy
  import jax.random
  import pytest
label_desc: Use image node_id 30 for calling a global function with example usage: test_spectral_norm(getkey) and returns: s