def diagonal(a, offset=0, axis1=0, axis2=1): # pylint: disable=missing-docstring
a = asarray(a).data
maybe_rank = a.shape.rank
if maybe_rank is not None and offset == 0 and (
axis1 == maybe_rank - 2
or axis1 == -2) and (axis2 == maybe_rank - 1 or axis2 == -1):
return np_utils.tensor_to_ndarray(array_ops.matrix_diag_part(a))
a = moveaxis(np_utils.tensor_to_ndarray(a), (axis1, axis2), (-2, -1)).data
a_shape = array_ops.shape(a)
def _zeros(): # pylint: disable=missing-docstring
return (array_ops.zeros(array_ops.concat([a_shape[:-1], [0]], 0),
dtype=a.dtype), 0)
# All zeros since diag_part doesn't handle all possible k (aka offset).
# Written this way since cond will run shape inference on both branches,
# and diag_part shape inference will fail when offset is out of bounds.
a, offset = np_utils.cond(
np_utils.logical_or(
np_utils.less_equal(offset, -1 * np_utils.getitem(a_shape, -2)),
np_utils.greater_equal(offset, np_utils.getitem(a_shape, -1)),
), _zeros, lambda: (a, offset))
a = np_utils.tensor_to_ndarray(array_ops.matrix_diag_part(a, k=offset))
return a
def sort(a, axis=-1, kind='quicksort', order=None): # pylint: disable=missing-docstring
if kind != 'quicksort':
raise ValueError("Only 'quicksort' is supported.")
if order is not None:
raise ValueError("'order' argument to sort is not supported.")
a = np_array_ops.array(a)
if axis is None:
result_t = sort_ops.sort(array_ops.reshape(a.data, [-1]), 0)
return np_utils.tensor_to_ndarray(result_t)
else:
return np_utils.tensor_to_ndarray(sort_ops.sort(a.data, axis))
def poisson(lam=1.0, size=None):
if size is None:
size = ()
elif np_utils.isscalar(size):
size = (size, )
return np_utils.tensor_to_ndarray(
random_ops.random_poisson(shape=size, lam=lam, dtype=np_dtypes.int64))
def einsum(subscripts, *operands, **kwargs): # pylint: disable=missing-docstring
casting = kwargs.get('casting', 'safe')
optimize = kwargs.get('optimize', False)
if casting == 'safe':
operands = np_array_ops._promote_dtype(*operands) # pylint: disable=protected-access
elif casting == 'no':
operands = [np_array_ops.asarray(x) for x in operands]
else:
raise ValueError('casting policy not supported: %s' % casting)
if not optimize:
# TF doesn't have a "no optimization" option.
# TODO(wangpeng): Print a warning that np and tf use different
# optimizations.
tf_optimize = 'greedy'
elif optimize == True: # pylint: disable=singleton-comparison,g-explicit-bool-comparison
tf_optimize = 'greedy'
elif optimize == 'greedy':
tf_optimize = 'greedy'
elif optimize == 'optimal':
tf_optimize = 'optimal'
else:
raise ValueError('`optimize` method not supported: %s' % optimize)
operands = [x.data for x in operands]
res = special_math_ops.einsum(subscripts, *operands, optimize=tf_optimize)
res = np_utils.tensor_to_ndarray(res)
return res
def _bin_op(tf_fun, a, b, promote=True):
if promote:
a, b = np_array_ops._promote_dtype(a, b) # pylint: disable=protected-access
else:
a = np_array_ops.array(a)
b = np_array_ops.array(b)
return np_utils.tensor_to_ndarray(tf_fun(a.data, b.data))
def diag(v, k=0): # pylint: disable=missing-docstring
"""Raises an error if input is not 1- or 2-d."""
v = asarray(v).data
v_rank = array_ops.rank(v)
v.shape.with_rank_at_most(2)
# TODO(nareshmodi): Consider a np_utils.Assert version that will fail during
# tracing time if the shape is known.
control_flow_ops.Assert(
np_utils.logical_or(math_ops.equal(v_rank, 1),
math_ops.equal(v_rank, 2)), [v_rank])
def _diag(v, k):
return np_utils.cond(
math_ops.equal(array_ops.size(v), 0),
lambda: array_ops.zeros([abs(k), abs(k)], dtype=v.dtype),
lambda: array_ops.matrix_diag(v, k=k))
def _diag_part(v, k):
v_shape = array_ops.shape(v)
v, k = np_utils.cond(
np_utils.logical_or(
np_utils.less_equal(k, -1 * np_utils.getitem(v_shape, 0)),
np_utils.greater_equal(k, np_utils.getitem(v_shape, 1)),
), lambda: (array_ops.zeros([0, 0], dtype=v.dtype), 0), lambda:
(v, k))
result = array_ops.matrix_diag_part(v, k=k)
return result
result = np_utils.cond(math_ops.equal(v_rank, 1), lambda: _diag(v, k),
lambda: _diag_part(v, k))
return np_utils.tensor_to_ndarray(result)
def tri(N, M=None, k=0, dtype=None): # pylint: disable=invalid-name,missing-docstring
M = M if M is not None else N
if dtype is not None:
dtype = np_utils.result_type(dtype)
else:
dtype = np_dtypes.default_float_type()
if k < 0:
lower = -k - 1
if lower > N:
r = array_ops.zeros([N, M], dtype)
else:
# Keep as tf bool, since we create an upper triangular matrix and invert
# it.
o = array_ops.ones([N, M], dtype=dtypes.bool)
r = math_ops.cast(
math_ops.logical_not(array_ops.matrix_band_part(o, lower, -1)),
dtype)
else:
o = array_ops.ones([N, M], dtype)
if k > M:
r = o
else:
r = array_ops.matrix_band_part(o, -1, k)
return np_utils.tensor_to_ndarray(r)
def split(ary, indices_or_sections, axis=0):
ary = asarray(ary)
if not isinstance(indices_or_sections, six.integer_types):
indices_or_sections = _boundaries_to_sizes(ary, indices_or_sections,
axis)
result = array_ops.split(ary.data, indices_or_sections, axis=axis)
return [np_utils.tensor_to_ndarray(a) for a in result]
def standard_normal(size=None):
# TODO(wangpeng): Use new stateful RNG
if size is None:
size = ()
elif np_utils.isscalar(size):
size = (size,)
dtype = np_dtypes.default_float_type()
return np_utils.tensor_to_ndarray(random_ops.random_normal(size, dtype=dtype))
def _argminmax(fn, a, axis=None):
a = np_array_ops.array(a)
if axis is None:
# When axis is None numpy flattens the array.
a_t = array_ops.reshape(a.data, [-1])
else:
a_t = np_array_ops.atleast_1d(a).data
return np_utils.tensor_to_ndarray(fn(input=a_t, axis=axis))
def uniform(low=0.0, high=1.0, size=None):
dtype = np_dtypes.default_float_type()
low = np_array_ops.asarray(low, dtype=dtype)
high = np_array_ops.asarray(high, dtype=dtype)
if size is None:
size = array_ops.broadcast_dynamic_shape(low.shape, high.shape)
return np_utils.tensor_to_ndarray(
random_ops.random_uniform(
shape=size, minval=low, maxval=high, dtype=dtype))
def moveaxis(a, source, destination): # pylint: disable=missing-docstring
"""Raises ValueError if source, destination not in (-ndim(a), ndim(a))."""
if not source and not destination:
return a
a = asarray(a).data
if isinstance(source, int):
source = (source, )
if isinstance(destination, int):
destination = (destination, )
a_rank = np_utils._maybe_static(array_ops.rank(a)) # pylint: disable=protected-access
def _correct_axis(axis, rank):
if axis < 0:
return axis + rank
return axis
source = tuple(_correct_axis(axis, a_rank) for axis in source)
destination = tuple(_correct_axis(axis, a_rank) for axis in destination)
if a.shape.rank is not None:
perm = [i for i in range(a_rank) if i not in source]
for dest, src in sorted(zip(destination, source)):
assert dest <= len(perm)
perm.insert(dest, src)
else:
r = math_ops.range(a_rank)
def _remove_indices(a, b):
"""Remove indices (`b`) from `a`."""
items = array_ops.unstack(sort_ops.sort(array_ops.stack(b)),
num=len(b))
i = 0
result = []
for item in items:
result.append(a[i:item])
i = item + 1
result.append(a[i:])
return array_ops.concat(result, 0)
minus_sources = _remove_indices(r, source)
minus_dest = _remove_indices(r, destination)
perm = array_ops.scatter_nd(array_ops.expand_dims(minus_dest, 1),
minus_sources, [a_rank])
perm = array_ops.tensor_scatter_update(
perm, array_ops.expand_dims(destination, 1), source)
a = array_ops.transpose(a, perm)
return np_utils.tensor_to_ndarray(a)
def where(condition, x=None, y=None):
"""Raises ValueError if exactly one of x or y is not None."""
condition = asarray(condition, dtype=np.bool_)
if x is None and y is None:
return nonzero(condition)
elif x is not None and y is not None:
x, y = _promote_dtype(x, y)
return np_utils.tensor_to_ndarray(
array_ops.where_v2(condition.data, x.data, y.data))
raise ValueError('Both x and y must be ndarrays, or both must be None.')
def _comparison(tf_fun, x1, x2, cast_bool_to_int=False):
dtype = np_utils.result_type(x1, x2)
# Cast x1 and x2 to the result_type if needed.
x1 = np_array_ops.array(x1, dtype=dtype)
x2 = np_array_ops.array(x2, dtype=dtype)
x1 = x1.data
x2 = x2.data
if cast_bool_to_int and x1.dtype == dtypes.bool:
x1 = math_ops.cast(x1, dtypes.int32)
x2 = math_ops.cast(x2, dtypes.int32)
return np_utils.tensor_to_ndarray(tf_fun(x1, x2))
def clip(a, a_min, a_max): # pylint: disable=missing-docstring
if a_min is None and a_max is None:
raise ValueError('Not more than one of `a_min` and `a_max` may be `None`.')
if a_min is None:
return minimum(a, a_max)
elif a_max is None:
return maximum(a, a_min)
else:
a, a_min, a_max = np_array_ops._promote_dtype(a, a_min, a_max) # pylint: disable=protected-access
return np_utils.tensor_to_ndarray(
clip_ops.clip_by_value(
*np_utils.tf_broadcast(a.data, a_min.data, a_max.data)))
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))
def cumsum(a, axis=None, dtype=None): # pylint: disable=missing-docstring
a = asarray(a, dtype=dtype)
if dtype is None:
a = _maybe_promote_to_int(a)
# If axis is None, the input is flattened.
if axis is None:
a = ravel(a)
axis = 0
elif axis < 0:
axis += array_ops.rank(a.data)
return np_utils.tensor_to_ndarray(math_ops.cumsum(a.data, axis))
def around(a, decimals=0): # pylint: disable=missing-docstring
a = asarray(a)
dtype = a.dtype
factor = math.pow(10, decimals)
# Use float as the working dtype instead of a.dtype, because a.dtype can be
# integer and `decimals` can be negative.
float_dtype = np_dtypes.default_float_type()
a = a.astype(float_dtype).data
factor = math_ops.cast(factor, float_dtype)
a = math_ops.multiply(a, factor)
a = math_ops.round(a)
a = math_ops.divide(a, factor)
return np_utils.tensor_to_ndarray(a).astype(dtype)
def randint(low, high=None, size=None, dtype=onp.int): # pylint: disable=missing-function-docstring
low = int(low)
if high is None:
high = low
low = 0
if size is None:
size = ()
elif isinstance(size, int):
size = (size,)
dtype = np_utils.result_type(dtype)
if dtype not in (onp.int32, onp.int64):
raise ValueError('Only np.int32 or np.int64 types are supported')
return np_utils.tensor_to_ndarray(
random_ops.random_uniform(
shape=size, minval=low, maxval=high, dtype=dtype))
def reshape(a, newshape):
"""Reshapes 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`.
newshape: 0-d or 1-d array_like.
Returns:
An ndarray with the contents and dtype of `a` and shape `newshape`.
"""
a = asarray(a)
if isinstance(newshape, np_arrays.ndarray):
newshape = newshape.data
return np_utils.tensor_to_ndarray(array_ops.reshape(a.data, newshape))
def swapaxes(a, axis1, axis2): # pylint: disable=missing-docstring
a = asarray(a)
a_rank = array_ops.rank(a)
if axis1 < 0:
axis1 += a_rank
if axis2 < 0:
axis2 += a_rank
perm = math_ops.range(a_rank)
perm = array_ops.tensor_scatter_update(perm, [[axis1], [axis2]],
[axis2, axis1])
a = array_ops.transpose(a, perm)
return np_utils.tensor_to_ndarray(a)
def trace(a, offset=0, axis1=0, axis2=1, dtype=None): # pylint: disable=missing-docstring
if dtype:
dtype = np_utils.result_type(dtype)
a = np_array_ops.asarray(a, dtype).data
if offset == 0:
a_shape = a.shape
if a_shape.rank is not None:
rank = len(a_shape)
if (axis1 == -2 or axis1 == rank - 2) and (axis2 == -1 or
axis2 == rank - 1):
return np_utils.tensor_to_ndarray(math_ops.trace(a))
a = np_array_ops.diagonal(a, offset, axis1, axis2)
return np_array_ops.sum(a, -1, dtype)
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))
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.
return np_utils.tensor_to_ndarray(math_ops.imag(a.data))
def maximum(x1, x2): # pylint: disable=missing-function-docstring
# Fast path for when maximum is used as relu.
if isinstance(x2, numbers.Real) and not isinstance(
x2, bool) and x2 == 0 and isinstance(
x1, np_arrays.ndarray) and not x1._is_boolean(): # pylint: disable=protected-access
return np_utils.tensor_to_ndarray(
nn_ops.relu(np_array_ops.asarray(x1).data))
def max_or_or(x1, x2):
if x1.dtype == dtypes.bool:
assert x2.dtype == dtypes.bool
return math_ops.logical_or(x1, x2)
return math_ops.maximum(x1, x2)
return _bin_op(max_or_or, x1, x2)
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.
return np_utils.tensor_to_ndarray(math_ops.real(val.data))
def randn(*args):
"""Returns samples from a normal distribution.
Uses `tf.random_normal`.
Args:
*args: The shape of the output array.
Returns:
An ndarray with shape `args` and dtype `float64`.
"""
# TODO(wangpeng): Use new stateful RNG
if np_utils.isscalar(args):
args = (args, )
return np_utils.tensor_to_ndarray(
random_ops.random_normal(args, dtype=DEFAULT_RANDN_DTYPE))
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.
"""
condition = asarray(condition, dtype=bool)
a = asarray(a)
if condition.ndim != 1:
raise ValueError('condition must be a 1-d array.')
# `np.compress` treats scalars as 1-d arrays.
if a.ndim == 0:
a = ravel(a)
if axis is None:
a = ravel(a)
axis = 0
if axis < 0:
axis += a.ndim
assert axis >= 0 and axis < a.ndim
# `tf.boolean_mask` requires the first dimensions of array and condition to
# match. `np.compress` pads condition with False when it is shorter.
condition_t = condition.data
a_t = a.data
if condition.shape[0] < a.shape[axis]:
padding = array_ops.fill([a.shape[axis] - condition.shape[0]], False)
condition_t = array_ops.concat([condition_t, padding], axis=0)
return np_utils.tensor_to_ndarray(
array_ops.boolean_mask(tensor=a_t, mask=condition_t, axis=axis))
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)
return np_utils.tensor_to_ndarray(array_ops.transpose(a=a.data, perm=axes))
def ravel(a):
"""Flattens `a` into a 1-d array.
If `a` is already a 1-d ndarray it is returned as is.
Uses `tf.reshape`.
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:
A 1-d ndarray.
"""
a = asarray(a)
if a.ndim == 1:
return a
return np_utils.tensor_to_ndarray(array_ops.reshape(a.data, [-1]))
