def testUpdateAddSubGradients(self): with self.cached_session(): indices = constant_op.constant([[3], [1]]) updates = constant_op.constant([9, 10], dtype=dtypes.float32) x = array_ops.ones([4], dtype=dtypes.float32) theoretical, numerical = gradient_checker_v2.compute_gradient( lambda x: array_ops.tensor_scatter_update(x, indices, updates), [x]) self.assertAllClose(theoretical, numerical, 5e-4, 5e-4) theoretical, numerical = gradient_checker_v2.compute_gradient( lambda x: array_ops.tensor_scatter_add(x, indices, updates), [x]) self.assertAllClose(theoretical, numerical, 5e-4, 5e-4) theoretical, numerical = gradient_checker_v2.compute_gradient( lambda x: array_ops.tensor_scatter_sub(x, indices, updates), [x]) self.assertAllClose(theoretical, numerical, 5e-4, 5e-4) theoretical, numerical = gradient_checker_v2.compute_gradient( lambda updates: array_ops.tensor_scatter_update( x, indices, updates), [updates]) self.assertAllClose(theoretical, numerical, 5e-4, 5e-4) theoretical, numerical = gradient_checker_v2.compute_gradient( lambda updates: array_ops.tensor_scatter_add( x, indices, updates), [updates]) self.assertAllClose(theoretical, numerical, 5e-4, 5e-4) theoretical, numerical = gradient_checker_v2.compute_gradient( lambda updates: array_ops.tensor_scatter_sub( x, indices, updates), [updates]) self.assertAllClose(theoretical, numerical, 5e-4, 5e-4)
def testDeterminism(self): a = array_ops.zeros([1]) indices = array_ops.zeros([100000, 1], dtypes.int32) values = np.random.randn(100000) val = self.evaluate(array_ops.tensor_scatter_update(a, indices, values)) for _ in range(5): val2 = self.evaluate(array_ops.tensor_scatter_update(a, indices, values)) self.assertAllEqual(val, val2)
def testUpdateAddSubGradients(self): with self.cached_session(): indices = constant_op.constant([[3], [1]]) updates = constant_op.constant([9, 10], dtype=dtypes.float32) x = array_ops.ones([4], dtype=dtypes.float32) assigned = array_ops.tensor_scatter_update(x, indices, updates) added = array_ops.tensor_scatter_add(x, indices, updates) subbed = array_ops.tensor_scatter_sub(x, indices, updates) err_assigned = gradient_checker.compute_gradient_error( x, [4], assigned, [4]) err_added = gradient_checker.compute_gradient_error( x, [4], added, [4]) err_subbed = gradient_checker.compute_gradient_error( x, [4], subbed, [4]) self.assertLess(err_assigned, 2e-4) self.assertLess(err_added, 2e-4) self.assertLess(err_subbed, 2e-4) err_assigned_wrt_updates = gradient_checker.compute_gradient_error( updates, [2], assigned, [4]) err_added_wrt_updates = gradient_checker.compute_gradient_error( updates, [2], added, [4]) err_subbed_wrt_updates = gradient_checker.compute_gradient_error( updates, [2], subbed, [4]) self.assertLess(err_assigned_wrt_updates, 2e-4) self.assertLess(err_added_wrt_updates, 2e-4) self.assertLess(err_subbed_wrt_updates, 2e-4)
def _TestFn(): indices = constant_op.constant([[4], [3], [1], [7]]) updates = constant_op.constant([9, 10, 11, 12], dtype=dtypes.float32) t = array_ops.ones([8], dtype=dtypes.float32) return array_ops.tensor_scatter_update(t, indices, updates)
def _TensorScatterUpdateGrad(op, grad): indices = op.inputs[1] updates_grad = array_ops.gather_nd(grad, indices) tensor_grad = array_ops.tensor_scatter_update( array_ops.identity(grad), indices, array_ops.zeros_like(op.inputs[2], dtype=grad.dtype)) return [tensor_grad, None, updates_grad]
def testUpdateAddSubGradients(self): with self.cached_session(): indices = constant_op.constant([[3], [1]]) updates = constant_op.constant([9, 10], dtype=dtypes.float32) x = array_ops.ones([4], dtype=dtypes.float32) assigned = array_ops.tensor_scatter_update(x, indices, updates) added = array_ops.tensor_scatter_add(x, indices, updates) subbed = array_ops.tensor_scatter_sub(x, indices, updates) err_assigned = gradient_checker.compute_gradient_error( x, [4], assigned, [4]) err_added = gradient_checker.compute_gradient_error(x, [4], added, [4]) err_subbed = gradient_checker.compute_gradient_error(x, [4], subbed, [4]) self.assertLess(err_assigned, 2e-4) self.assertLess(err_added, 2e-4) self.assertLess(err_subbed, 2e-4) err_assigned_wrt_updates = gradient_checker.compute_gradient_error( updates, [2], assigned, [4]) err_added_wrt_updates = gradient_checker.compute_gradient_error( updates, [2], added, [4]) err_subbed_wrt_updates = gradient_checker.compute_gradient_error( updates, [2], subbed, [4]) self.assertLess(err_assigned_wrt_updates, 2e-4) self.assertLess(err_added_wrt_updates, 2e-4) self.assertLess(err_subbed_wrt_updates, 2e-4)
def _TensorScatterUpdateGrad(op, grad): indices = op.inputs[1] updates_grad = array_ops.gather_nd(grad, indices) tensor_grad = array_ops.tensor_scatter_update( array_ops.identity(grad), indices, array_ops.zeros_like(op.inputs[2], dtype=grad.dtype)) return [tensor_grad, None, updates_grad]
def testUpdateRepeatedIndices1D(self): if test_util.is_gpu_available(): self.skipTest("Duplicate indices scatter is non-deterministic on GPU") a = array_ops.zeros([10, 1]) b = array_ops.tensor_scatter_update(a, [[5], [5]], [[4], [8]]) self.assertAllEqual( b, constant_op.constant([[0.], [0.], [0.], [0.], [0.], [8.], [0.], [0.], [0.], [0.]]))
def moveaxis(a, source, destination): # pylint: disable=missing-docstring """Raises ValueError if source, destination not in (-ndim(a), ndim(a)).""" if not source and not destination: return a a = asarray(a).data if isinstance(source, int): source = (source, ) if isinstance(destination, int): destination = (destination, ) a_rank = np_utils._maybe_static(array_ops.rank(a)) # pylint: disable=protected-access def _correct_axis(axis, rank): if axis < 0: return axis + rank return axis source = tuple(_correct_axis(axis, a_rank) for axis in source) destination = tuple(_correct_axis(axis, a_rank) for axis in destination) if a.shape.rank is not None: perm = [i for i in range(a_rank) if i not in source] for dest, src in sorted(zip(destination, source)): assert dest <= len(perm) perm.insert(dest, src) else: r = math_ops.range(a_rank) def _remove_indices(a, b): """Remove indices (`b`) from `a`.""" items = array_ops.unstack(sort_ops.sort(array_ops.stack(b)), num=len(b)) i = 0 result = [] for item in items: result.append(a[i:item]) i = item + 1 result.append(a[i:]) return array_ops.concat(result, 0) minus_sources = _remove_indices(r, source) minus_dest = _remove_indices(r, destination) perm = array_ops.scatter_nd(array_ops.expand_dims(minus_dest, 1), minus_sources, [a_rank]) perm = array_ops.tensor_scatter_update( perm, array_ops.expand_dims(destination, 1), source) a = array_ops.transpose(a, perm) return np_utils.tensor_to_ndarray(a)
def _descending_sort(values, axis, return_argsort=False): """Sorts values in reverse using `top_k`. Args: values: Tensor of numeric values. axis: Index of the axis which values should be sorted along. return_argsort: If False, return the sorted values. If True, return the indices that would sort the values. Returns: The sorted values. """ # TODO(b/190410105): replace with a proper sort kernel. k = array_ops.shape(values)[axis] rank = array_ops.rank(values) static_rank = values.shape.ndims # Fast path: sorting the last axis. if axis == -1 or axis + 1 == values.get_shape().ndims: top_k_input = values transposition = None else: # Otherwise, transpose the array. Swap axes `axis` and `rank - 1`. if axis < 0: # Calculate the actual axis index if counting from the end. Use the static # rank if available, or else make the axis back into a tensor. axis += static_rank or rank if static_rank is not None: # Prefer to calculate the transposition array in NumPy and make it a # constant. transposition = constant_op.constant( np.r_[ # Axes up to axis are unchanged. np.arange(axis), # Swap axis and rank - 1. [static_rank - 1], # Axes in [axis + 1, rank - 1) are unchanged. np.arange(axis + 1, static_rank - 1), # Swap axis and rank - 1. [axis]], name='transposition') else: # Generate the transposition array from the tensors. transposition = array_ops.tensor_scatter_update( math_ops.range(rank), [[axis], [rank-1]], [rank-1, axis]) top_k_input = array_ops.transpose(values, transposition) values, indices = nn_ops.top_k(top_k_input, k) return_value = indices if return_argsort else values if transposition is not None: # transposition contains a single cycle of length 2 (swapping 2 elements), # so it is an involution (it is its own inverse). return_value = array_ops.transpose(return_value, transposition) return return_value
def testUpdateRepeatedIndices2D(self): if test_util.is_gpu_available(): self.skipTest("Duplicate indices scatter is non-deterministic on GPU") a = array_ops.zeros([10, 10]) b = array_ops.tensor_scatter_update( a, [[5], [6], [6]], [math_ops.range(10), math_ops.range(11, 21), math_ops.range(10, 20)]) self.assertAllEqual( b[6], constant_op.constant( [10., 11., 12., 13., 14., 15., 16., 17., 18., 19.]))
def testUpdateAddSub(self): indices = constant_op.constant([[4], [3], [1], [7]]) updates = constant_op.constant([9, 10, 11, 12], dtype=dtypes.float32) t = array_ops.ones([8], dtype=dtypes.float32) assigned = array_ops.tensor_scatter_update(t, indices, updates) added = array_ops.tensor_scatter_add(t, indices, updates) subbed = array_ops.tensor_scatter_sub(t, indices, updates) self.assertAllEqual(assigned, constant_op.constant([1, 11, 1, 10, 9, 1, 1, 12])) self.assertAllEqual(added, constant_op.constant([1, 12, 1, 11, 10, 1, 1, 13])) self.assertAllEqual(subbed, constant_op.constant([1, -10, 1, -9, -8, 1, 1, -11]))
def testUpdateMinMax(self): indices = constant_op.constant([[4], [3], [1], [7]]) updates = constant_op.constant([0, 2, -1, 1.2], dtype=dtypes.float32) t = array_ops.ones([8], dtype=dtypes.float32) assigned = array_ops.tensor_scatter_update(t, indices, updates) min_result = array_ops.tensor_scatter_min(t, indices, updates) max_result = array_ops.tensor_scatter_max(t, indices, updates) self.assertAllEqual(assigned, constant_op.constant([1, -1, 1, 2, 0, 1, 1, 1.2])) self.assertAllEqual(min_result, constant_op.constant([1, -1, 1, 1, 0, 1, 1, 1])) self.assertAllEqual(max_result, constant_op.constant([1, 1, 1, 2, 1, 1, 1, 1.2]))
def testUpdateAddSub(self): indices = constant_op.constant([[4], [3], [1], [7]]) updates = constant_op.constant([9, 10, 11, 12], dtype=dtypes.float32) t = array_ops.ones([8], dtype=dtypes.float32) assigned = array_ops.tensor_scatter_update(t, indices, updates) added = array_ops.tensor_scatter_add(t, indices, updates) subbed = array_ops.tensor_scatter_sub(t, indices, updates) self.assertAllEqual(assigned, constant_op.constant([1, 11, 1, 10, 9, 1, 1, 12])) self.assertAllEqual(added, constant_op.constant([1, 12, 1, 11, 10, 1, 1, 13])) self.assertAllEqual( subbed, constant_op.constant([1, -10, 1, -9, -8, 1, 1, -11]))
def swapaxes(a, axis1, axis2): # pylint: disable=missing-docstring a = asarray(a) a_rank = array_ops.rank(a) if axis1 < 0: axis1 += a_rank if axis2 < 0: axis2 += a_rank perm = math_ops.range(a_rank) perm = array_ops.tensor_scatter_update(perm, [[axis1], [axis2]], [axis2, axis1]) a = array_ops.transpose(a, perm) return np_utils.tensor_to_ndarray(a)
def testTensorScatterUpdateWithStrings(self): indices = constant_op.constant([[4], [3], [1], [7]]) updates = constant_op.constant(["there", "there", "there", "12"], dtype=dtypes.string) tensor = constant_op.constant([ "hello", "hello", "hello", "hello", "hello", "hello", "hello", "hello" ], dtype=dtypes.string) updated = array_ops.tensor_scatter_update(tensor, indices, updates) self.assertAllEqual( updated, constant_op.constant([ "hello", "there", "hello", "there", "there", "hello", "hello", "12" ]))
def _ragged_tensor_scatter_nd_update(params, indices, updates): """Version of tensor_scatter_nd_update() where the values are ragged.""" # Create a RT in the shape of `params` and containing the "global" positions. # Here "global" means the element position in the flat values Tensor. global_positions_flat = math_ops.range(array_ops.size(params.flat_values)) global_positions = params.with_flat_values(global_positions_flat) global_indices = array_ops.batch_gather(global_positions, indices) update_indices = global_indices.flat_values update_indices = array_ops.expand_dims(update_indices, -1) update_indices = math_ops.cast(update_indices, params.dtype) params_flat = params.flat_values update_values = math_ops.cast(updates.flat_values, params_flat.dtype) results_flat = array_ops.tensor_scatter_update(params_flat, update_indices, update_values) return params.with_flat_values(results_flat)
def get_selectable(self, input_ids, axis): """See `get_selectable()` in superclass.""" selectable = super(FirstNItemSelector, self).get_selectable(input_ids, axis) axis = array_ops.get_positive_axis( axis, input_ids.ragged_rank + input_ids.flat_values.shape.rank) # Create a positions RT and mask out positions that are not selectable positions_flat = math_ops.range(array_ops.size(input_ids.flat_values)) positions = input_ids.with_flat_values(positions_flat) selectable_positions = ragged_array_ops.boolean_mask( positions, selectable) # merge to the desired axis selectable_positions = selectable_positions.merge_dims( 1, axis) if axis > 1 else selectable_positions # Get a selection mask based off of how many items are desired for selection merged_axis = axis - (axis - 1) selection_mask = _get_selection_mask(selectable_positions, self._num_to_select, merged_axis) # Mask out positions that were not selected. selected_positions = ragged_array_ops.boolean_mask( selectable_positions, selection_mask) # Now that we have all the positions which were chosen, we recreate a mask # (matching the original input's shape) where the value is True if it was # selected. We do this by creating a "all false" RT and scattering true # values to the positions chosen for selection. all_true = selected_positions.with_flat_values( array_ops.ones_like(selected_positions.flat_values)) all_false = math_ops.cast( array_ops.zeros(array_ops.shape(input_ids.flat_values)), dtypes.int32) results_flat = array_ops.tensor_scatter_update( all_false, array_ops.expand_dims(selected_positions.flat_values, -1), all_true.flat_values) results = input_ids.with_flat_values(results_flat) results = math_ops.cast(results, dtypes.bool) # Reduce until input.shape[:axis] for _ in range(input_ids.shape.ndims - axis - 1): results = math_ops.reduce_all(results, -1) return results
def _TestFn(): indices = constant_op.constant([[4], [3], [1], [7]]) updates = constant_op.constant([9, 10, 11, 12], dtype=dtypes.float32) t = array_ops.ones([8], dtype=dtypes.float32) return array_ops.tensor_scatter_update(t, indices, updates)
def _TestFn(): indices = constant_op.constant([[4], [3], [1], [7]]) updates = constant_op.constant([9, 10, 11, 12], dtype=dtype) # pylint: disable=cell-var-from-loop t = array_ops.ones([8], dtype=dtype) # pylint: disable=cell-var-from-loop return array_ops.tensor_scatter_update(t, indices, updates)