def np_abstract_eval(f):
  """Abstract evaluation in numpy by running the real function on 0s."""
  def abstract_f(*args, **kwargs):
    real_args = [nested_map(lambda x: np.zeros(x.shape, x.dtype), a)
                 for a in args]
    real_res = f(*real_args, **kwargs)
    return signature(real_res)
  return abstract_f
Abstract evaluation in numpy by running the real function on 0s.
trax/fastmath/numpy.py