code: def _expand_scalars(output):
      """If rewards are scalar, expand them."""
      return tree.map_structure(
          tf.experimental.numpy.atleast_1d, output
      )
  
file_path: acme/acme/agents/jax/mbop/dataset.py
imports: import functools
  import itertools
  from typing import Iterator, Optional
  from acme import types
  from acme.jax import running_statistics
  import jax
  import jax.numpy
  import rlds
  import tensorflow
  import tree
label_desc: Use image node_id 3 for calling a global function with example usage: _expand_scalars(output) and returns: tree