code: def update(
      state: RunningStatisticsState,
      batch: types.NestedArray,
      *,
      config: NestStatisticsConfig = NestStatisticsConfig(),
      weights: Optional[jnp.ndarray] = None,
      std_min_value: float = 1e-6,
      std_max_value: float = 1e6,
      pmap_axis_name: Optional[str] = None,
      validate_shapes: bool = True,
  ) -> RunningStatisticsState:
      """Updates the running statistics with the given batch of data.
  
      Note: data batch and state elements (mean, etc.) must have the same structure.
  
      Note: by default will use int32 for counts and float32 for accumulated
      variance. This results in an integer overflow after 2^31 data points and
      degrading precision after 2^24 batch updates or even earlier if variance
      updates have large dynamic range.
      To improve precision, consider setting jax_enable_x64 to True, see
      https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
  
      Arguments:
        state: The running statistics before the update.
        batch: The data to be used to update the running statistics.
        config: The config that specifies which leaves of the nested structure
          should the running statistics be computed for.
        weights: Weights of the batch data. Should match the batch dimensions.
          Passing a weight of 2. should be equivalent to updating on the
          corresponding data point twice.
        std_min_value: Minimum value for the standard deviation.
        std_max_value: Maximum value for the standard deviation.
        pmap_axis_name: Name of the pmapped axis, if any.
        validate_shapes: If true, the shapes of all leaves of the batch will be
          validated. Enabled by default. Doesn't impact performance when jitted.
  
      Returns:
        Updated running statistics.
      """
      # We require exactly the same structure to avoid issues when flattened
      # batch and state have different order of elements.
      tree.assert_same_structure(batch, state.mean)
      batch_shape = tree.flatten(batch)[0].shape
      # We assume the batch dimensions always go first.
      batch_dims = batch_shape[
          : len(batch_shape) - tree.flatten(state.mean)[0].ndim
      ]
      batch_axis = range(len(batch_dims))
      if weights is None:
          step_increment = np.prod(batch_dims)
      else:
          step_increment = jnp.sum(weights)
      if pmap_axis_name is not None:
          step_increment = jax.lax.psum(
              step_increment, axis_name=pmap_axis_name
          )
      count = state.count + step_increment
  
      # Validation is important. If the shapes don't match exactly, but are
      # compatible, arrays will be silently broadcasted resulting in incorrect
      # statistics.
      if validate_shapes:
          if weights is not None:
              if weights.shape != batch_dims:
                  raise ValueError(f"{weights.shape} != {batch_dims}")
          _validate_batch_shapes(batch, state.mean, batch_dims)
  
      def _compute_node_statistics(
          path: Path,
          mean: jnp.ndarray,
          summed_variance: jnp.ndarray,
          batch: jnp.ndarray,
      ) -> Tuple[jnp.ndarray, jnp.ndarray]:
          assert isinstance(mean, jnp.ndarray), type(mean)
          assert isinstance(summed_variance, jnp.ndarray), type(
              summed_variance
          )
          if not _is_path_included(config, path):
              # Return unchanged.
              return mean, summed_variance
          # The mean and the sum of past variances are updated with Welford's
          # algorithm using batches (see https://stackoverflow.com/q/56402955).
          diff_to_old_mean = batch - mean
          if weights is not None:
              expanded_weights = jnp.reshape(
                  weights,
                  list(weights.shape)
                  + [1] * (batch.ndim - weights.ndim),
              )
              diff_to_old_mean = diff_to_old_mean * expanded_weights
          mean_update = (
              jnp.sum(diff_to_old_mean, axis=batch_axis) / count
          )
          if pmap_axis_name is not None:
              mean_update = jax.lax.psum(
                  mean_update, axis_name=pmap_axis_name
              )
          mean = mean + mean_update
  
          diff_to_new_mean = batch - mean
          variance_update = diff_to_old_mean * diff_to_new_mean
          variance_update = jnp.sum(variance_update, axis=batch_axis)
          if pmap_axis_name is not None:
              variance_update = jax.lax.psum(
                  variance_update, axis_name=pmap_axis_name
              )
          summed_variance = summed_variance + variance_update
          return mean, summed_variance
  
      updated_stats = tree_utils.fast_map_structure_with_path(
          _compute_node_statistics,
          state.mean,
          state.summed_variance,
          batch,
      )
      # map_structure_up_to is slow, so shortcut if we know the input is not
      # structured.
      if isinstance(state.mean, jnp.ndarray):
          mean, summed_variance = updated_stats
      else:
          # Reshape the updated stats from `nest(mean, summed_variance)` to
          # `nest(mean), nest(summed_variance)`.
          mean, summed_variance = [
              tree.map_structure_up_to(
                  state.mean, lambda s, i=idx: s[i], updated_stats
              )
              for idx in range(2)
          ]
  
      def compute_std(
          path: Path, summed_variance: jnp.ndarray, std: jnp.ndarray
      ) -> jnp.ndarray:
          assert isinstance(summed_variance, jnp.ndarray)
          if not _is_path_included(config, path):
              return std
          # Summed variance can get negative due to rounding errors.
          summed_variance = jnp.maximum(summed_variance, 0)
          std = jnp.sqrt(summed_variance / count)
          std = jnp.clip(std, std_min_value, std_max_value)
          return std
  
      std = tree_utils.fast_map_structure_with_path(
          compute_std, summed_variance, state.std
      )
  
      return RunningStatisticsState(
          count=count,
          mean=mean,
          summed_variance=summed_variance,
          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 7 for calling a global function with example usage: update(state, batch) and returns: RunningStatisticsState, mean,  summed_variance, std, mean,  summed_variance, std