def nested_map_multiarg(f, *objs, ignore_nones=True):
  """Maps multi-arg `f` recursively inside any dicts/lists/tuples in `objs`.

  Args:
    f: A function taking len(objs) inputs. f's input must NOT be a
        dict, list, or tuple, or any subclass of those.
    *objs: Either input objects to f or some nested structure of collections
        of (collections of ...) input objects to f.
    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 `objs[0]`, but with each input
    object `x` replaced by `f(*xs)`.
  """
  if isinstance(objs[0], list):
    return [nested_map_multiarg(f, *[o[i] for o in objs])
            for i in range(len(objs[0]))]
  if isinstance(objs[0], tuple):
    return tuple([nested_map_multiarg(f, *[o[i] for o in objs])
                  for i in range(len(objs[0]))])
  if isinstance(objs[0], dict):
    return {k: nested_map_multiarg(f, *[o[k] for o in objs])
            for k in objs[0]}
  if ignore_nones and _is_made_of_nones(objs):
    return None
  return f(*objs)
Maps multi-arg `f` recursively inside any dicts/lists/tuples in `objs`.

  Args:
    f: A function taking len(objs) inputs. f's input must NOT be a
        dict, list, or tuple, or any subclass of those.
    *objs: Either input objects to f or some nested structure of collections
        of (collections of ...) input objects to f.
    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 `objs[0]`, but with each input
    object `x` replaced by `f(*xs)`.
  
trax/fastmath/numpy.py