code: def normalize(
      batch: types.NestedArray,
      mean_std: NestedMeanStd,
      max_abs_value: Optional[float] = None,
  ) -> types.NestedArray:
      """Normalizes data using running statistics."""
  
      def normalize_leaf(
          data: jnp.ndarray, mean: jnp.ndarray, std: jnp.ndarray
      ) -> jnp.ndarray:
          # Only normalize inexact types.
          if not jnp.issubdtype(data.dtype, jnp.inexact):
              return data
          data = (data - mean) / std
          if max_abs_value is not None:
              # TODO(b/124318564): remove pylint directive
              data = jnp.clip(
                  data, -max_abs_value, +max_abs_value
              )  # pylint: disable=invalid-unary-operand-type
          return data
  
      return tree_utils.fast_map_structure(
          normalize_leaf, batch, mean_std.mean, mean_std.std
      )
  
imports: import dataclasses
  from typing import Any, Optional, Tuple, Union
  from acme import types
  from acme.utils import tree_utils
  import chex
  import jax
  import jax.numpy
  import numpy
  import tree
label_desc: Use image node_id 8 for calling a global function with example usage: normalize(batch, mean_std, max_abs_value) and returns: tree_utils, data, data