def _Batch(x, batch_axes):
  """Combine all but batch_axes last axes into batch to avoid shape problems."""
  if batch_axes is None:
    return x
  if isinstance(x, list) and not x:
    return []
  return tl.BatchLeadingAxes(x, n_last_axes_to_keep=batch_axes)
Combine all but batch_axes last axes into batch to avoid shape problems.
trax/models/rl.py