code: def add_batch_dim(nest: types.NestedTensor) -> types.NestedTensor:
      """Adds a batch dimension to each leaf of a nested structure of Tensors."""
      return tree.map_structure(
          lambda x: tf.expand_dims(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 1 for calling a global function with example usage: add_batch_dim(nest) and returns: tree