code: def BatchNorm(
      axis=(0, 1, 2),
      epsilon=1e-5,
      center=True,
      scale=True,
      beta_init=zeros,
      gamma_init=ones,
  ):
      """Layer construction function for a batch normalization layer."""
      _beta_init = (
          lambda rng, shape: beta_init(rng, shape) if center else ()
      )
      _gamma_init = (
          lambda rng, shape: gamma_init(rng, shape) if scale else ()
      )
      axis = (axis,) if jnp.isscalar(axis) else axis
  
      def init_fun(rng, input_shape):
          shape = tuple(
              d for i, d in enumerate(input_shape) if i not in axis
          )
          k1, k2 = random.split(rng)
          beta, gamma = _beta_init(k1, shape), _gamma_init(k2, shape)
          return input_shape, (beta, gamma)
  
      def apply_fun(params, x, **kwargs):
          beta, gamma = params
          # TODO(phawkins): jnp.expand_dims should accept an axis tuple.
          # (https://github.com/numpy/numpy/issues/12290)
          ed = tuple(
              None if i in axis else slice(None)
              for i in range(jnp.ndim(x))
          )
          z = standardize(x, axis, epsilon=epsilon)
          if center and scale:
              return gamma[ed] * z + beta[ed]
          if center:
              return z + beta[ed]
          if scale:
              return gamma[ed] * z
          return z
  
      return init_fun, apply_fun
  
imports: import functools
  import operator
  from jax import lax
  from jax import random
  import jax.numpy
  from jax.nn import relu, log_softmax, softmax, softplus, sigmoid, elu, leaky_relu, selu, gelu, standardize
  from jax.nn.initializers import glorot_normal, normal, ones, zeros
label_desc: Use image node_id 4 for calling a global function with example usage: BatchNorm(axis, epsilon, center, scale, beta_init, gamma_init) and returns: init_fun,  apply_fun, input_shape,  beta,  gamma, z, unknown, unknown, unknown