def _np_signature(f):
  """An enhanced funcsigs.signature that can handle numpy.ufunc."""
  if not isinstance(f, np.ufunc):
    try:
      return funcsigs.signature(f)
    except ValueError:
      return None
  def names_from_num(prefix, n):
    if n <= 0:
      return []
    elif n == 1:
      return [prefix]
    else:
      return [prefix + str(i + 1) for i in range(n)]
  input_names = names_from_num('x', f.nin)
  output_names = names_from_num('out', f.nout)
  keyword_only_params = [
      ('where', True),
      ('casting', 'same_kind'),
      ('order', 'K'),
      ('dtype', None),
      ('subok', True),
      ('signature', None),
      ('extobj', None)]
  params = []
  params += [funcsigs.Parameter(name, funcsigs.Parameter.POSITIONAL_ONLY)
             for name in input_names]
  if f.nout > 1:
    params += [funcsigs.Parameter(name, funcsigs.Parameter.POSITIONAL_ONLY,
                                  default=None)
               for name in output_names]
  params += [funcsigs.Parameter(
      'out', funcsigs.Parameter.POSITIONAL_OR_KEYWORD,
      default=None if f.nout == 1 else (None,) * f.nout)]
  params += [funcsigs.Parameter(name, funcsigs.Parameter.KEYWORD_ONLY,
                                default=default)
             for name, default in keyword_only_params]
  return funcsigs.Signature(params)
An enhanced funcsigs.signature that can handle numpy.ufunc.
trax/tf_numpy/numpy_impl/utils.py