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)