code: def _normalize(
      mdl: nn.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.
  
      A seperate bias and scale is learned for each feature as specified by
      feature_axes.
      """
      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)
      if use_scale:
          scale = mdl.param_with_axes(
              "scale",
              scale_init,
              reduced_feature_shape,
              param_dtype,
              axes=("embed",),
          ).reshape(feature_shape)
          mul *= scale
      y *= mul
      if use_bias:
          bias = mdl.param_with_axes(
              "bias",
              bias_init,
              reduced_feature_shape,
              param_dtype,
              axes=("embed",),
          ).reshape(feature_shape)
          y += bias
      return jnp.asarray(y, dtype)
  
imports: import dataclasses
  from typing import Any, Callable, Iterable, Optional, Tuple, Union
  import jax.numpy
  from jax import lax
  from flax import linen
  from flax.linen import initializers
  from flax.linen.linear import DotGeneralT, PrecisionLike
  from flax.linen.partitioning import param_with_axes, with_sharding_constraint
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