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)
  ]
Stacks successive game frames along their last dimension.
trax/models/atari_cnn.py