[TF-numpy] Adds @np_doc/@np_doc_only to all public ops.
PiperOrigin-RevId: 316748465 Change-Id: I40a49431d6075e47ac05f2946c745ab1c1222214
This commit is contained in:
parent
4db7ec5201
commit
14ee01957c
|
@ -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
|
||||
|
||||
|
|
Loading…
Reference in New Issue