code: def test_batch_norm(getkey):
      x0 = jrandom.uniform(getkey(), (5,))
      x1 = jrandom.uniform(getkey(), (10, 5))
      x2 = jrandom.uniform(getkey(), (10, 5, 6))
      x3 = jrandom.uniform(getkey(), (10, 5, 7, 8))
  
      # Test that it works with a single vmap'd axis_name
  
      bn = eqx.nn.BatchNorm(5, "batch")
      state = eqx.nn.State(bn)
      vbn = jax.vmap(
          bn, axis_name="batch", in_axes=(0, None), out_axes=(0, None)
      )
  
      for x in (x1, x2, x3):
          out, state = vbn(x, state)
          assert out.shape == x.shape
          running_mean, running_var = state.get(bn.state_index)
          assert running_mean.shape == (5,)
          assert running_var.shape == (5,)
  
      # Test that it fails without any vmap'd axis_name
  
      with pytest.raises(NameError):
          jax.vmap(bn, in_axes=(0, None), out_axes=(0, None))(x1, state)
      state = eqx.nn.State(bn)
  
      with pytest.raises(NameError):
          bn(x0, state)
      state = eqx.nn.State(bn)
  
      # Test that it vmaps with other vmaps without axis_name
  
      out, state = jax.vmap(
          jax.vmap(
              bn,
              axis_name="batch",
              in_axes=(1, None),
              out_axes=(1, None),
          ),
          in_axes=(0, None),
      )(x2, state)
      assert out.shape == x2.shape
      running_mean, running_var = state.get(bn.state_index)
      assert running_mean.shape == (10, 5)
      assert running_var.shape == (10, 5)
  
      # Test that it handles multiple axis_names
  
      vvbn = eqx.nn.BatchNorm(6, ("batch1", "batch2"))
      vvstate = eqx.nn.State(vvbn)
      for axis_name in ("batch1", "batch2"):
          vvbn = jax.vmap(
              vvbn,
              axis_name=axis_name,
              in_axes=(0, None),
              out_axes=(0, None),
          )
      out, out_vvstate = vvbn(x2, vvstate)
      assert out.shape == x2.shape
      running_mean, running_var = out_vvstate.get(vvbn.state_index)
      assert running_mean.shape == (6,)
      assert running_var.shape == (6,)
  
      # Test that it normalises
  
      x1alt = jrandom.normal(
          jrandom.PRNGKey(5678), (10, 5)
      )  # avoid flakey test
      bn = eqx.nn.BatchNorm(5, "batch", channelwise_affine=False)
      state = eqx.nn.State(bn)
      vbn = jax.vmap(
          bn, axis_name="batch", in_axes=(0, None), out_axes=(0, None)
      )
      out, state = vbn(x1alt, state)
      true_out = (
          x1alt - jnp.mean(x1alt, axis=0, keepdims=True)
      ) / jnp.sqrt(jnp.var(x1alt, axis=0, keepdims=True) + 1e-5)
      assert jnp.allclose(out, true_out)
  
      # Test that the statistics update during training
      out, state = vbn(x1, state)
      running_mean, running_var = state.get(bn.state_index)
      out, state = vbn(3 * x1 + 10, state)
      running_mean2, running_var2 = state.get(bn.state_index)
      assert not jnp.allclose(running_mean, running_mean2)
      assert not jnp.allclose(running_var, running_var2)
  
      # Test that the statistics don't update at inference
  
      ibn = eqx.nn.inference_mode(bn, value=True)
      vibn = jax.vmap(
          ibn, axis_name="batch", in_axes=(0, None), out_axes=(0, None)
      )
      out, state = vibn(4 * x1 + 20, state)
      running_mean3, running_var3 = state.get(bn.state_index)
      assert jnp.array_equal(running_mean2, running_mean3)
      assert jnp.array_equal(running_var2, running_var3)
  
      # Test that we can differentiate through it
  
      @jax.grad
      def f(x):
          out, _ = vbn(x, state)
          return jnp.sum(out)
  
      f(jrandom.normal(getkey(), (1, 5)))
  
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 29 for calling a global function with example usage: test_batch_norm(getkey) and returns: jnp