def _length_fn(example, length_axis, length_keys):
  """Length is the maximum of shape on length_axis over length_keys."""
  if isinstance(example, (list, tuple)):
    return max([example[i].shape[length_axis] for i in length_keys])
  return example.shape[length_axis]
Length is the maximum of shape on length_axis over length_keys.
trax/data/inputs.py