def _nested_reduce(f, x):
  """Fold the function f to the nested structure x (dicts, tuples, lists)."""
  if isinstance(x, list):
    return f([_nested_reduce(f, y) for y in x])
  if isinstance(x, tuple):
    return f([_nested_reduce(f, y) for y in x])
  if isinstance(x, dict):
    return f([_nested_reduce(f, v) for (_, v) in x.items()])
  return x
Fold the function f to the nested structure x (dicts, tuples, lists).
trax/supervised/trainer_lib.py