def METHOD_NAME(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)