code: def _validate_batch_shapes(
      batch: types.NestedArray,
      reference_sample: types.NestedArray,
      batch_dims: Tuple[int, ...],
  ) -> None:
      """Verifies shapes of the batch leaves against the reference sample.
  
      Checks that batch dimensions are the same in all leaves in the batch.
      Checks that non-batch dimensions for all leaves in the batch are the same
      as in the reference sample.
  
      Arguments:
        batch: the nested batch of data to be verified.
        reference_sample: the nested array to check non-batch dimensions.
        batch_dims: a Tuple of indices of batch dimensions in the batch shape.
  
      Returns:
        None.
      """
  
      def validate_node_shape(
          reference_sample: jnp.ndarray, batch: jnp.ndarray
      ) -> None:
          expected_shape = batch_dims + reference_sample.shape
          assert (
              batch.shape == expected_shape
          ), f"{batch.shape} != {expected_shape}"
  
      tree_utils.fast_map_structure(
          validate_node_shape, reference_sample, batch
      )
  
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 6 for calling a global function with example usage: _validate_batch_shapes(batch, reference_sample, batch_dims) without return types