def benchmarkBatchMatMulBroadcast(self):
for (a_shape, b_shape) in self.shape_pairs:
with compat.forward_compatibility_horizon(2019, 4, 26):
with ops.Graph().as_default(), \
session.Session(config=benchmark.benchmark_config()) as sess, \
ops.device("/cpu:0"):
matrix_a = variables.Variable(
GetRandomNormalInput(a_shape, np.float32))
matrix_b = variables.Variable(
GetRandomNormalInput(b_shape, np.float32))
variables.global_variables_initializer().run()
# Use batch matmul op's internal broadcasting.
self.run_op_benchmark(
sess,
math_ops.matmul(matrix_a, matrix_b),
min_iters=50,
name="batch_matmul_cpu_{}_{}".format(a_shape, b_shape))
# Manually broadcast the input matrices using the broadcast_to op.
broadcasted_batch_shape = array_ops.broadcast_static_shape(
matrix_a.shape[:-2], matrix_b.shape[:-2])
broadcasted_a_shape = broadcasted_batch_shape.concatenate(
matrix_a.shape[-2:])
broadcasted_b_shape = broadcasted_batch_shape.concatenate(
matrix_b.shape[-2:])
self.run_op_benchmark(
sess,
math_ops.matmul(
array_ops.broadcast_to(matrix_a, broadcasted_a_shape),
array_ops.broadcast_to(matrix_b, broadcasted_b_shape)),
min_iters=50,
name="batch_matmul_manual_broadcast_cpu_{}_{}".format(
a_shape, b_shape))
def testBroadcastScalarToNonScalar(self):
with self.session(use_gpu=True):
x = np.array(1.0, dtype=np.float)
v_tf = array_ops.broadcast_to(constant_op.constant(1.0), [2, 3, 4,
1, 1, 1])
v_np = np.broadcast_to(x, [2, 3, 4, 1, 1, 1])
self.assertAllEqual(v_tf.eval(), v_np)
def testBroadcastToBasic(self):
for dtype in [np.uint8, np.uint16, np.int8, np.int16, np.int32, np.int64]:
with self.test_session(use_gpu=True):
x = np.array([1, 2, 3], dtype=dtype)
v_tf = array_ops.broadcast_to(constant_op.constant(x), [3, 3])
v_np = np.broadcast_to(x, [3, 3])
self.assertAllEqual(v_tf.eval(), v_np)
def testBroadcastToBadOutputShape(self):
with context.eager_mode():
with self.assertRaisesRegexp(errors.InvalidArgumentError,
"Unable to broadcast tensor of shape"):
self.evaluate(
array_ops.broadcast_to(
constant_op.constant([0, 1]), constant_op.constant([2, 1])))
def _broadcast_to_uniform_shape(rt_input, shape, broadcast_inner_dimensions):
"""Broadcasts rt_input to the uniform shape `shape`."""
if isinstance(rt_input, ragged_tensor.RaggedTensor):
raise ValueError('Incompatible with shape: ragged rank mismatch')
if broadcast_inner_dimensions:
return array_ops.broadcast_to(rt_input, shape.inner_dim_sizes)
else:
return rt_input
def testBroadcastToShapeLargerDim2(self):
input_shape = [2, 1, 3, 2, 2, 2, 1, 1, 1]
output_shape = [1, 1, 1, 2, 5, 3, 2, 2, 2, 3, 3, 3]
with self.cached_session(use_gpu=True):
x = np.array(np.random.randint(5, size=input_shape), dtype=np.int32)
v_tf = array_ops.broadcast_to(constant_op.constant(x), output_shape)
v_np = np.broadcast_to(x, output_shape)
self.assertAllEqual(v_tf, v_np)
def testGradientForScalar(self):
x = constant_op.constant(1, dtype=dtypes.float32)
v = array_ops.broadcast_to(x, [2, 4, 3])
out = 2 * v
with self.cached_session():
err = gradient_checker.compute_gradient_error(x, x.get_shape(), out,
out.get_shape())
self.assertLess(err, 1e-4)
def testGradientWithBroadcastAllDimensions(self):
x = constant_op.constant([[1, 2, 3], [4, 5, 6]], dtype=dtypes.float32)
v = array_ops.broadcast_to(x, [5, 4, 6])
out = 2 * v
with self.test_session():
err = gradient_checker.compute_gradient_error(x, x.get_shape(),
out, out.get_shape())
self.assertLess(err, 1e-4)
def testBroadcastToBadOutputShape(self):
with context.eager_mode():
with self.assertRaisesRegexp(
errors.InvalidArgumentError,
"Unable to broadcast tensor of shape"):
self.evaluate(
array_ops.broadcast_to(constant_op.constant([0, 1]),
constant_op.constant([2, 1])))
def testBroadcastToShapeLargerDim2(self):
input_shape = [2, 1, 3, 2, 2, 2, 1, 1, 1]
output_shape = [1, 1, 1, 2, 5, 3, 2, 2, 2, 3, 3, 3]
with self.cached_session(use_gpu=True):
x = np.array(np.random.randint(5, size=input_shape), dtype=np.int32)
v_tf = array_ops.broadcast_to(constant_op.constant(x), output_shape)
v_np = np.broadcast_to(x, output_shape)
self.assertAllEqual(v_tf.eval(), v_np)
def testGradientWithBroadcastAllDimensions(self):
x = constant_op.constant([1], dtype=dtypes.float32)
v = array_ops.broadcast_to(x, [5, 2, 3])
out = 2 * v
with self.cached_session():
err = gradient_checker.compute_gradient_error(x, x.get_shape(),
out, out.get_shape())
self.assertLess(err, 1e-4)
def _broadcast_to_uniform_shape(rt_input, shape, broadcast_inner_dimensions):
"""Broadcasts rt_input to the uniform shape `shape`."""
if isinstance(rt_input, ragged_tensor.RaggedTensor):
raise ValueError('Incompatible with shape: ragged rank mismatch')
if broadcast_inner_dimensions:
return array_ops.broadcast_to(rt_input, shape.inner_dim_sizes)
else:
return rt_input
def testGradientForScalar(self):
x = constant_op.constant(1, dtype=dtypes.float32)
v = array_ops.broadcast_to(x, [2, 4, 3])
out = 2 * v
with self.cached_session():
err = gradient_checker.compute_gradient_error(x, x.get_shape(), out,
out.get_shape())
self.assertLess(err, 1e-4)
def _verifyLu(self, x, output_idx_type=dtypes.int64):
# Verify that Px = LU.
lu, perm = linalg_ops.lu(x, output_idx_type=output_idx_type)
# Prepare the lower factor of shape num_rows x num_rows
lu_shape = np.array(lu.shape.as_list())
batch_shape = lu_shape[:-2]
num_rows = lu_shape[-2]
num_cols = lu_shape[-1]
lower = array_ops.matrix_band_part(lu, -1, 0)
if num_rows > num_cols:
eye = linalg_ops.eye(num_rows,
batch_shape=batch_shape,
dtype=lower.dtype)
lower = array_ops.concat([lower, eye[..., num_cols:]], axis=-1)
elif num_rows < num_cols:
lower = lower[..., :num_rows]
# Fill the diagonal with ones.
ones_diag = array_ops.ones(np.append(batch_shape, num_rows),
dtype=lower.dtype)
lower = array_ops.matrix_set_diag(lower, ones_diag)
# Prepare the upper factor.
upper = array_ops.matrix_band_part(lu, 0, -1)
verification = math_ops.matmul(lower, upper)
# Permute the rows of product of the Cholesky factors.
if num_rows > 0:
# Reshape the product of the triangular factors and permutation indices
# to a single batch dimension. This makes it easy to apply
# invert_permutation and gather_nd ops.
perm_reshaped = array_ops.reshape(perm, [-1, num_rows])
verification_reshaped = array_ops.reshape(verification,
[-1, num_rows, num_cols])
# Invert the permutation in each batch.
inv_perm_reshaped = map_fn.map_fn(array_ops.invert_permutation,
perm_reshaped)
batch_size = perm_reshaped.shape.as_list()[0]
# Prepare the batch indices with the same shape as the permutation.
# The corresponding batch index is paired with each of the `num_rows`
# permutation indices.
batch_indices = math_ops.cast(array_ops.broadcast_to(
math_ops.range(batch_size)[:, None], perm_reshaped.shape),
dtype=output_idx_type)
permuted_verification_reshaped = array_ops.gather_nd(
verification_reshaped,
array_ops.stack([batch_indices, inv_perm_reshaped], axis=-1))
# Reshape the verification matrix back to the original shape.
verification = array_ops.reshape(permuted_verification_reshaped,
lu_shape)
self._verifyLuBase(x, lower, upper, perm, verification,
output_idx_type)
def _verifyLu(self, x, output_idx_type=dtypes.int64):
# Verify that Px = LU.
lu, perm = linalg_ops.lu(x, output_idx_type=output_idx_type)
# Prepare the lower factor of shape num_rows x num_rows
lu_shape = np.array(lu.shape.as_list())
batch_shape = lu_shape[:-2]
num_rows = lu_shape[-2]
num_cols = lu_shape[-1]
lower = array_ops.matrix_band_part(lu, -1, 0)
if num_rows > num_cols:
eye = linalg_ops.eye(
num_rows, batch_shape=batch_shape, dtype=lower.dtype)
lower = array_ops.concat([lower, eye[..., num_cols:]], axis=-1)
elif num_rows < num_cols:
lower = lower[..., :num_rows]
# Fill the diagonal with ones.
ones_diag = array_ops.ones(
np.append(batch_shape, num_rows), dtype=lower.dtype)
lower = array_ops.matrix_set_diag(lower, ones_diag)
# Prepare the upper factor.
upper = array_ops.matrix_band_part(lu, 0, -1)
verification = math_ops.matmul(lower, upper)
# Permute the rows of product of the Cholesky factors.
if num_rows > 0:
# Reshape the product of the triangular factors and permutation indices
# to a single batch dimension. This makes it easy to apply
# invert_permutation and gather_nd ops.
perm_reshaped = array_ops.reshape(perm, [-1, num_rows])
verification_reshaped = array_ops.reshape(verification,
[-1, num_rows, num_cols])
# Invert the permutation in each batch.
inv_perm_reshaped = map_fn.map_fn(array_ops.invert_permutation,
perm_reshaped)
batch_size = perm_reshaped.shape.as_list()[0]
# Prepare the batch indices with the same shape as the permutation.
# The corresponding batch index is paired with each of the `num_rows`
# permutation indices.
batch_indices = math_ops.cast(
array_ops.broadcast_to(
math_ops.range(batch_size)[:, None], perm_reshaped.shape),
dtype=output_idx_type)
permuted_verification_reshaped = array_ops.gather_nd(
verification_reshaped,
array_ops.stack([batch_indices, inv_perm_reshaped], axis=-1))
# Reshape the verification matrix back to the original shape.
verification = array_ops.reshape(permuted_verification_reshaped,
lu_shape)
self._verifyLuBase(x, lower, upper, perm, verification,
output_idx_type)
def benchmarkBatchMatMulBroadcast(self):
for (a_shape, b_shape) in self.shape_pairs:
with compat.forward_compatibility_horizon(2019, 4, 19):
with ops.Graph().as_default(), \
session.Session(config=benchmark.benchmark_config()) as sess, \
ops.device("/cpu:0"):
matrix_a = variables.Variable(
GetRandomNormalInput(a_shape, np.float32))
matrix_b = variables.Variable(
GetRandomNormalInput(b_shape, np.float32))
variables.global_variables_initializer().run()
# Use batch matmul op's internal broadcasting.
self.run_op_benchmark(
sess,
math_ops.matmul(matrix_a, matrix_b),
min_iters=50,
name="batch_matmul_cpu_{}_{}".format(a_shape, b_shape))
# Manually broadcast the input matrices using the broadcast_to op.
broadcasted_batch_shape = array_ops.broadcast_static_shape(
matrix_a.shape[:-2], matrix_b.shape[:-2])
broadcasted_a_shape = broadcasted_batch_shape.concatenate(
matrix_a.shape[-2:])
broadcasted_b_shape = broadcasted_batch_shape.concatenate(
matrix_b.shape[-2:])
self.run_op_benchmark(
sess,
math_ops.matmul(
array_ops.broadcast_to(matrix_a, broadcasted_a_shape),
array_ops.broadcast_to(matrix_b, broadcasted_b_shape)),
min_iters=50,
name="batch_matmul_manual_broadcast_cpu_{}_{}".format(
a_shape, b_shape))
# Use linear_operator_util.matmul_with_broadcast.
name_template = (
"batch_matmul_manual_broadcast_with_linear_operator_util"
"_cpu_{}_{}"
)
self.run_op_benchmark(
sess,
linear_operator_util.matmul_with_broadcast(matrix_a, matrix_b),
min_iters=50,
name=name_template.format(a_shape, b_shape))
def _diag_part(self):
if self.diagonals_format == _MATRIX:
return array_ops.matrix_diag_part(self.diagonals)
elif self.diagonals_format == _SEQUENCE:
diagonal = self.diagonals[1]
return array_ops.broadcast_to(
diagonal, self.shape_tensor()[:-1])
else:
return self.diagonals[..., 1, :]
def testGradientWithIncreasingRank(self):
x = constant_op.constant([[1], [2]],
dtype=dtypes.float32)
v = array_ops.broadcast_to(x, [5, 2, 3])
out = 2 * v
with self.test_session():
err = gradient_checker.compute_gradient_error(x, x.get_shape(),
out, out.get_shape())
self.assertLess(err, 1e-4)
def testBroadcastToBasic(self):
for dtype in [
np.uint8, np.uint16, np.int8, np.int16, np.int32, np.int64
]:
with self.session(use_gpu=True):
x = np.array([1, 2, 3], dtype=dtype)
v_tf = array_ops.broadcast_to(constant_op.constant(x), [3, 3])
v_np = np.broadcast_to(x, [3, 3])
self.assertAllEqual(v_tf.eval(), v_np)
def testGradientWithSameRank(self):
x = constant_op.constant(np.reshape(np.arange(6), (2, 1, 3)),
dtype=dtypes.float32)
v = array_ops.broadcast_to(x, [2, 5, 3])
out = 2 * v
with self.test_session():
err = gradient_checker.compute_gradient_error(x, x.get_shape(),
out, out.get_shape())
self.assertLess(err, 1e-4)
def testGradientWithSameRank(self):
x = constant_op.constant(np.reshape(np.arange(6), (2, 1, 3)),
dtype=dtypes.float32)
v = array_ops.broadcast_to(x, [2, 5, 3])
out = 2 * v
with self.cached_session():
err = gradient_checker.compute_gradient_error(
x, x.get_shape(), out, out.get_shape())
self.assertLess(err, 1e-4)
def testBroadcastToShape(self):
for input_dim in range(1, 6):
for output_dim in range(input_dim, 6):
with self.test_session(use_gpu=True):
input_shape = [2] * input_dim
output_shape = [2] * output_dim
x = np.array(np.random.randint(5, size=input_shape), dtype=np.int32)
v_tf = array_ops.broadcast_to(constant_op.constant(x), output_shape)
v_np = np.broadcast_to(x, output_shape)
self.assertAllEqual(v_tf.eval(), v_np)
def testBroadcastToShape(self):
for input_dim in range(1, 6):
for output_dim in range(input_dim, 6):
with self.cached_session():
input_shape = [2] * input_dim
output_shape = [2] * output_dim
x = np.array(np.random.randint(5, size=input_shape), dtype=np.int32)
v_tf = array_ops.broadcast_to(constant_op.constant(x), output_shape)
v_np = np.broadcast_to(x, output_shape)
self.assertAllEqual(v_tf, v_np)
def true_branch():
cur_layer = 2 * (1 + math_ops.cast(
array_ops.reshape(lod, [-1])[0], dtypes.int32))
cutoff = tf_utils.smart_cond(
random_ops.random_uniform([], 0.0, 1.0) < self.mixing_prob,
lambda: random_ops.random_uniform([
], 1, cur_layer, dtypes.int32), lambda: cur_layer)
return array_ops.where(
array_ops.broadcast_to(self.layer_idx < cutoff,
array_ops.shape(latent1)), latent1,
latent2)
def testGradientForScalar(self):
# TODO(alextp): There is a bug with broadcast_to on GPU from scalars,
# hence we make this test cpu-only.
with ops.device("cpu:0"):
x = constant_op.constant(1, dtype=dtypes.float32)
v = array_ops.broadcast_to(x, [2, 4, 3])
out = 2 * v
with self.test_session():
err = gradient_checker.compute_gradient_error(x, x.get_shape(),
out, out.get_shape())
self.assertLess(err, 1e-4)
def f(p, x):
if p.shape.rank == 0:
p = array_ops.reshape(p, [1])
p = array_ops.unstack(p)
# TODO(wangpeng): Make tf version take a tensor for p instead of a list.
y = math_ops.polyval(p, x)
# If the polynomial is 0-order, numpy requires the result to be broadcast to
# `x`'s shape.
if len(p) == 1:
y = array_ops.broadcast_to(y, x.shape)
return y
def testGradientForScalar(self):
# TODO(alextp): There is a bug with broadcast_to on GPU from scalars,
# hence we make this test cpu-only.
with ops.device("cpu:0"):
x = constant_op.constant(1, dtype=dtypes.float32)
v = array_ops.broadcast_to(x, [2, 4, 3])
out = 2 * v
with self.test_session():
err = gradient_checker.compute_gradient_error(
x, x.get_shape(), out, out.get_shape())
self.assertLess(err, 1e-4)
def testBroadcastToShapeTypeAndInference(self):
for dtype in [dtypes.int32, dtypes.int64]:
with self.cached_session():
x = np.array([1, 2, 3])
v_tf = array_ops.broadcast_to(
constant_op.constant(x), constant_op.constant([3, 3], dtype=dtype))
shape = v_tf.get_shape().as_list()
v_np = np.broadcast_to(x, [3, 3])
self.assertAllEqual(v_tf, v_np)
# check shape inference when shape input is constant
self.assertAllEqual(shape, v_np.shape)
def testGradientWithLargeDim(self):
input_shape = [2, 1, 3, 2, 2, 2, 1, 1, 1]
output_shape = [1, 1, 1, 2, 5, 3, 2, 2, 2, 3, 3, 3]
x = constant_op.constant(np.array(np.random.randn(*input_shape),
dtype=np.float32))
v = array_ops.broadcast_to(x, output_shape)
out = 2 * v
with self.cached_session():
err = gradient_checker.compute_gradient_error(x, x.get_shape(),
out, out.get_shape())
self.assertLess(err, 1e-4)
def testGradientWithLargeDim(self):
input_shape = [2, 1, 3, 2, 2, 2, 1, 1, 1]
output_shape = [1, 1, 1, 2, 5, 3, 2, 2, 2, 3, 3, 3]
x = constant_op.constant(
np.array(np.random.randn(*input_shape), dtype=np.float32))
v = array_ops.broadcast_to(x, output_shape)
out = 2 * v
with self.cached_session():
err = gradient_checker.compute_gradient_error(
x, x.get_shape(), out, out.get_shape())
self.assertLess(err, 1e-4)
def testBroadcastToShapeTypeAndInference(self):
for dtype in [dtypes.int32, dtypes.int64]:
with self.test_session(use_gpu=True):
x = np.array([1, 2, 3])
v_tf = array_ops.broadcast_to(
constant_op.constant(x),
constant_op.constant([3, 3], dtype=dtype))
shape = v_tf.get_shape().as_list()
v_np = np.broadcast_to(x, [3, 3])
self.assertAllEqual(v_tf.eval(), v_np)
# check shape inference when shape input is constant
self.assertAllEqual(shape, v_np.shape)
def ann_with_grads(db, out_grads):
with backprop.GradientTape() as tape:
val, idx = nn_ops.approx_max_k(db, k)
result_in_grads = tape.gradient(val, db, out_grads)
lifted_k_idx = array_ops.reshape(idx, [num_rows, k, 1])
iota_idx = array_ops.broadcast_to(
array_ops.reshape(math_ops.range(num_rows), [num_rows, 1, 1]),
[num_rows, k, 1])
lifted_idx = array_ops.concat([iota_idx, lifted_k_idx], axis=2)
k_idx_s = array_ops.reshape(lifted_idx, [num_rows * k, 2])
k_gra_s = array_ops.reshape(out_grads, [num_rows * k])
expected_in_grads = array_ops.scatter_nd(k_idx_s, k_gra_s,
[num_rows, row_size])
return [expected_in_grads, result_in_grads]
def full_like(a, fill_value, dtype=None, order='K', subok=True, shape=None): # pylint: disable=missing-docstring,redefined-outer-name
"""order, subok and shape arguments mustn't be changed."""
if order != 'K':
raise ValueError('Non-standard orders are not supported.')
if not subok:
raise ValueError('subok being False is not supported.')
if shape:
raise ValueError('Overriding the shape is not supported.')
a = asarray(a).data
dtype = dtype or np_utils.result_type(a)
fill_value = asarray(fill_value, dtype=dtype)
return np_arrays.tensor_to_ndarray(
array_ops.broadcast_to(fill_value.data, array_ops.shape(a)))
def _matmul(self, x, adjoint=False, adjoint_arg=False):
perm = ops.convert_to_tensor_v2_with_dispatch(self.perm)
if adjoint and not self.is_self_adjoint:
# TODO(srvasude): invert_permutation doesn't work on batches so we use
# argsort.
perm = sort_ops.argsort(perm, axis=-1)
x = linalg.adjoint(x) if adjoint_arg else x
# We need to broadcast x and the permutation since tf.gather doesn't
# broadcast.
broadcast_shape = array_ops.broadcast_dynamic_shape(
array_ops.shape(x)[:-1], array_ops.shape(perm))
k = array_ops.shape(x)[-1]
broadcast_x_shape = array_ops.concat([broadcast_shape, [k]], axis=-1)
x = array_ops.broadcast_to(x, broadcast_x_shape)
perm = array_ops.broadcast_to(perm, broadcast_shape)
m = array_ops.shape(x)[-2]
x = array_ops.reshape(x, [-1, m, k])
perm = array_ops.reshape(perm, [-1, m])
y = array_ops.gather(x, perm, axis=-2, batch_dims=1)
return array_ops.reshape(y, broadcast_x_shape)
def tf_broadcast(*args):
"""Broadcast tensors.
Args:
*args: a list of tensors whose shapes are broadcastable against each other.
Returns:
Tensors broadcasted to the common shape.
"""
if len(args) <= 1:
return args
sh = array_ops.shape(args[0])
for arg in args[1:]:
sh = array_ops.broadcast_dynamic_shape(sh, array_ops.shape(arg))
return [array_ops.broadcast_to(arg, sh) for arg in args]
def _to_dense(self):
total_shape = array_ops.broadcast_dynamic_shape(
array_ops.shape(self.row), array_ops.shape(self.col))
n = array_ops.shape(self.row)[-1]
row = array_ops.broadcast_to(self.row, total_shape)
col = array_ops.broadcast_to(self.col, total_shape)
# We concatenate the column in reverse order to the row.
# This gives us 2*n + 1 elements.
elements = array_ops.concat(
[array_ops.reverse(col, axis=[-1]), row[..., 1:]], axis=-1)
# Given the above vector, the i-th row of the Toeplitz matrix
# is the last n elements of the above vector shifted i right
# (hence the first row is just the row vector provided, and
# the first element of each row will belong to the column vector).
# We construct these set of indices below.
indices = math_ops.mod(
# How much to shift right. This corresponds to `i`.
math_ops.range(0, n) +
# Specifies the last `n` indices.
math_ops.range(n - 1, -1, -1)[..., array_ops.newaxis],
# Mod out by the total number of elements to ensure the index is
# non-negative (for tf.gather) and < 2 * n - 1.
2 * n - 1)
return array_ops.gather(elements, indices, axis=-1)
def _BroadcastToGrad(op, grad):
input_value = op.inputs[0]
broadcast_shape = op.inputs[1]
# Assign ids for each position in input_value.
input_value_shape = array_ops.shape(input_value)
input_value_size = array_ops.size(input_value)
ids = array_ops.reshape(math_ops.range(input_value_size),
input_value_shape)
broadcast_ids = array_ops.broadcast_to(ids, broadcast_shape)
# Group by ids and sum its gradients.
grad_flatten = array_ops.reshape(grad, [-1])
broadcast_ids_flatten = array_ops.reshape(broadcast_ids, [-1])
updates_grad_flatten = math_ops.unsorted_segment_sum(
grad_flatten, broadcast_ids_flatten, input_value_size)
updates_grad = array_ops.reshape(updates_grad_flatten, input_value_shape)
return [updates_grad, None]
def _BroadcastToGrad(op, grad):
input_value = op.inputs[0]
broadcast_shape = op.inputs[1]
# Assign ids for each position in input_value.
input_value_shape = array_ops.shape(input_value)
input_value_size = array_ops.size(input_value)
ids = array_ops.reshape(math_ops.range(input_value_size), input_value_shape)
broadcast_ids = array_ops.broadcast_to(ids, broadcast_shape)
# Group by ids and sum its gradients.
grad_flatten = array_ops.reshape(grad, [-1])
broadcast_ids_flatten = array_ops.reshape(broadcast_ids, [-1])
updates_grad_flatten = math_ops.unsorted_segment_sum(grad_flatten,
broadcast_ids_flatten,
input_value_size)
updates_grad = array_ops.reshape(updates_grad_flatten, input_value_shape)
return [updates_grad, None]
def _get_selection_mask(original, num_to_select, axis=-1):
"""Get a selection mask given how many items to select."""
num_to_select = ops.convert_to_tensor(num_to_select)
num_to_select = array_ops.reshape(num_to_select, [-1])
row_lengths = _get_row_lengths_merged_to_axis(original, axis)
num_to_select = array_ops.broadcast_to(num_to_select,
array_ops.shape(row_lengths))
num_to_select = math_ops.cast(num_to_select, row_lengths.dtype)
num_to_select = math_ops.minimum(num_to_select, row_lengths)
ones = array_ops.ones_like(ragged_math_ops.range(num_to_select))
ones = math_ops.cast(ones, dtypes.int32)
zeros_row_length = row_lengths - num_to_select
zeros = math_ops.cast(
array_ops.zeros_like(ragged_math_ops.range(zeros_row_length)),
dtypes.int32)
results = array_ops.concat([ones, zeros], 1)
results = math_ops.cast(results, dtypes.bool)
return results
def testStateBasedSentenceBreaker(self, test_description, doc,
expected_fragment_text):
input = constant_op.constant(doc) # pylint: disable=redefined-builtin
sentence_breaker = (
state_based_sentence_breaker_op.StateBasedSentenceBreaker())
fragment_text, fragment_starts, fragment_ends = (
sentence_breaker.break_sentences_with_offsets(input))
texts, starts, ends = self.evaluate(
(fragment_text, fragment_starts, fragment_ends))
self.assertAllEqual(expected_fragment_text, fragment_text)
for d, text, start, end in zip(doc, texts.to_list(), starts.to_list(),
ends.to_list()):
# broadcast d to match start/end's shape
start = constant_op.constant(start)
end = constant_op.constant(end)
d = array_ops.broadcast_to(d, start.shape)
self.assertAllEqual(string_ops.substr(d, start, end - start), text)
def operator_and_matrix(self,
build_info,
dtype,
use_placeholder,
ensure_self_adjoint_and_pd=False):
shape = list(build_info.shape)
perm = math_ops.range(0, shape[-1])
perm = array_ops.broadcast_to(perm, shape[:-1])
perm = random_ops.random_shuffle(perm)
if use_placeholder:
perm = array_ops.placeholder_with_default(perm, shape=None)
operator = permutation.LinearOperatorPermutation(perm, dtype=dtype)
matrix = math_ops.cast(
math_ops.equal(math_ops.range(0, shape[-1]),
perm[..., array_ops.newaxis]), dtype)
return operator, matrix
def _BroadcastToGrad(op, grad):
input_value = op.inputs[0]
broadcast_shape = op.inputs[1]
# Assign ids for each position in input_value.
input_value_shape = array_ops.shape(input_value)
input_value_size = array_ops.size(input_value)
ids = array_ops.reshape(math_ops.range(input_value_size),
input_value_shape)
broadcast_ids = array_ops.broadcast_to(ids, broadcast_shape)
# Group by ids and sum its gradients.
grad_flatten = array_ops.reshape(grad, [-1])
broadcast_ids_flatten = array_ops.reshape(broadcast_ids, [-1])
# TODO(apassos): Use reduce_sum for gradient now that we only support
# the usual numpy broadcast semantics.
updates_grad_flatten = math_ops.unsorted_segment_sum(
grad_flatten, broadcast_ids_flatten, input_value_size)
updates_grad = array_ops.reshape(updates_grad_flatten, input_value_shape)
return [updates_grad, None]
def triu(m, k=0): # pylint: disable=missing-docstring
m = asarray(m).data
m_shape = m.shape.as_list()
if len(m_shape) < 2:
raise ValueError('Argument to triu must have rank at least 2')
if m_shape[-1] is None or m_shape[-2] is None:
raise ValueError(
'Currently, the last two dimensions of the input array '
'need to be known.')
z = constant_op.constant(0, m.dtype)
mask = tri(*m_shape[-2:], k=k - 1, dtype=bool)
return np_utils.tensor_to_ndarray(
array_ops.where_v2(array_ops.broadcast_to(mask, array_ops.shape(m)), z,
m))
def _broadcast_parameter_with_batch_shape(
param, param_ndims_to_matrix_ndims, batch_shape):
"""Broadcasts `param` with the given batch shape, recursively."""
if hasattr(param, 'batch_shape_tensor'):
# Recursively broadcast every parameter inside the operator.
override_dict = {}
for name, ndims in param._experimental_parameter_ndims_to_matrix_ndims.items(): # pylint:disable=protected-access,line-too-long
sub_param = getattr(param, name)
override_dict[name] = nest.map_structure_up_to(
sub_param, functools.partial(
_broadcast_parameter_with_batch_shape,
batch_shape=batch_shape), sub_param, ndims)
parameters = dict(param.parameters, **override_dict)
return type(param)(**parameters)
base_shape = array_ops.concat(
[batch_shape, array_ops.ones(
[param_ndims_to_matrix_ndims], dtype=dtypes.int32)], axis=0)
return array_ops.broadcast_to(
param,
array_ops.broadcast_dynamic_shape(base_shape, array_ops.shape(param)))
def _spectrum_for_symmetric_circulant(
spectrum_shape,
d,
ensure_self_adjoint_and_pd,
dtype,
):
"""Spectrum for d-dimensional real/symmetric circulant."""
grid_shape = spectrum_shape[-d:]
if grid_shape == (0, ) * d:
kernel = array_ops.reshape(math_ops.cast([], dtype), grid_shape)
else:
kernel = exponential_power_convolution_kernel(
grid_shape=grid_shape,
# power=2 with this scale and no inflation will have some negative
# spectra. It will still be real/symmetric.
length_scale=math_ops.cast([0.2] * d, dtype.real_dtype),
power=1 if ensure_self_adjoint_and_pd else 2,
zero_inflation=0.2 if ensure_self_adjoint_and_pd else None,
)
spectrum = linear_operator_circulant._FFT_OP[d](_to_complex(kernel))
spectrum = math_ops.cast(spectrum, dtype)
return array_ops.broadcast_to(spectrum, spectrum_shape)
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`.
"""
fill_value = asarray(fill_value, dtype=dtype)
if np_utils.isscalar(shape):
shape = array_ops.reshape(shape, [1])
return np_arrays.tensor_to_ndarray(
array_ops.broadcast_to(fill_value.data, shape))
def testBroadcastToString(self):
with self.test_session(use_gpu=True):
x = np.array([b"1", b"2", b"3"])
v_tf = array_ops.broadcast_to(constant_op.constant(x), [3, 3])
v_np = np.broadcast_to(x, [3, 3])
self.assertAllEqual(v_tf.eval(), v_np)
def broadcast(x, dims, name=None):
x = ops.convert_to_tensor(x)
shape = array_ops.concat(
[constant_op.constant(dims),
array_ops.shape(x)], axis=0)
return array_ops.broadcast_to(x, shape, name=name)
def to_tensor(rt_input, default_value=None, name=None):
"""Converts a `RaggedTensor` into a `Tensor`.
Example:
```python
>>> rt = ragged.constant([[9, 8, 7], [], [6, 5], [4]])
>>> print ragged.to_tensor(rt).eval()
[[9 8 7]
[0 0 0]
[6 5 0]
[4 0 0]]
```
Args:
rt_input: The input `RaggedTensor`.
default_value: Value to set for indices not specified in `rt_input`.
Defaults to zero. `default_value` must be broadcastable to
`rt_input.shape[rt_input.ragged_rank + 1:]`.
name: A name prefix for the returned tensors (optional).
Returns:
A `Tensor` with shape `ragged.bounding_shape(rt_input)` and the
values specified by the non-empty values in `rt_input`. Empty values are
assigned `default_value`.
"""
with ops.name_scope(name, 'RaggedToTensor', [rt_input, default_value]):
rt_input = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
rt_input, name='rt_input')
if not ragged_tensor.is_ragged(rt_input):
return rt_input # already dense
if default_value is not None:
default_value = ops.convert_to_tensor(
default_value, name='default_value', dtype=rt_input.dtype)
# If ragged_rank > 1, then recursively convert the ragged values into a
# `Tensor` before we proceed.
values = rt_input.values
if ragged_tensor.is_ragged(values):
values = to_tensor(values, default_value)
# Tile the default value, if necessary.
if default_value is not None:
if values.shape.ndims is not None:
default_value.shape.with_rank_at_most(values.shape.ndims - 1)
if (values.shape.ndims is None or default_value.shape.ndims is None or
values.shape.ndims != default_value.shape.ndims + 1):
value_shape = array_ops.shape(values)[1:]
default_value = array_ops.broadcast_to(default_value, value_shape)
default_value.shape.assert_is_compatible_with(values.shape[1:])
# Get the expected dense shape ([nrows, ncols] + value_shape).
rt_row_lengths = [rt_input.row_splits[1:] - rt_input.row_splits[:-1]]
nrows = array_ops.shape(rt_input.row_splits, out_type=dtypes.int64)[0] - 1
ncols = math_ops.maximum(math_ops.reduce_max(rt_row_lengths), 0)
values_shape = array_ops.shape(values, out_type=dtypes.int64)
value_shape = values_shape[1:]
nvals = values_shape[0]
# Build a default value if none was supplied.
if default_value is None:
default_value = array_ops.zeros(value_shape, dtype=values.dtype)
default_value.shape.assert_is_compatible_with(values.shape[1:])
default_value.set_shape(values.shape[1:])
# Get the row start indices, and expand to shape=[nrows, 1].
starts = array_ops.expand_dims(rt_input.row_splits[:-1], 1)
# Get the row limit indices, and expand to shape=[nrows, 1].
limits = array_ops.expand_dims(rt_input.row_splits[1:], 1)
# Get the column indices, and expand to shape=[1, ncols].
columns = array_ops.expand_dims(math_ops.range(0, ncols), 0)
# Build a list containing the values plus the default value. We will use
# tf.gather to collect values from this list for the `Tensor` (using
# nvals as the index for the default value).
values_and_default = array_ops.concat(
[values, array_ops.stack([default_value])], axis=0)
# Construct a matrix "indices" pointing into values_and_default. I.e.,
# output[r, c] = values_and_default[indices[r, c].
nondefault_index = starts + columns
has_value = nondefault_index < limits
default_index = array_ops.fill(array_ops.stack([nrows, ncols]), nvals)
indices = array_ops.where(has_value, nondefault_index, default_index)
# Gather the results into a `Tensor`.
return array_ops.gather(values_and_default, indices)
def broadcast_matrix_batch_dims(batch_matrices, name=None):
"""Broadcast leading dimensions of zero or more [batch] matrices.
Example broadcasting one batch dim of two simple matrices.
```python
x = [[1, 2],
[3, 4]] # Shape [2, 2], no batch dims
y = [[[1]]] # Shape [1, 1, 1], 1 batch dim of shape [1]
x_bc, y_bc = broadcast_matrix_batch_dims([x, y])
x_bc
==> [[[1, 2],
[3, 4]]] # Shape [1, 2, 2], 1 batch dim of shape [1].
y_bc
==> same as y
```
Example broadcasting many batch dims
```python
x = tf.random.normal(shape=(2, 3, 1, 4, 4))
y = tf.random.normal(shape=(1, 3, 2, 5, 5))
x_bc, y_bc = broadcast_matrix_batch_dims([x, y])
x_bc.shape
==> (2, 3, 2, 4, 4)
y_bc.shape
==> (2, 3, 2, 5, 5)
```
Args:
batch_matrices: Iterable of `Tensor`s, each having two or more dimensions.
name: A string name to prepend to created ops.
Returns:
bcast_matrices: List of `Tensor`s, with `bcast_matricies[i]` containing
the values from `batch_matrices[i]`, with possibly broadcast batch dims.
Raises:
ValueError: If any input `Tensor` is statically determined to have less
than two dimensions.
"""
with ops.name_scope(
name or "broadcast_matrix_batch_dims", values=batch_matrices):
check_ops.assert_proper_iterable(batch_matrices)
batch_matrices = list(batch_matrices)
for i, mat in enumerate(batch_matrices):
batch_matrices[i] = ops.convert_to_tensor(mat)
assert_is_batch_matrix(batch_matrices[i])
if len(batch_matrices) < 2:
return batch_matrices
# Try static broadcasting.
# bcast_batch_shape is the broadcast batch shape of ALL matrices.
# E.g. if batch_matrices = [x, y], with
# x.shape = [2, j, k] (batch shape = [2])
# y.shape = [3, 1, l, m] (batch shape = [3, 1])
# ==> bcast_batch_shape = [3, 2]
bcast_batch_shape = batch_matrices[0].get_shape()[:-2]
for mat in batch_matrices[1:]:
bcast_batch_shape = array_ops.broadcast_static_shape(
bcast_batch_shape,
mat.get_shape()[:-2])
if bcast_batch_shape.is_fully_defined():
for i, mat in enumerate(batch_matrices):
if mat.get_shape()[:-2] != bcast_batch_shape:
bcast_shape = array_ops.concat(
[bcast_batch_shape.as_list(), array_ops.shape(mat)[-2:]], axis=0)
batch_matrices[i] = array_ops.broadcast_to(mat, bcast_shape)
return batch_matrices
# Since static didn't work, do dynamic, which always copies data.
bcast_batch_shape = array_ops.shape(batch_matrices[0])[:-2]
for mat in batch_matrices[1:]:
bcast_batch_shape = array_ops.broadcast_dynamic_shape(
bcast_batch_shape,
array_ops.shape(mat)[:-2])
for i, mat in enumerate(batch_matrices):
batch_matrices[i] = array_ops.broadcast_to(
mat,
array_ops.concat(
[bcast_batch_shape, array_ops.shape(mat)[-2:]], axis=0))
return batch_matrices
def testBroadcastToBool(self):
with self.test_session(use_gpu=True):
x = np.array([True, False, True], dtype=np.bool)
v_tf = array_ops.broadcast_to(constant_op.constant(x), [3, 3])
v_np = np.broadcast_to(x, [3, 3])
self.assertAllEqual(v_tf.eval(), v_np)
def repeat(data, repeats, axis, name=None):
"""Repeats elements of `data`.
Args:
data: An `N`-dimensional tensor.
repeats: A 1-D integer tensor specifying how many times each element in
`axis` should be repeated. `len(repeats)` must equal `data.shape[axis]`.
Supports broadcasting from a scalar value.
axis: `int`. The axis along which to repeat values. Must be less than
`max(N, 1)`.
name: A name for the operation.
Returns:
A tensor with `max(N, 1)` dimensions. Has the same shape as `data`,
except that dimension `axis` has size `sum(repeats)`.
#### Examples:
```python
>>> repeat(['a', 'b', 'c'], repeats=[3, 0, 2], axis=0)
['a', 'a', 'a', 'c', 'c']
>>> repeat([[1, 2], [3, 4]], repeats=[2, 3], axis=0)
[[1, 2], [1, 2], [3, 4], [3, 4], [3, 4]]
>>> repeat([[1, 2], [3, 4]], repeats=[2, 3], axis=1)
[[1, 1, 2, 2, 2], [3, 3, 4, 4, 4]]
```
"""
if not isinstance(axis, int):
raise TypeError("axis must be an int; got %s" % type(axis).__name__)
with ops.name_scope(name, "Repeat", [data, repeats]):
data = ops.convert_to_tensor(data, name="data")
repeats = convert_to_int_tensor(repeats, name="repeats")
repeats.shape.with_rank_at_most(1)
# If `data` is a scalar, then upgrade it to a vector.
data = _with_nonzero_rank(data)
data_shape = array_ops.shape(data)
# If `axis` is negative, then convert it to a positive value.
axis = get_positive_axis(axis, data.shape.ndims)
# Check data Tensor shapes.
if repeats.shape.ndims == 1:
data.shape.dims[axis].assert_is_compatible_with(repeats.shape[0])
# If we know that `repeats` is a scalar, then we can just tile & reshape.
if repeats.shape.ndims == 0:
expanded = array_ops.expand_dims(data, axis + 1)
tiled = tile_one_dimension(expanded, axis + 1, repeats)
result_shape = array_ops.concat(
[data_shape[:axis], [-1], data_shape[axis + 1:]], axis=0)
return array_ops.reshape(tiled, result_shape)
# Broadcast the `repeats` tensor so rank(repeats) == axis + 1.
if repeats.shape.ndims != axis + 1:
repeats_shape = array_ops.shape(repeats)
repeats_ndims = array_ops.rank(repeats)
broadcast_shape = array_ops.concat(
[data_shape[:axis + 1 - repeats_ndims], repeats_shape], axis=0)
repeats = array_ops.broadcast_to(repeats, broadcast_shape)
repeats.set_shape([None] * (axis + 1))
# Create a "sequence mask" based on `repeats`, where slices across `axis`
# contain one `True` value for each repetition. E.g., if
# `repeats = [3, 1, 2]`, then `mask = [[1, 1, 1], [1, 0, 0], [1, 1, 0]]`.
max_repeat = math_ops.maximum(0, math_ops.reduce_max(repeats))
mask = array_ops.sequence_mask(repeats, max_repeat)
# Add a new dimension around each value that needs to be repeated, and
# then tile that new dimension to match the maximum number of repetitions.
expanded = array_ops.expand_dims(data, axis + 1)
tiled = tile_one_dimension(expanded, axis + 1, max_repeat)
# Use `boolean_mask` to discard the extra repeated values. This also
# flattens all dimensions up through `axis`.
masked = array_ops.boolean_mask(tiled, mask)
# Reshape the output tensor to add the outer dimensions back.
if axis == 0:
result = masked
else:
result_shape = array_ops.concat(
[data_shape[:axis], [-1], data_shape[axis + 1:]], axis=0)
result = array_ops.reshape(masked, result_shape)
# Preserve shape information.
if data.shape.ndims is not None:
new_axis_size = 0 if repeats.shape[0] == 0 else None
result.set_shape(data.shape[:axis].concatenate(
[new_axis_size]).concatenate(data.shape[axis + 1:]))
return result
def batch_gather_with_default(params,
indices,
default_value='',
name=None):
"""Same as `batch_gather` but inserts `default_value` for invalid indices.
This operation is similar to `batch_gather` except that it will substitute
the value for invalid indices with `default_value` as the contents.
See `batch_gather` for more details.
Args:
params: A potentially ragged tensor with shape `[B1...BN, P1...PM]` (`N>=0`,
`M>0`).
indices: A potentially ragged tensor with shape `[B1...BN, I]` (`N>=0`).
default_value: A value to be inserted in places where `indices` are out of
bounds. Must be the same dtype as params and either a scalar or rank 1.
name: A name for the operation (optional).
Returns:
A potentially ragged tensor with shape `[B1...BN, I, P2...PM]`.
`result.ragged_rank = max(indices.ragged_rank, params.ragged_rank)`.
#### Example:
```python
>>> params = tf.ragged.constant([
['a', 'b', 'c'],
['d'],
[],
['e']])
>>> indices = tf.ragged.constant([[1, 2, -1], [], [], [0, 10]])
>>> batch_gather_with_default(params, indices, 'FOO')
[['b', 'c', 'FOO'], [], [], ['e', 'FOO']]
```
"""
with ops.name_scope(name, 'RaggedBatchGatherWithDefault'):
params = ragged_tensor.convert_to_tensor_or_ragged_tensor(
params, name='params',
)
indices = ragged_tensor.convert_to_tensor_or_ragged_tensor(
indices, name='indices',
)
default_value = ragged_tensor.convert_to_tensor_or_ragged_tensor(
default_value, name='default_value',
)
# TODO(hterry): lift this restriction and support default_values of
# of rank > 1
if (default_value.shape.ndims is not 0
and default_value.shape.ndims is not 1):
raise ValueError('"default_value" must be a scalar or vector')
upper_bounds = None
if indices.shape.ndims is None:
raise ValueError('Indices must have a known rank.')
if params.shape.ndims is None:
raise ValueError('Params must have a known rank.')
num_batch_dimensions = indices.shape.ndims - 1
pad = None
# The logic for this works as follows:
# - create a padded params, where:
# padded_params[b1...bn, 0] = default_value
# padded_params[b1...bn, i] = params[b1...bn, i-1] (i>0)
# - create an `upper_bounds` Tensor that contains the number of elements
# in each innermost rank. Broadcast `upper_bounds` to be the same shape
# as `indices`.
# - check to see which index in `indices` are out of bounds and substitute
# it with the index containing `default_value` (the first).
# - call batch_gather with the indices adjusted.
with ops.control_dependencies([
check_ops.assert_greater_equal(array_ops.rank(params),
array_ops.rank(indices))]):
if ragged_tensor.is_ragged(params):
row_lengths = ragged_array_ops.expand_dims(
params.row_lengths(axis=num_batch_dimensions),
axis=-1)
upper_bounds = math_ops.cast(row_lengths, indices.dtype)
pad_shape = _get_pad_shape(params, indices)
pad = ragged_tensor_shape.broadcast_to(
default_value, pad_shape)
else:
params_shape = array_ops.shape(params)
pad_shape = array_ops.concat([
params_shape[:num_batch_dimensions],
[1],
params_shape[num_batch_dimensions + 1:params.shape.ndims]
], 0)
upper_bounds = params_shape[num_batch_dimensions]
pad = array_ops.broadcast_to(default_value, pad_shape)
# Add `default_value` as the first value in the innermost (ragged) rank.
pad = math_ops.cast(pad, params.dtype)
padded_params = array_ops.concat(
[pad, params], axis=num_batch_dimensions)
# Adjust the indices by substituting out-of-bound indices to the
# default-value index (which is the first element)
shifted_indices = indices + 1
is_out_of_bounds = (indices < 0) | (indices > upper_bounds)
adjusted_indices = ragged_where_op.where(
is_out_of_bounds,
x=array_ops.zeros_like(indices), y=shifted_indices,
)
return array_ops.batch_gather(
params=padded_params, indices=adjusted_indices, name=name)
def loop_fn(i):
x1 = array_ops.gather(x, i)
return (array_ops.broadcast_to(x1, [2, 2, 3]),
array_ops.broadcast_to(x1, [1, 2, 1, 3]))
def testBroadcastToScalar(self):
with self.test_session(use_gpu=True):
x = np.array(1, dtype=np.int32)
v_tf = array_ops.broadcast_to(constant_op.constant(x), [3, 3])
v_np = np.broadcast_to(x, [3, 3])
self.assertAllEqual(v_tf.eval(), v_np)
