code: def _normalize(
      x: Array,
      mean: Array,
      var: Array,
      scale: tp.Optional[Array],
      bias: tp.Optional[Array],
      reduction_axes: Axes,
      feature_axes: Axes,
      dtype: Dtype,
      epsilon: float,
  ):
      """ "Normalizes the input of a normalization layer and optionally applies a learned scale and bias.
  
      Arguments:
        x: The input.
        mean: Mean to use for normalization.
        var: Variance to use for normalization.
        reduction_axes: The axes in ``x`` to reduce.
        feature_axes: Axes containing features. A separate bias and scale is learned
          for each specified feature.
        dtype: The dtype of the result (default: infer from input and params).
        epsilon: Normalization epsilon.
  
      Returns:
        The normalized input.
      """
      reduction_axes = _canonicalize_axes(x.ndim, reduction_axes)
      feature_axes = _canonicalize_axes(x.ndim, feature_axes)
      stats_shape = list(x.shape)
      for axis in reduction_axes:
          stats_shape[axis] = 1
      mean = mean.reshape(stats_shape)
      var = var.reshape(stats_shape)
      feature_shape = [1] * x.ndim
      reduced_feature_shape = []
      for ax in feature_axes:
          feature_shape[ax] = x.shape[ax]
          reduced_feature_shape.append(x.shape[ax])
      y = x - mean
      mul = lax.rsqrt(var + epsilon)
      args = [x]
      if scale is not None:
          scale = scale.reshape(feature_shape)
          mul *= scale
          args.append(scale)
      y *= mul
      if bias is not None:
          bias = bias.reshape(feature_shape)
          y += bias
          args.append(bias)
      dtype = dtypes.canonicalize_dtype(*args, dtype=dtype)
      return jnp.asarray(y, dtype)
  
file_path: flax/flax/experimental/nnx/nnx/nn/normalization.py
imports: import typing
  import jax
  import jax.numpy
  from jax import lax
  from flax.experimental import nnx
  from flax.experimental.nnx.nnx import flaglib, rnglib
  from flax.experimental.nnx.nnx.module import Module, first_from
  from flax.experimental.nnx.nnx.nn import dtypes, initializers
label_desc: Use image node_id 4 for calling a global function with example usage: _normalize(x, mean, var, scale, bias, reduction_axes, feature_axes, dtype, epsilon) and returns: jnp