def nested_zip(objs):
  """Zips the leaves of each nested structure in `objs`.

  Args:
    objs: List of nested structures to zip.

  Returns:
    An object with the same nested structure as each element of `objs`, with
    leaves zipped together into tuples.
  """
  assert isinstance(objs, (list, tuple))
  assert objs, 'Cannot zip an empty sequence.'

  if _is_at_level(objs, 1):
    return tuple(objs)

  if _is_namedtuple_instance(objs[0]):
    return type(objs[0])(*nested_zip(list(map(list, objs))))
  if isinstance(objs[0], list):
    return [nested_zip([obj[i] for obj in objs]) for i in range(len(objs[0]))]
  if isinstance(objs[0], tuple):
    return nested_zip(list(map(list, objs)))
  if isinstance(objs[0], dict):
    return {k: nested_zip([obj[k] for obj in objs]) for k in objs[0]}

  raise ValueError('Non-exhaustive pattern match for {}.'.format(objs[0]))
Zips the leaves of each nested structure in `objs`.

  Args:
    objs: List of nested structures to zip.

  Returns:
    An object with the same nested structure as each element of `objs`, with
    leaves zipped together into tuples.
  
trax/fastmath/numpy.py