def new_shape(_, old_shape):
# pylint: disable=g-long-lambda
ndim_ = tf.size(old_shape)
return utils.cond(
ndim_ == 0, lambda: tf.constant([1, 1, 1], dtype=tf.int32),
lambda: utils.cond(
ndim_ == 1, lambda: tf.pad(old_shape, [[1, 1]], constant_values=1),
lambda: tf.pad(old_shape, [[0, 1]], constant_values=1)))
def f(a, b): # pylint: disable=missing-docstring
return utils.cond(
utils.logical_or(tf.rank(a) == 0,
tf.rank(b) == 0),
lambda: a * b,
lambda: utils.cond( # pylint: disable=g-long-lambda
tf.rank(b) == 1, lambda: tf.tensordot(a, b, axes=[[-1], [-1]]),
lambda: tf.tensordot(a, b, axes=[[-1], [-2]])))
def f(x1, x2):
try:
return utils.cond(tf.rank(x2) == 1,
lambda: tf.tensordot(x1, x2, axes=1),
lambda: utils.cond(tf.rank(x1) == 1, # pylint: disable=g-long-lambda
lambda: tf.tensordot( # pylint: disable=g-long-lambda
x1, x2, axes=[[0], [-2]]),
lambda: tf.matmul(x1, x2)))
except tf.errors.InvalidArgumentError as err:
six.reraise(ValueError, ValueError(str(err)), sys.exc_info()[2])
def f(a, b): # pylint: disable=missing-docstring
# We can't assign to captured variable `axisa`, so make a new variable
axis_a = axisa
axis_b = axisb
axis_c = axisc
if axis is not None:
axis_a = axis
axis_b = axis
axis_c = axis
if axis_a < 0:
axis_a = utils.add(axis_a, tf.rank(a))
if axis_b < 0:
axis_b = utils.add(axis_b, tf.rank(b))
def maybe_move_axis_to_last(a, axis):
def move_axis_to_last(a, axis):
return tf.transpose(
a, tf.concat(
[tf.range(axis), tf.range(axis + 1, tf.rank(a)), [axis]],
axis=0))
return utils.cond(
axis == utils.subtract(tf.rank(a), 1),
lambda: a,
lambda: move_axis_to_last(a, axis))
a = maybe_move_axis_to_last(a, axis_a)
b = maybe_move_axis_to_last(b, axis_b)
a_dim = utils.getitem(tf.shape(a), -1)
b_dim = utils.getitem(tf.shape(b), -1)
def maybe_pad_0(a, size_of_last_dim):
def pad_0(a):
return tf.pad(a, tf.concat([tf.zeros([tf.rank(a) - 1, 2], tf.int32),
tf.constant([[0, 1]], tf.int32)], axis=0))
return utils.cond(size_of_last_dim == 2,
lambda: pad_0(a),
lambda: a)
a = maybe_pad_0(a, a_dim)
b = maybe_pad_0(b, b_dim)
c = tf.linalg.cross(*utils.tf_broadcast(a, b))
if axis_c < 0:
axis_c = utils.add(axis_c, tf.rank(c))
def move_last_to_axis(a, axis):
r = tf.rank(a)
return tf.transpose(
a, tf.concat(
[tf.range(axis), [r - 1], tf.range(axis, r - 1)], axis=0))
c = utils.cond(
(a_dim == 2) & (b_dim == 2),
lambda: c[..., 2],
lambda: utils.cond( # pylint: disable=g-long-lambda
axis_c == utils.subtract(tf.rank(c), 1),
lambda: c,
lambda: move_last_to_axis(c, axis_c)))
return c
def f(x):
# pylint: disable=g-long-lambda
x = array_creation.asarray(x)
return array_creation.asarray(
utils.cond(
utils.greater(n, tf.rank(x)), lambda: array_methods.reshape(
x, new_shape(n, tf.shape(x.data))).data, lambda: x.data))
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 utils.tensor_to_ndarray(tf.linalg.diag_part(a))
a = moveaxis(utils.tensor_to_ndarray(a), (axis1, axis2), (-2, -1)).data
a_shape = tf.shape(a)
def _zeros(): # pylint: disable=missing-docstring
return (tf.zeros(tf.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 = utils.cond(
utils.logical_or(
utils.less_equal(offset, -1 * utils.getitem(a_shape, -2)),
utils.greater_equal(offset, utils.getitem(a_shape, -1)),
), _zeros, lambda: (a, offset))
a = utils.tensor_to_ndarray(tf.linalg.diag_part(a, k=offset))
return a
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 = tf.rank(v)
v.shape.with_rank_at_most(2)
# TODO(nareshmodi): Consider a utils.Assert version that will fail during
# tracing time if the shape is known.
tf.debugging.Assert(
utils.logical_or(tf.equal(v_rank, 1), tf.equal(v_rank, 2)), [v_rank])
def _diag(v, k):
return utils.cond(
tf.equal(tf.size(v), 0),
lambda: tf.zeros([abs(k), abs(k)], dtype=v.dtype),
lambda: tf.linalg.diag(v, k=k))
def _diag_part(v, k):
v_shape = tf.shape(v)
v, k = utils.cond(
utils.logical_or(
utils.less_equal(k, -1 * utils.getitem(v_shape, 0)),
utils.greater_equal(k, utils.getitem(v_shape, 1)),
), lambda: (tf.zeros([0, 0], dtype=v.dtype), 0), lambda: (v, k))
result = tf.linalg.diag_part(v, k=k)
return result
result = utils.cond(
tf.equal(v_rank, 1), lambda: _diag(v, k), lambda: _diag_part(v, k))
return utils.tensor_to_ndarray(result)
def maybe_pad_0(a, size_of_last_dim):
def pad_0(a):
return tf.pad(a, tf.concat([tf.zeros([tf.rank(a) - 1, 2], tf.int32),
tf.constant([[0, 1]], tf.int32)], axis=0))
return utils.cond(size_of_last_dim == 2,
lambda: pad_0(a),
lambda: a)
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.
"""
if dtype:
dtype = utils.result_type(dtype)
if isinstance(val, arrays_lib.ndarray):
result_t = val.data
else:
result_t = val
if copy and isinstance(result_t, tf.Tensor):
# Note: In eager mode, a copy of `result_t` is made only if it is not on
# the context device.
result_t = tf.identity(result_t)
if not isinstance(result_t, tf.Tensor):
if not dtype:
dtype = utils.result_type(result_t)
# We can't call `convert_to_tensor(result_t, dtype=dtype)` here because
# convert_to_tensor doesn't allow incompatible arguments such as (5.5, int)
# while np.array allows them. We need to convert-then-cast.
def maybe_data(x):
if isinstance(x, arrays_lib.ndarray):
return x.data
return x
# Handles lists of ndarrays
result_t = tf.nest.map_structure(maybe_data, result_t)
result_t = arrays_lib.convert_to_tensor(result_t)
result_t = tf.cast(result_t, dtype=dtype)
elif dtype:
result_t = tf.cast(result_t, dtype)
ndims = tf.rank(result_t)
def true_fn():
old_shape = tf.shape(result_t)
new_shape = tf.concat([tf.ones(ndmin - ndims, tf.int32), old_shape],
axis=0)
return tf.reshape(result_t, new_shape)
result_t = utils.cond(utils.greater(ndmin, ndims), true_fn,
lambda: result_t)
return arrays_lib.tensor_to_ndarray(result_t)
def average(a, axis=None, weights=None, returned=False): # pylint: disable=missing-docstring
if axis is not None and not isinstance(axis, six.integer_types):
# TODO(wangpeng): Support tuple of ints as `axis`
raise ValueError('`axis` must be an integer. Tuple of ints is not '
'supported yet. Got type: %s' % type(axis))
a = array_ops.array(a)
if weights is None: # Treat all weights as 1
if not np.issubdtype(a.dtype, np.inexact):
a = a.astype(
utils.result_type(a.dtype, dtypes.default_float_type()))
avg = tf.reduce_mean(a.data, axis=axis)
if returned:
if axis is None:
weights_sum = tf.size(a.data)
else:
weights_sum = tf.shape(a.data)[axis]
weights_sum = tf.cast(weights_sum, a.data.dtype)
else:
if np.issubdtype(a.dtype, np.inexact):
out_dtype = utils.result_type(a.dtype, weights)
else:
out_dtype = utils.result_type(a.dtype, weights,
dtypes.default_float_type())
a = array_ops.array(a, out_dtype).data
weights = array_ops.array(weights, out_dtype).data
def rank_equal_case():
tf.debugging.Assert(
tf.reduce_all(tf.shape(a) == tf.shape(weights)),
[tf.shape(a), tf.shape(weights)])
weights_sum = tf.reduce_sum(weights, axis=axis)
avg = tf.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():
tf.debugging.Assert(tf.rank(weights) == 1, [tf.rank(weights)])
weights_sum = tf.reduce_sum(weights)
axes = tf.convert_to_tensor([[axis], [0]])
avg = tf.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, utils.cond will run shape checking on the true branch,
# which will raise a shape-checking error.
avg, weights_sum = utils.cond(
tf.rank(a) == tf.rank(weights), rank_equal_case,
rank_not_equal_case)
avg = array_ops.array(avg)
if returned:
weights_sum = array_ops.broadcast_to(weights_sum, tf.shape(avg.data))
return avg, weights_sum
return avg
def hstack(tup):
arrays = [atleast_1d(a) for a in tup]
arrays = _promote_dtype(*arrays) # pylint: disable=protected-access
unwrapped_arrays = [
a.data if isinstance(a, arrays_lib.ndarray) else a for a in arrays
]
rank = tf.rank(unwrapped_arrays[0])
return utils.cond(rank == 1, lambda: tf.concat(unwrapped_arrays, axis=0),
lambda: tf.concat(unwrapped_arrays, axis=1))
def _diag_part(v, k):
v_shape = tf.shape(v)
v, k = utils.cond(
utils.logical_or(
utils.less_equal(k, -1 * utils.getitem(v_shape, 0)),
utils.greater_equal(k, utils.getitem(v_shape, 1)),
), lambda: (tf.zeros([0, 0], dtype=v.dtype), 0), lambda: (v, k))
result = tf.linalg.diag_part(v, k=k)
return result
def maybe_move_axis_to_last(a, axis):
def move_axis_to_last(a, axis):
return tf.transpose(
a, tf.concat(
[tf.range(axis), tf.range(axis + 1, tf.rank(a)), [axis]],
axis=0))
return utils.cond(
axis == utils.subtract(tf.rank(a), 1),
lambda: a,
lambda: move_axis_to_last(a, axis))
def trace(a, offset=0, axis1=0, axis2=1, dtype=None): # pylint: disable=missing-docstring
a = array_ops.asarray(a).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 utils.tensor_to_ndarray(tf.linalg.trace(a))
a_rank = tf.rank(a)
if axis1 < 0:
axis1 += a_rank
if axis2 < 0:
axis2 += a_rank
minaxis = tf.minimum(axis1, axis2)
maxaxis = tf.maximum(axis1, axis2)
# Move axes of interest to the end.
range_rank = tf.range(a_rank)
perm = tf.concat([
range_rank[0:minaxis], range_rank[minaxis + 1:maxaxis],
range_rank[maxaxis + 1:], [axis1, axis2]
],
axis=0)
a = tf.transpose(a, perm)
a_shape = tf.shape(a)
# 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 = utils.cond(
utils.logical_or(
utils.less_equal(offset, -1 * utils.getitem(a_shape, -2)),
utils.greater_equal(offset, utils.getitem(a_shape, -1)),
), lambda: (tf.zeros_like(a), 0), lambda: (a, offset))
a = utils.tensor_to_ndarray(tf.linalg.diag_part(a, k=offset))
return array_ops.sum(a, -1, dtype)
def f(a, b):
return utils.cond(utils.logical_or(tf.rank(a) == 0,
tf.rank(b) == 0), lambda: a * b,
lambda: tf.tensordot(a, b, axes=[[-1], [-1]]))
def _diag(v, k):
return utils.cond(
tf.equal(tf.size(v), 0),
lambda: tf.zeros([abs(k), abs(k)], dtype=v.dtype),
lambda: tf.linalg.diag(v, k=k))
