[TF-numpy] Adds @np_doc/@np_doc_only to all public ops.

PiperOrigin-RevId: 316748465
Change-Id: I40a49431d6075e47ac05f2946c745ab1c1222214

tensorflow/python/ops/numpy_ops/np_array_ops.py

@@ -39,51 +39,18 @@ from tensorflow.python.ops.numpy_ops import np_utils
 from tensorflow.python.util import nest
 
 
+@np_utils.np_doc(np.empty)
 def empty(shape, dtype=float):  # pylint: disable=redefined-outer-name
-  """Returns an empty array with the specified shape and dtype.
-
-  Args:
-    shape: A fully defined shape. Could be - NumPy array or a python scalar,
-      list or tuple of integers, - TensorFlow tensor/ndarray of integer type and
-      rank <=1.
-    dtype: Optional, defaults to float. The type of the resulting ndarray. Could
-      be a python type, a NumPy type or a TensorFlow `DType`.
-
-  Returns:
-    An ndarray.
-  """
   return zeros(shape, dtype)
 
 
+@np_utils.np_doc(np.empty_like)
 def empty_like(a, dtype=None):
-  """Returns an empty array with the shape and possibly type of the input array.
-
-  Args:
-    a: array_like. Could be an ndarray, a Tensor or any object that can be
-      converted to a Tensor using `tf.convert_to_tensor`.
-    dtype: Optional, defaults to dtype of the input array. The type of the
-      resulting ndarray. Could be a python type, a NumPy type or a TensorFlow
-      `DType`.
-
-  Returns:
-    An ndarray.
-  """
   return zeros_like(a, dtype)
 
 
+@np_utils.np_doc(np.zeros)
 def zeros(shape, dtype=float):  # pylint: disable=redefined-outer-name
-  """Returns an ndarray with the given shape and type filled with zeros.
-
-  Args:
-    shape: A fully defined shape. Could be - NumPy array or a python scalar,
-      list or tuple of integers, - TensorFlow tensor/ndarray of integer type and
-      rank <=1.
-    dtype: Optional, defaults to float. The type of the resulting ndarray. Could
-      be a python type, a NumPy type or a TensorFlow `DType`.
-
-  Returns:
-    An ndarray.
-  """
   dtype = (
       np_utils.result_type(dtype) if dtype else np_dtypes.default_float_type())
   if isinstance(shape, np_arrays.ndarray):
@@ -91,19 +58,8 @@ def zeros(shape, dtype=float):  # pylint: disable=redefined-outer-name
   return np_arrays.tensor_to_ndarray(array_ops.zeros(shape, dtype=dtype))
 
 
-def zeros_like(a, dtype=None):
+@np_utils.np_doc(np.zeros_like)
-  """Returns an array of zeros with the shape and type of the input array.
+def zeros_like(a, dtype=None):  # pylint: disable=missing-docstring
-
-  Args:
-    a: array_like. Could be an ndarray, a Tensor or any object that can be
-      converted to a Tensor using `tf.convert_to_tensor`.
-    dtype: Optional, defaults to dtype of the input array. The type of the
-      resulting ndarray. Could be a python type, a NumPy type or a TensorFlow
-      `DType`.
-
-  Returns:
-    An ndarray.
-  """
   if isinstance(a, np_arrays.ndarray):
     a = a.data
   if dtype is None:
@@ -117,19 +73,8 @@ def zeros_like(a, dtype=None):
   return np_arrays.tensor_to_ndarray(array_ops.zeros_like(a, dtype))
 
 
+@np_utils.np_doc(np.ones)
 def ones(shape, dtype=float):  # pylint: disable=redefined-outer-name
-  """Returns an ndarray with the given shape and type filled with ones.
-
-  Args:
-    shape: A fully defined shape. Could be - NumPy array or a python scalar,
-      list or tuple of integers, - TensorFlow tensor/ndarray of integer type and
-      rank <=1.
-    dtype: Optional, defaults to float. The type of the resulting ndarray. Could
-      be a python type, a NumPy type or a TensorFlow `DType`.
-
-  Returns:
-    An ndarray.
-  """
   if dtype:
     dtype = np_utils.result_type(dtype)
   if isinstance(shape, np_arrays.ndarray):
