def np_doc(np_fun):
  """Attachs numpy docstring to a function.

  Args:
    np_fun: the numpy function whose docstring will be used.

  Returns:
    A function decorator that attaches the docstring from `np_fun` to the
    decorated function.
  """
  np_sig = _np_signature(np_fun)
  def decorator(f):
    """The decorator."""
    unsupported_params = []
    if np_sig is not None:
      sig = funcsigs.signature(f)
      for name in np_sig.parameters:
        if name not in sig.parameters:
          unsupported_params.append(name)
    f.__doc__ = _np_doc_helper(f, np_fun, unsupported_params)
    return f
  return decorator
Attachs numpy docstring to a function.

  Args:
    np_fun: the numpy function whose docstring will be used.

  Returns:
    A function decorator that attaches the docstring from `np_fun` to the
    decorated function.
  
trax/tf_numpy/numpy_impl/utils.py