code: def _FrameStack(n_frames):
      """Stacks successive game frames along their last dimension."""
      # Input shape: (B, T, ..., C).
      # Output shape: (B, T, ..., C * n_frames).
      assert n_frames >= 1
      if n_frames == 1:
          return []  # No-op; just let the data flow through.
      return [
          # Create copies of input sequence, shift right by [0, ..., n_frames - 1]
          # frames, and concatenate along the channel dimension.
          tl.Branch(*map(_shift_right, range(n_frames))),
          tl.Concatenate(n_items=n_frames, axis=-1),
      ]
  
imports: from trax import layers
label_desc: Use image node_id 1 for calling a global function with example usage: _FrameStack(n_frames) and returns: list, list