@@ -137,19 +82,8 @@ def ones(shape, dtype=float):  # pylint: disable=redefined-outer-name
   return np_arrays.tensor_to_ndarray(array_ops.ones(shape, dtype=dtype))
 
 
+@np_utils.np_doc(np.ones_like)
 def ones_like(a, dtype=None):
-  """Returns an array of ones with the shape and type of the input array.
-
-  Args:
-    a: array_like. Could be an ndarray, a Tensor or any object that can be
-      converted to a Tensor using `tf.convert_to_tensor`.
-    dtype: Optional, defaults to dtype of the input array. The type of the
-      resulting ndarray. Could be a python type, a NumPy type or a TensorFlow
-      `DType`.
-
-  Returns:
-    An ndarray.
-  """
   if isinstance(a, np_arrays.ndarray):
     a = a.data
   if dtype is None:
@@ -191,38 +125,13 @@ def eye(N, M=None, k=0, dtype=float):  # pylint: disable=invalid-name,missing-do
       array_ops.matrix_diag(diagonal=diagonal_, num_rows=N, num_cols=M, k=k))
 
 
+@np_utils.np_doc(np.identity)
 def identity(n, dtype=float):
-  """Returns a square array with ones on the main diagonal and zeros elsewhere.
-
-  Args:
-    n: number of rows/cols.
-    dtype: Optional, defaults to float. The type of the resulting ndarray. Could
-      be a python type, a NumPy type or a TensorFlow `DType`.
-
-  Returns:
-    An ndarray of shape (n, n) and requested type.
-  """
   return eye(N=n, M=n, dtype=dtype)
 
 
+@np_utils.np_doc(np.full)
 def full(shape, fill_value, dtype=None):  # pylint: disable=redefined-outer-name
-  """Returns an array with given shape and dtype filled with `fill_value`.
-
-  Args:
-    shape: A valid shape object. Could be a native python object or an object of
-      type ndarray, numpy.ndarray or tf.TensorShape.
-    fill_value: array_like. Could be an ndarray, a Tensor or any object that can
-      be converted to a Tensor using `tf.convert_to_tensor`.
-    dtype: Optional, defaults to dtype of the `fill_value`. The type of the
-      resulting ndarray. Could be a python type, a NumPy type or a TensorFlow
-      `DType`.
-
-  Returns:
-    An ndarray.
-
-  Raises:
-    ValueError: if `fill_value` can not be broadcast to shape `shape`.
-  """
   if not isinstance(shape, np_arrays.ndarray):
     shape = asarray(np_arrays.convert_to_tensor(shape, dtype_hint=np.int32))
   shape = atleast_1d(shape).data
@@ -251,26 +160,13 @@ def full_like(a, fill_value, dtype=None, order='K', subok=True, shape=None):  #
 
 
 # TODO(wangpeng): investigate whether we can make `copy` default to False.
-# TODO(wangpeng): np_utils.np_doc can't handle np.array because np.array is a
+# pylint: disable=g-short-docstring-punctuation,g-no-space-after-docstring-summary,g-doc-return-or-yield,g-doc-args
-#   builtin function. Make np_utils.np_doc support builtin functions.
+@np_utils.np_doc_only(np.array)
 def array(val, dtype=None, copy=True, ndmin=0):  # pylint: disable=redefined-outer-name
-  """Creates an ndarray with the contents of val.
+  """Since Tensors are immutable, a copy is made only if val is placed on a
-
+  different device than the current one. Even if `copy` is False, a new Tensor
-  Args:
+  may need to be built to satisfy `dtype` and `ndim`. This is used only if `val`
-    val: array_like. Could be an ndarray, a Tensor or any object that can be
+  is an ndarray or a Tensor."""  # pylint:disable=g-docstring-missing-newline
-      converted to a Tensor using `tf.convert_to_tensor`.
-    dtype: Optional, defaults to dtype of the `val`. The type of the resulting
-      ndarray. Could be a python type, a NumPy type or a TensorFlow `DType`.
-    copy: Determines whether to create a copy of the backing buffer. Since
-      Tensors are immutable, a copy is made only if val is placed on a different
-      device than the current one. Even if `copy` is False, a new Tensor may
-      need to be built to satisfy `dtype` and `ndim`. This is used only if `val`
-      is an ndarray or a Tensor.
-    ndmin: The minimum rank of the returned array.
-
-  Returns:
-    An ndarray.
-  """
   if dtype:
     dtype = np_utils.result_type(dtype)
   if isinstance(val, np_arrays.ndarray):
