code: def _append_n_step_return(output, n_step_return):
      """Append n-step return to an output step."""
      output[N_STEP_RETURN] = n_step_return
      return 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 1 for calling a global function with example usage: _append_n_step_return(output, n_step_return) and returns: output