code: def squeeze_batch_dim(nest: types.NestedTensor) -> types.NestedTensor:
      """Squeezes out a batch dimension from each leaf of a nested structure."""
      return tree.map_structure(lambda x: tf.squeeze(x, axis=0), nest)
  
imports: import functools
  from typing import List, Optional, Union
  from acme import types
  from acme.utils import tree_utils
  import sonnet
  import tensorflow
  import tree
label_desc: Use image node_id 2 for calling a global function with example usage: squeeze_batch_dim(nest) and returns: tree