@@ -319,6 +215,7 @@ def array(val, dtype=None, copy=True, ndmin=0):  # pylint: disable=redefined-out
   result_t = np_utils.cond(
       np_utils.greater(ndmin, ndims), true_fn, lambda: result_t)
   return np_arrays.tensor_to_ndarray(result_t)
+# pylint: enable=g-short-docstring-punctuation,g-no-space-after-docstring-summary,g-doc-return-or-yield,g-doc-args
 
 
 @np_utils.np_doc(np.asarray)
@@ -341,6 +238,7 @@ def ascontiguousarray(a, dtype=None):
 
 
 # Numerical ranges.
+@np_utils.np_doc(np.arange)
 def arange(start, stop=None, step=1, dtype=None):
   """Returns `step`-separated values in the range [start, stop).
 
@@ -448,20 +346,8 @@ def diagonal(a, offset=0, axis1=0, axis2=1):  # pylint: disable=missing-docstrin
   return a
 
 
+@np_utils.np_doc(np.diagflat)
 def diagflat(v, k=0):
-  """Returns a 2-d array with flattened `v` as diagonal.
-
-  Args:
-    v: array_like of any rank. Gets flattened when setting as diagonal. Could be
-      an ndarray, a Tensor or any object that can be converted to a Tensor using
-      `tf.convert_to_tensor`.
-    k: Position of the diagonal. Defaults to 0, the main diagonal. Positive
-      values refer to diagonals shifted right, negative values refer to
-      diagonals shifted left.
-
-  Returns:
-    2-d ndarray.
-  """
   v = asarray(v)
   return diag(array_ops.reshape(v.data, [-1]), k)
 
@@ -471,69 +357,22 @@ def _promote_dtype(*arrays):
   return [asarray(a, dtype=dtype) for a in arrays]
 
 
+@np_utils.np_doc(np.all)
 def all(a, axis=None, keepdims=None):  # pylint: disable=redefined-builtin
-  """Whether all array elements or those along an axis evaluate to true.
-
-  Casts the array to bool type if it is not already and uses `tf.reduce_all` to
-  compute the result.
-
-  Args:
-    a: array_like. Could be an ndarray, a Tensor or any object that can be
-      converted to a Tensor using `tf.convert_to_tensor`.
-    axis: Optional. Could be an int or a tuple of integers. If not specified,
-      the reduction is performed over all array indices.
-    keepdims: If true, retains reduced dimensions with length 1.
-
-  Returns:
-    An ndarray. Note that unlike NumPy this does not return a scalar bool if
-    `axis` is None.
-  """
   a = asarray(a, dtype=bool)
   return np_utils.tensor_to_ndarray(
       math_ops.reduce_all(input_tensor=a.data, axis=axis, keepdims=keepdims))
 
 
+@np_utils.np_doc(np.any)
 def any(a, axis=None, keepdims=None):  # pylint: disable=redefined-builtin
-  """Whether any element in the entire array or in an axis evaluates to true.
-
-  Casts the array to bool type if it is not already and uses `tf.reduce_any` to
-  compute the result.
-
-  Args:
-    a: array_like. Could be an ndarray, a Tensor or any object that can be
-      converted to a Tensor using `tf.convert_to_tensor`.
-    axis: Optional. Could be an int or a tuple of integers. If not specified,
-      the reduction is performed over all array indices.
-    keepdims: If true, retains reduced dimensions with length 1.
-
-  Returns:
-    An ndarray. Note that unlike NumPy this does not return a scalar bool if
-    `axis` is None.
-  """
   a = asarray(a, dtype=bool)
   return np_utils.tensor_to_ndarray(
       math_ops.reduce_any(input_tensor=a.data, axis=axis, keepdims=keepdims))
 
 
-def compress(condition, a, axis=None):
+@np_utils.np_doc(np.compress)
-  """Compresses `a` by selecting values along `axis` with `condition` true.
+def compress(condition, a, axis=None):  # pylint: disable=redefined-outer-name,missing-function-docstring
-
-  Uses `tf.boolean_mask`.
-
-  Args:
-    condition: 1-d array of bools. If `condition` is shorter than the array axis
-      (or the flattened array if axis is None), it is padded with False.
-    a: array_like. Could be an ndarray, a Tensor or any object that can be
-      converted to a Tensor using `tf.convert_to_tensor`.
-    axis: Optional. Axis along which to select elements. If None, `condition` is
-      applied on flattened array.
-
-  Returns:
-    An ndarray.
-
-  Raises:
-    ValueError: if `condition` is not of rank 1.
-  """
   condition = asarray(condition, dtype=bool)
   a = asarray(a)
 
@@ -563,8 +402,8 @@ def compress(condition, a, axis=None):
       array_ops.boolean_mask(tensor=a_t, mask=condition_t, axis=axis))
 
 
+@np_utils.np_doc(np.copy)
 def copy(a):
-  """Returns a copy of the array."""
   return array(a, copy=True)
 
 
@@ -611,18 +450,8 @@ def cumsum(a, axis=None, dtype=None):  # pylint: disable=missing-docstring
   return np_utils.tensor_to_ndarray(math_ops.cumsum(a.data, axis))
 
 
+@np_utils.np_doc(np.imag)
 def imag(a):
-  """Returns imaginary parts of all elements in `a`.
-
-  Uses `tf.imag`.
-
-  Args:
-    a: array_like. Could be an ndarray, a Tensor or any object that can be
-      converted to a Tensor using `tf.convert_to_tensor`.
-
-  Returns:
-    An ndarray with the same shape as `a`.
-  """
   a = asarray(a)
   # TODO(srbs): np.imag returns a scalar if a is a scalar, whereas we always
   # return an ndarray.
@@ -760,6 +589,7 @@ def amin(a, axis=None, keepdims=None):
       preserve_bool=True)
 
 
+@np_utils.np_doc(np.var)
 def var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=None):  # pylint: disable=missing-docstring
   if dtype:
     working_dtype = np_utils.result_type(a, dtype)
@@ -829,18 +659,8 @@ def ravel(a):  # pylint: disable=missing-docstring
 setattr(np_arrays.ndarray, 'ravel', ravel)
 
 
+@np_utils.np_doc(np.real)
 def real(val):
-  """Returns real parts of all elements in `a`.
-
-  Uses `tf.real`.
-
-  Args:
-    val: array_like. Could be an ndarray, a Tensor or any object that can be
-      converted to a Tensor using `tf.convert_to_tensor`.
-
-  Returns:
-    An ndarray with the same shape as `a`.
-  """
   val = asarray(val)
   # TODO(srbs): np.real returns a scalar if val is a scalar, whereas we always
   # return an ndarray.
@@ -897,7 +717,6 @@ def around(a, decimals=0):  # pylint: disable=missing-docstring
   return np_utils.tensor_to_ndarray(a).astype(dtype)
 
 
-round_ = around
 setattr(np_arrays.ndarray, '__round__', around)
 
 
@@ -933,51 +752,20 @@ def _reshape_method_wrapper(a, *newshape, **kwargs):
   return reshape(a, newshape, order=order)
 
 
+@np_utils.np_doc(np.expand_dims)
 def expand_dims(a, axis):
-  """Expand the shape of an array.
-
-  Args:
-    a: array_like. Could be an ndarray, a Tensor or any object that can be
-      converted to a Tensor using `tf.convert_to_tensor`.
-    axis: int. axis on which to expand the shape.
-
-  Returns:
-    An ndarray with the contents and dtype of `a` and shape expanded on axis.
-  """
   a = asarray(a)
   return np_utils.tensor_to_ndarray(array_ops.expand_dims(a.data, axis=axis))
 
 
+@np_utils.np_doc(np.squeeze)
 def squeeze(a, axis=None):
-  """Removes single-element axes from the array.
-
-  Args:
-    a: array_like. Could be an ndarray, a Tensor or any object that can be
-      converted to a Tensor using `tf.convert_to_tensor`.
-    axis: scalar or list/tuple of ints.
-  TODO(srbs): tf.squeeze throws error when axis is a Tensor eager execution is
-    enabled. So we cannot allow axis to be array_like here. Fix.
-
-  Returns:
-    An ndarray.
-  """
   a = asarray(a)
   return np_utils.tensor_to_ndarray(array_ops.squeeze(a, axis))
 
 
+@np_utils.np_doc(np.transpose)
 def transpose(a, axes=None):
-  """Permutes dimensions of the array.
-
-  Args:
-    a: array_like. Could be an ndarray, a Tensor or any object that can be
-      converted to a Tensor using `tf.convert_to_tensor`.
-    axes: array_like. A list of ints with length rank(a) or None specifying the
-      order of permutation. The i'th dimension of the output array corresponds
-      to axes[i]'th dimension of the `a`. If None, the axes are reversed.
-
-  Returns:
-    An ndarray.
-  """
   a = asarray(a)
   if axes is not None:
     axes = asarray(axes)
@@ -1113,37 +901,16 @@ def _setitem(arr, index, value):
         [prefix_t, array_ops.expand_dims(subarray.data, 0), postfix_t], 0)
 
 
+# TODO(wangpeng): Make a custom `setattr` that also sets docstring for the
+#   method.
 setattr(np_arrays.ndarray, 'transpose', transpose)
 setattr(np_arrays.ndarray, 'reshape', _reshape_method_wrapper)
 setattr(np_arrays.ndarray, '__setitem__', _setitem)
 
 
+@np_utils.np_doc(np.pad)
 def pad(ary, pad_width, mode, constant_values=0):
-  """Pads an array.
+  """Only supports modes 'constant', 'reflect' and 'symmetric' currently."""
-
-  Args:
-    ary: array_like of rank N. Input array.
-    pad_width: {sequence, array_like, int}. Number of values padded to the edges
-      of each axis. ((before_1, after_1), ... (before_N, after_N)) unique pad
-      widths for each axis. ((before, after),) yields same before and after pad
-      for each axis. (pad,) or int is a shortcut for before = after = pad width
-      for all axes.
-    mode: string. One of the following string values: 'constant' Pads with a
-      constant value. 'reflect' Pads with the reflection of the vector mirrored
-      on the first and last values of the vector along each axis. 'symmetric'
-      Pads with the reflection of the vector mirrored along the edge of the
-      array.
-      **NOTE**: The supported list of `mode` does not match that of numpy's.
-    constant_values: scalar with same dtype as `array`. Used in 'constant' mode
-      as the pad value.  Default is 0.
-
-  Returns:
-    An ndarray padded array of rank equal to `array` with shape increased
-    according to `pad_width`.
-
-  Raises:
-    ValueError if `mode` is not supported.
-  """
   if not (mode == 'constant' or mode == 'reflect' or mode == 'symmetric'):
     raise ValueError('Unsupported padding mode: ' + mode)
   mode = mode.upper()
@@ -1214,24 +981,19 @@ def select(condlist, choicelist, default=0):  # pylint: disable=missing-docstrin
   return output
 
 
+@np_utils.np_doc(np.shape)
 def shape(a):
-  """Return the shape of an array.
-
-  Args:
-    a: array_like. Input array.
-
-  Returns:
-    Tuple of ints.
-  """
   a = asarray(a)
   return a.shape
 
 
+@np_utils.np_doc(np.ndim)
 def ndim(a):
   a = asarray(a)
   return a.ndim
 
 
+@np_utils.np_doc(np.isscalar)
 def isscalar(a):
   return ndim(a) == 0