code: def _normalize(
      mdl: Module,
      x: Array,
      mean: Array,
      var: Array,
      reduction_axes: Axes,
      feature_axes: Axes,
      dtype: Dtype,
      param_dtype: Dtype,
      epsilon: float,
      use_bias: bool,
      use_scale: bool,
      bias_init: Callable[[PRNGKey, Shape, Dtype], Array],
      scale_init: Callable[[PRNGKey, Shape, Dtype], Array],
  ):
      """Normalizes the input of a normalization layer and optionally applies a learned scale and bias.
  
      Arguments:
        mdl: Module to apply the normalization in (normalization params will reside
          in this module).
        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).
        param_dtype: The dtype of the parameters.
        epsilon: Normalization epsilon.
        use_bias: If true, add a bias term to the output.
        use_scale: If true, scale the output.
        bias_init: Initialization function for the bias term.
        scale_init: Initialization function for the scaling function.
  
      Returns:
        The normalized input.
      """
      reduction_axes = _canonicalize_axes(x.ndim, reduction_axes)
      feature_axes = _canonicalize_axes(x.ndim, feature_axes)
      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])
  
      mean = jnp.expand_dims(mean, reduction_axes)
      var = jnp.expand_dims(var, reduction_axes)
      y = x - mean
      mul = lax.rsqrt(var + epsilon)
      args = [x]
      if use_scale:
          scale = mdl.param(
              "scale", scale_init, reduced_feature_shape, param_dtype
          ).reshape(feature_shape)
          mul *= scale
          args.append(scale)
      y *= mul
      if use_bias:
          bias = mdl.param(
              "bias", bias_init, reduced_feature_shape, param_dtype
          ).reshape(feature_shape)
          y += bias
          args.append(bias)
      dtype = dtypes.canonicalize_dtype(*args, dtype=dtype)
      return jnp.asarray(y, dtype)
  
imports: import dataclasses
  import functools
  from typing import Any, Callable, Iterable, Optional, Sequence, Tuple, Union
  import jax
  import jax.numpy
  from jax import lax
  from jax.nn import initializers
  from flax.linen import dtypes, module, transforms
label_desc: Use image node_id 4 for calling a global function with example usage: _normalize(mdl, x, mean, var, reduction_axes, feature_axes, dtype, param_dtype, epsilon, use_bias, use_scale, bias_init, scale_init) and returns: jnp