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

PiperOrigin-RevId: 316748465
Change-Id: I40a49431d6075e47ac05f2946c745ab1c1222214
This commit is contained in:
Peng Wang 2020-06-16 13:37:06 -07:00 committed by TensorFlower Gardener
parent 4db7ec5201
commit 14ee01957c
1 changed files with 36 additions and 274 deletions

View File

@ -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):
"""Returns an array of zeros 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.
"""
@np_utils.np_doc(np.zeros_like)
def zeros_like(a, dtype=None): # pylint: disable=missing-docstring
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
# builtin function. Make np_utils.np_doc support builtin functions.
# pylint: disable=g-short-docstring-punctuation,g-no-space-after-docstring-summary,g-doc-return-or-yield,g-doc-args
@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.
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`.
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.
"""
"""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.""" # pylint:disable=g-docstring-missing-newline
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):
"""Compresses `a` by selecting values along `axis` with `condition` true.
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.
"""
@np_utils.np_doc(np.compress)
def compress(condition, a, axis=None): # pylint: disable=redefined-outer-name,missing-function-docstring
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.
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.
"""
"""Only supports modes 'constant', 'reflect' and 'symmetric' currently."""
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