def average(a, axis=None, weights=None, returned=False): # pylint: disable=missing-docstring
if axis is not None and not isinstance(axis, int):
# TODO(wangpeng): Support tuple of ints as `axis`
raise ValueError('Argument `axis` must be an integer. '
f'Received axis={axis} (of type {type(axis)})')
a = np_array_ops.array(a)
if weights is None: # Treat all weights as 1
if not np.issubdtype(a.dtype.as_numpy_dtype, np.inexact):
a = a.astype(
np_utils.result_type(a.dtype, np_dtypes.default_float_type()))
avg = math_ops.reduce_mean(a, axis=axis)
if returned:
if axis is None:
weights_sum = array_ops.size(a)
else:
weights_sum = array_ops.shape(a)[axis]
weights_sum = math_ops.cast(weights_sum, a.dtype)
else:
if np.issubdtype(a.dtype.as_numpy_dtype, np.inexact):
out_dtype = np_utils.result_type(a.dtype, weights)
else:
out_dtype = np_utils.result_type(a.dtype, weights,
np_dtypes.default_float_type())
a = np_array_ops.array(a, out_dtype)
weights = np_array_ops.array(weights, out_dtype)
def rank_equal_case():
control_flow_ops.Assert(
math_ops.reduce_all(array_ops.shape(a) == array_ops.shape(weights)),
[array_ops.shape(a), array_ops.shape(weights)])
weights_sum = math_ops.reduce_sum(weights, axis=axis)
avg = math_ops.reduce_sum(a * weights, axis=axis) / weights_sum
return avg, weights_sum
if axis is None:
avg, weights_sum = rank_equal_case()
else:
def rank_not_equal_case():
control_flow_ops.Assert(
array_ops.rank(weights) == 1, [array_ops.rank(weights)])
weights_sum = math_ops.reduce_sum(weights)
axes = ops.convert_to_tensor([[axis], [0]])
avg = math_ops.tensordot(a, weights, axes) / weights_sum
return avg, weights_sum
# We condition on rank rather than shape equality, because if we do the
# latter, when the shapes are partially unknown but the ranks are known
# and different, np_utils.cond will run shape checking on the true branch,
# which will raise a shape-checking error.
avg, weights_sum = np_utils.cond(
math_ops.equal(array_ops.rank(a), array_ops.rank(weights)),
rank_equal_case, rank_not_equal_case)
avg = np_array_ops.array(avg)
if returned:
weights_sum = np_array_ops.broadcast_to(weights_sum, array_ops.shape(avg))
return avg, weights_sum
return avg
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 testAllowF64TruePreferF32True(self):
np_dtypes.set_allow_float64(True)
np_dtypes.set_prefer_float32(True)
self.assertEqual(dtypes.float32, np_dtypes.default_float_type())
self.assertEqual(dtypes.float32, np_dtypes._result_type(1.1))
self.assertEqual(dtypes.float64,
np_dtypes._result_type(np.zeros([], np.float64), 1.1))
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 random_ops.random_uniform(
shape=size, minval=low, maxval=high, dtype=dtype)
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 random_ops.random_normal(size, dtype=dtype)
def convert_to_tensor(value, dtype=None):
# A safer version of `tf.convert_to_tensor` to work around b/149876037.
# TODO(wangpeng): Remove this function once the bug is fixed.
if (dtype is None and isinstance(value, six.integer_types)
and value >= 2**63):
dtype = dtypes.uint64
elif (dtype is None and isinstance(value, float)):
dtype = np_dtypes.default_float_type()
return ops.convert_to_tensor(value, dtype=dtype)
def heaviside(x1, x2): # pylint: disable=missing-function-docstring
def f(x1, x2):
return array_ops.where_v2(
x1 < 0, constant_op.constant(0, dtype=x2.dtype),
array_ops.where_v2(x1 > 0, constant_op.constant(1, dtype=x2.dtype),
x2))
y = _bin_op(f, x1, x2)
if not np.issubdtype(y.dtype.as_numpy_dtype, np.inexact):
y = y.astype(np_dtypes.default_float_type())
return y
def f(x1, x2):
if x1.dtype == dtypes.bool:
assert x2.dtype == dtypes.bool
float_ = np_dtypes.default_float_type()
x1 = math_ops.cast(x1, float_)
x2 = math_ops.cast(x2, float_)
if not np_dtypes.is_allow_float64():
# math_ops.truediv in Python3 produces float64 when both inputs are int32
# or int64. We want to avoid that when is_allow_float64() is False.
x1, x2 = _avoid_float64(x1, x2)
return math_ops.truediv(x1, x2)
def heaviside(x1, x2):
def f(x1, x2):
return array_ops.where_v2(
x1 < 0, constant_op.constant(0, dtype=x2.dtype),
array_ops.where_v2(x1 > 0, constant_op.constant(1, dtype=x2.dtype),
x2))
y = _bin_op(f, x1, x2)
if not np.issubdtype(y.dtype, np.inexact):
y = y.astype(np_dtypes.default_float_type())
return y
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 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, )
dtype = np_dtypes.default_float_type()
return random_ops.random_normal(args, dtype=dtype)
def _scalar(tf_fn, x, promote_to_float=False):
"""Computes the tf_fn(x) for each element in `x`.
Args:
tf_fn: function that takes a single Tensor argument.
x: array_like. Could be an ndarray, a Tensor or any object that can be
converted to a Tensor using `ops.convert_to_tensor`.
promote_to_float: whether to cast the argument to a float dtype
(`np_dtypes.default_float_type`) if it is not already.
Returns:
An ndarray with the same shape as `x`. The default output dtype is
determined by `np_dtypes.default_float_type`, unless x is an ndarray with a
floating point type, in which case the output type is same as x.dtype.
"""
x = np_array_ops.asarray(x)
if promote_to_float and not np.issubdtype(x.dtype, np.inexact):
x = x.astype(np_dtypes.default_float_type())
return np_utils.tensor_to_ndarray(tf_fn(x.data))
def convert_to_tensor(value, dtype=None, dtype_hint=None):
"""Wrapper over `tf.convert_to_tensor`.
Args:
value: value to convert
dtype: (optional) the type we would like it to be converted to.
dtype_hint: (optional) soft preference for the type we would like it to
be converted to. `tf.convert_to_tensor` will attempt to convert value
to this type first, but will not fail if conversion is not possible
falling back to inferring the type instead.
"""
# A safer version of `tf.convert_to_tensor` to work around b/149876037.
# TODO(wangpeng): Remove this function once the bug is fixed.
if (dtype is None and isinstance(value, six.integer_types)
and value >= 2**63):
dtype = dtypes.uint64
elif (dtype is None and isinstance(value, float)):
dtype = np_dtypes.default_float_type()
return ops.convert_to_tensor(value, dtype=dtype, dtype_hint=dtype_hint)
def _reduce(tf_fn,
a,
axis=None,
dtype=None,
keepdims=None,
promote_int=_TO_INT64,
tf_bool_fn=None,
preserve_bool=False):
"""A general reduction function.
Args:
tf_fn: the TF reduction function.
a: the array to be reduced.
axis: (optional) the axis along which to do the reduction. If None, all
dimensions are reduced.
dtype: (optional) the dtype of the result.
keepdims: (optional) whether to keep the reduced dimension(s).
promote_int: how to promote integer and bool inputs. There are three
choices: (1) _TO_INT64: always promote them to int64 or uint64; (2)
_TO_FLOAT: always promote them to a float type (determined by
dtypes.default_float_type); (3) None: don't promote.
tf_bool_fn: (optional) the TF reduction function for bool inputs. It will
only be used if `dtype` is explicitly set to `np.bool_` or if `a`'s dtype
is `np.bool_` and `preserve_bool` is True.
preserve_bool: a flag to control whether to use `tf_bool_fn` if `a`'s dtype
is `np.bool_` (some reductions such as np.sum convert bools to integers,
while others such as np.max preserve bools.
Returns:
An ndarray.
"""
if dtype:
dtype = np_utils.result_type(dtype)
if keepdims is None:
keepdims = False
a = asarray(a, dtype=dtype)
if ((dtype == np.bool_ or preserve_bool and a.dtype == np.bool_)
and tf_bool_fn is not None):
return np_utils.tensor_to_ndarray(
tf_bool_fn(input_tensor=a.data, axis=axis, keepdims=keepdims))
if dtype is None:
dtype = a.dtype
if np.issubdtype(dtype, np.integer) or dtype == np.bool_:
if promote_int == _TO_INT64:
# If a is an integer/bool type and whose bit width is less than 64,
# numpy up-casts it to 64-bit.
if dtype == np.bool_:
is_signed = True
width = 8 # We can use any number here that is less than 64
else:
is_signed = np.issubdtype(dtype, np.signedinteger)
width = np.iinfo(dtype).bits
if width < 64:
if is_signed:
dtype = np.int64
else:
dtype = np.uint64
a = a.astype(dtype)
elif promote_int == _TO_FLOAT:
a = a.astype(np_dtypes.default_float_type())
return np_utils.tensor_to_ndarray(
tf_fn(input_tensor=a.data, axis=axis, keepdims=keepdims))
def testAllowF64TruePreferF32False(self):
np_dtypes.set_allow_float64(True)
np_dtypes.set_prefer_float32(False)
self.assertEqual(dtypes.float64, np_dtypes.default_float_type())
self.assertEqual(dtypes.float64, np_dtypes._result_type(1.1))
self.assertEqual(dtypes.complex128, np_dtypes._result_type(1.j))
def testAllowF64False(self, prefer_f32):
np_dtypes.set_allow_float64(False)
np_dtypes.set_prefer_float32(prefer_f32)
self.assertEqual(dtypes.float32, np_dtypes.default_float_type())
self.assertEqual(dtypes.float32,
np_dtypes._result_type(np.zeros([], np.float64), 1.1))
