def nested_map(f, obj, level=0, ignore_nones=True):
  """Maps `f` recursively inside any dicts/lists/tuples in `obj`.

  Args:
    f: A function taking a single object as input. f's input must NOT be a
        dict, list, or tuple, or any subclass of those.
    obj: Either an input object to f or some nested structure of collections
        of (collections of ...) input objects to f.
    level: Level in the nested structure to stop at, counted from the leaves -
        so level 0 is the leaf, level 1 is such that all of its children are at
        level 0 etc.
    ignore_nones: Whether to ignore Nones in the structure, i.e. return None
        without calling `f`.

  Returns:
    An object with the same nested structure as `obj`, but with each input
    object `x` replaced by `f(x)`.
  """
  if _is_at_level(obj, level):
    if ignore_nones and _is_made_of_nones(obj):
      return None
    else:
      return f(obj)

  if _is_namedtuple_instance(obj):
    return type(obj)(*nested_map(f, list(obj), level=level))
  if isinstance(obj, list):
    return [nested_map(f, y, level=level) for y in obj]
  if isinstance(obj, tuple):
    return tuple([nested_map(f, y, level=level) for y in obj])
  if isinstance(obj, dict):
    return {k: nested_map(f, v, level=level) for (k, v) in obj.items()}

  raise ValueError('Non-exhaustive pattern match for {}.'.format(obj))
Maps `f` recursively inside any dicts/lists/tuples in `obj`.

  Args:
    f: A function taking a single object as input. f's input must NOT be a
        dict, list, or tuple, or any subclass of those.
    obj: Either an input object to f or some nested structure of collections
        of (collections of ...) input objects to f.
    level: Level in the nested structure to stop at, counted from the leaves -
        so level 0 is the leaf, level 1 is such that all of its children are at
        level 0 etc.
    ignore_nones: Whether to ignore Nones in the structure, i.e. return None
        without calling `f`.

  Returns:
    An object with the same nested structure as `obj`, but with each input
    object `x` replaced by `f(x)`.
  
trax/fastmath/numpy.py