Replace references to ops.IndexedSlices with indexed_slices.IndexedSl…

…ices.

(IndexedSlices was moved to its own module, but most references had not yet been updated to refer to the new location, and still used a backwards-compatibility alias.)

PiperOrigin-RevId: 405919446
Change-Id: Iaa0bc835d9a20e67666c78583203e769f8fc20ae

tensorflow/compiler/tests/eager_test.py

@@ -24,6 +24,7 @@
 from tensorflow.python.eager import function
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
 from tensorflow.python.framework import ops
 from tensorflow.python.layers import convolutional
 from tensorflow.python.layers import pooling
@@ -278,7 +279,7 @@ def testAdamSparse(self):
         embedding = embedding_ops.embedding_lookup(embedding_matrix, [1])
         y = math_ops.reduce_sum(embedding)
       dy_dx = tape.gradient(y, embedding_matrix)
-      self.assertIsInstance(dy_dx, ops.IndexedSlices)
+      self.assertIsInstance(dy_dx, indexed_slices.IndexedSlices)
       optimizer = adam.AdamOptimizer(0.1)
       # The gradient application operations will run on CPU because optimizer
       # updates are always collocated with the variable.

tensorflow/python/client/session.py

@@ -31,6 +31,7 @@
 from tensorflow.python.framework import device
 from tensorflow.python.framework import error_interpolation
 from tensorflow.python.framework import errors
+from tensorflow.python.framework import indexed_slices
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import sparse_tensor
 from tensorflow.python.ops import session_ops
@@ -73,7 +74,7 @@ def partial_run(self, handle, fetches, feed_dict=None):
 
 
 def _get_indexed_slices_value_from_fetches(fetched_vals):
-  return ops.IndexedSlicesValue(
+  return indexed_slices.IndexedSlicesValue(
       fetched_vals[0], fetched_vals[1],
       fetched_vals[2] if len(fetched_vals) == 3 else None)
 
@@ -119,7 +120,7 @@ def _get_feeds_for_indexed_slices(feed, feed_val):
      lambda feed: [feed.indices, feed.values, feed.dense_shape]),
     # IndexedSlices are fetched as IndexedSlicesValues. They can be fed
     # IndexedSlicesValues or normal tuples.
-    (ops.IndexedSlices,
+    (indexed_slices.IndexedSlices,
      lambda fetch: ([fetch.values, fetch.indices] if fetch.dense_shape is None
                     else [fetch.values, fetch.indices, fetch.dense_shape
                          ], _get_indexed_slices_value_from_fetches),

tensorflow/python/client/session_test.py

@@ -38,6 +38,7 @@
 from tensorflow.python.framework import errors
 from tensorflow.python.framework import function
 from tensorflow.python.framework import importer
+from tensorflow.python.framework import indexed_slices
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import sparse_tensor
 from tensorflow.python.framework import tensor_util
@@ -848,7 +849,7 @@ def testFetchIndexedSlices(self):
       indices = np.array([[3, 2, 0], [4, 5, 1]]).astype(np.int64)
       values = np.array([1.0, 2.0]).astype(np.float32)
       dense_shape = np.array([7, 9, 2]).astype(np.int64)
-      ind = ops.IndexedSlices(
+      ind = indexed_slices.IndexedSlices(
           constant_op.constant(values), constant_op.constant(indices),
           constant_op.constant(dense_shape))
       # Single fetch, use as tuple
@@ -883,15 +884,16 @@ def testFeedIndexedSlices(self):
       values = np.array([1.0, 2.0]).astype(np.float32)
       indices = np.array([[3, 2, 0], [4, 5, 1]]).astype(np.int64)
       dense_shape = np.array([7, 9, 2]).astype(np.int64)
-      ind = ops.IndexedSlices(
+      ind = indexed_slices.IndexedSlices(
           array_ops.placeholder(dtype=np.float32, shape=(2,)),
           array_ops.placeholder(dtype=np.int64, shape=(2, 3)),
           array_ops.placeholder(dtype=np.int64, shape=(3,)),
       )
       ind_values = array_ops.identity(ind.values)
       ind_indices = array_ops.identity(ind.indices)
       ind_dense_shape = array_ops.identity(ind.dense_shape)
-      ind2 = ops.IndexedSlices(ind_values, ind_indices, ind_dense_shape)
+      ind2 = indexed_slices.IndexedSlices(ind_values, ind_indices,
+                                          ind_dense_shape)
       # Feed with tuple
       values_out, indices_out, dense_shape_out = s.run(
           [ind_values, ind_indices, ind_dense_shape], {
@@ -901,16 +903,15 @@ def testFeedIndexedSlices(self):
       self.assertAllEqual(indices_out, indices)
       self.assertAllEqual(dense_shape_out, dense_shape)
       # Feed with IndexedSlicesValue
-      values_out, indices_out, dense_shape_out = s.run(
-          [ind_values, ind_indices, ind_dense_shape], {
-              ind: ops.IndexedSlicesValue(values, indices, dense_shape)
-          })
+      values_out, indices_out, dense_shape_out = s.run([
+          ind_values, ind_indices, ind_dense_shape
+      ], {ind: indexed_slices.IndexedSlicesValue(values, indices, dense_shape)})
       self.assertAllEqual(values_out, values)
       self.assertAllEqual(indices_out, indices)
       self.assertAllEqual(dense_shape_out, dense_shape)
       # Feed with IndexedSlicesValue, fetch IndexedSlicesValue
       ind2_out = s.run(ind2, {
-          ind: ops.IndexedSlicesValue(values, indices, dense_shape)
+          ind: indexed_slices.IndexedSlicesValue(values, indices, dense_shape)
       })
       self.assertAllEqual(ind2_out.values, values)
       self.assertAllEqual(ind2_out.indices, indices)
@@ -921,7 +922,7 @@ def testFetchIndexedSlicesWithoutDenseShape(self):
       indices = np.array([[3, 2, 0], [4, 5, 1]]).astype(np.int64)
       values = np.array([1.0, 2.0]).astype(np.float32)
       dense_shape = None
-      ind = ops.IndexedSlices(
+      ind = indexed_slices.IndexedSlices(
           constant_op.constant(values), constant_op.constant(indices), None)
       # Single fetch, use as tuple
       ind_out = s.run(ind)
@@ -955,12 +956,12 @@ def testFeedIndexedSlicesWithoutDenseShape(self):
       values = np.array([1.0, 2.0]).astype(np.float32)
       indices = np.array([[3, 2, 0], [4, 5, 1]]).astype(np.int64)
       dense_shape = None
-      ind = ops.IndexedSlices(
+      ind = indexed_slices.IndexedSlices(
           array_ops.placeholder(dtype=np.float32, shape=(2,)),
           array_ops.placeholder(dtype=np.int64, shape=(2, 3)), None)
       ind_values = array_ops.identity(ind.values)
       ind_indices = array_ops.identity(ind.indices)
-      ind2 = ops.IndexedSlices(ind_values, ind_indices)
+      ind2 = indexed_slices.IndexedSlices(ind_values, ind_indices)
       # Feed with tuple
       values_out, indices_out = s.run([ind_values, ind_indices], {
           ind: (values, indices)
@@ -969,13 +970,13 @@ def testFeedIndexedSlicesWithoutDenseShape(self):
       self.assertAllEqual(indices_out, indices)
       # Feed with IndexedSlicesValue
       values_out, indices_out = s.run([ind_values, ind_indices], {
-          ind: ops.IndexedSlicesValue(values, indices, dense_shape)
+          ind: indexed_slices.IndexedSlicesValue(values, indices, dense_shape)
       })
       self.assertAllEqual(values_out, values)
       self.assertAllEqual(indices_out, indices)
       # Feed with IndexedSlicesValue, fetch IndexedSlicesValue
       ind2_out = s.run(ind2, {
-          ind: ops.IndexedSlicesValue(values, indices, dense_shape)
+          ind: indexed_slices.IndexedSlicesValue(values, indices, dense_shape)
       })
       self.assertAllEqual(ind2_out.values, values)
       self.assertAllEqual(ind2_out.indices, indices)

tensorflow/python/distribute/BUILD

@@ -76,6 +76,7 @@ pytype_strict_library(
         "//tensorflow/python:resource_variable_ops",
         "//tensorflow/python/eager:backprop",
         "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:indexed_slices",
         "//tensorflow/python/framework:tensor_spec",
         "//tensorflow/python/types",
     ],

tensorflow/python/distribute/cross_device_ops.py

@@ -34,6 +34,7 @@
 from tensorflow.python.distribute import values_util
 from tensorflow.python.eager import context
 from tensorflow.python.eager import def_function
+from tensorflow.python.framework import indexed_slices
 from tensorflow.python.framework import kernels
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor_util
@@ -66,7 +67,7 @@ def validate_destinations(destinations):
   """Validates the `destination` is one of expected types."""
   if not isinstance(
       destinations,
-      (value_lib.DistributedValues, ops.Tensor, ops.IndexedSlices,
+      (value_lib.DistributedValues, ops.Tensor, indexed_slices.IndexedSlices,
        ps_values.AggregatingVariable, six.string_types,
        tpu_values.TPUMirroredVariable
       )) and not resource_variable_ops.is_resource_variable(destinations):
@@ -340,7 +341,7 @@ def _gather(self, per_replica_value, destinations, axis, options=None):
         `tf.distribute.DistributedValues` or if destinations is not a string,
         `tf.Variable` or `tf.distribute.DistributedValues`.
     """
-    if isinstance(per_replica_value, ops.IndexedSlices):
+    if isinstance(per_replica_value, indexed_slices.IndexedSlices):
       raise NotImplementedError("gather/all_gather does not support "
                                 "IndexedSlices")
     if options is None:
@@ -1175,7 +1176,7 @@ def _all_reduce(self, reduce_op, value, replica_id, options):
       if reduce_op == reduce_util.ReduceOp.MEAN:
         for i, v in enumerate(sparse_results):
           with ops.device(self._devices[replica_id]):
-            sparse_results[i] = ops.IndexedSlices(
+            sparse_results[i] = indexed_slices.IndexedSlices(
                 values=sparse_results[i].values / self._group_size,
                 indices=sparse_results[i].indices,
                 dense_shape=sparse_results[i].dense_shape)

tensorflow/python/distribute/cross_device_utils.py

@@ -23,6 +23,7 @@
 from tensorflow.python.eager import backprop
 from tensorflow.python.eager import context
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor_spec
 from tensorflow.python.ops import array_ops
@@ -495,8 +496,9 @@ def all_gather(
 
   def all_reduce_indexed_slices(
       self,
-      input_slices: ops.IndexedSlices,
-      options: Optional[collective_util.Options] = None) -> ops.IndexedSlices:
+      input_slices: indexed_slices.IndexedSlices,
+      options: Optional[collective_util.Options] = None
+  ) -> indexed_slices.IndexedSlices:
     """All-reduce an IndexedSlices.
 
     This method must be called inside a tf.function.
@@ -529,7 +531,7 @@ def all_reduce_indexed_slices(
       def all_gather_indexed_slices(
           all_gather_fn: Callable[
               [core.TensorLike, Optional[collective_util.Options]], core.Tensor]
-      ) -> ops.IndexedSlices:
+      ) -> indexed_slices.IndexedSlices:
         """Use all_gather_fn to aggregate `IndexedSlices`."""
         all_values = all_gather_fn(input_slices.values, options)
         # Add control dependency to order the all-gather.
@@ -540,7 +542,7 @@ def all_gather_indexed_slices(
           control = []
         with ops.control_dependencies(control):
           all_indices = all_gather_fn(input_slices.indices, options)
-        return ops.IndexedSlices(
+        return indexed_slices.IndexedSlices(
             values=all_values,
             indices=all_indices,
             dense_shape=input_slices.dense_shape)
@@ -572,42 +574,44 @@ def all_gather_with_padding(
 
 def aggregate_tensors_or_indexed_slices(values, accumulation_fn=math_ops.add_n):
   """Aggregate tensors using `accumulation_fn` and IndexedSlices via concat."""
-  if any(isinstance(v, ops.IndexedSlices) for v in values):
+  if any(isinstance(v, indexed_slices.IndexedSlices) for v in values):
     return backprop.aggregate_indexed_slices_gradients(values)
   else:
     return accumulation_fn(values)
 
 
 def divide_by_n_tensors_or_indexed_slices(value, n):
-  if isinstance(value, ops.IndexedSlices):
+  if isinstance(value, indexed_slices.IndexedSlices):
     value = backprop.flatten_nested_indexed_slices(value)
-    return ops.IndexedSlices(
-        value.values / n, value.indices, value.dense_shape)
+    return indexed_slices.IndexedSlices(value.values / n, value.indices,
+                                        value.dense_shape)
   else:
     return value / n
 
 
 def copy_tensor_or_indexed_slices_to_device(value, device):
   """Copies a tensor or IndexedSlices to a device."""
   with ops.device(device):
-    if isinstance(value, ops.IndexedSlices):
+    if isinstance(value, indexed_slices.IndexedSlices):
       copied_values = array_ops.identity(value.values)
       copied_indices = array_ops.identity(value.indices)
       if value.dense_shape is not None:
         copied_shape = array_ops.identity(value.dense_shape)
       else:
         copied_shape = None
-      result = ops.IndexedSlices(copied_values, copied_indices, copied_shape)
+      result = indexed_slices.IndexedSlices(copied_values, copied_indices,
+                                            copied_shape)
     else:
       result = array_ops.identity(value)
   return result
 
 
 def is_indexed_slices(value):
-  if isinstance(value, ops.IndexedSlices):
+  if isinstance(value, indexed_slices.IndexedSlices):
     return True
   if isinstance(value, value_lib.DistributedValues):
-    return all(isinstance(v, ops.IndexedSlices) for v in value.values)
+    return all(
+        isinstance(v, indexed_slices.IndexedSlices) for v in value.values)
   return False
 
 

tensorflow/python/distribute/cross_device_utils_test.py

@@ -22,6 +22,7 @@
 from tensorflow.python.eager import test
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
@@ -51,7 +52,7 @@ def testAggregateIndexedSlices(self):
         constant_op.constant([[0., 0.], [5, 6], [7., 8.]]))
     total = constant_op.constant([[1., 2.], [5, 6], [10., 12.]])
     result = cross_device_utils.aggregate_tensors_or_indexed_slices([t0, t1])
-    self.assertIsInstance(result, ops.IndexedSlices)
+    self.assertIsInstance(result, indexed_slices.IndexedSlices)
     self._assert_values_equal(total, result)
 
   @test_util.run_in_graph_and_eager_modes
@@ -69,7 +70,7 @@ def testDivideIndexedSlices(self):
     n = 2
     expected = constant_op.constant([[0.5, 1.], [0, 0], [1.5, 2.]])
     result = cross_device_utils.divide_by_n_tensors_or_indexed_slices(t, n)
-    self.assertIsInstance(result, ops.IndexedSlices)
+    self.assertIsInstance(result, indexed_slices.IndexedSlices)
     self._assert_values_equal(expected, result)
 
   @test_util.run_in_graph_and_eager_modes
@@ -103,7 +104,7 @@ def testCopyIndexedSlices(self):
     result = cross_device_utils.copy_tensor_or_indexed_slices_to_device(
         t, destination)
 
-    self.assertIsInstance(result, ops.IndexedSlices)
+    self.assertIsInstance(result, indexed_slices.IndexedSlices)
     self._assert_values_equal(t, result)
     self.assertEqual(
         device_util.resolve(destination), device_util.resolve(result.device))
@@ -112,13 +113,13 @@ def testCopyIndexedSlices(self):
       combinations.combine(mode=["graph", "eager"], required_gpus=1))
   def testCopyIndexedSlicesNoDenseShape(self):
     with ops.device("/cpu:0"):
-      t = ops.IndexedSlices(
+      t = indexed_slices.IndexedSlices(
           indices=array_ops.identity([0]), values=array_ops.identity([1.]))
     destination = "/gpu:0"
     result = cross_device_utils.copy_tensor_or_indexed_slices_to_device(
         t, destination)
 
-    self.assertIsInstance(result, ops.IndexedSlices)
+    self.assertIsInstance(result, indexed_slices.IndexedSlices)
     self.assertAllEqual(t.indices, result.indices)
     self.assertAllEqual(t.values, result.values)
     self.assertEqual(

tensorflow/python/distribute/distribute_lib.py

@@ -208,6 +208,7 @@
 from tensorflow.python.eager import monitoring
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor_shape
 from tensorflow.python.framework import tensor_util
@@ -1851,7 +1852,7 @@ def gather(self, value, axis):
                                                        error_message)
     dst = device_util.current(
     ) or self._extended._default_device or "/device:CPU:0"
-    if isinstance(value, ops.IndexedSlices):
+    if isinstance(value, indexed_slices.IndexedSlices):
       raise NotImplementedError("gather does not support IndexedSlices")
     return self._extended._local_results(
         self._extended._gather_to(value, dst, axis))[0]
@@ -3234,7 +3235,7 @@ def all_reduce(self, reduce_op, value, options=None):
     has_indexed_slices = False
 
     for v in flattened_value:
-      if isinstance(v, ops.IndexedSlices):
+      if isinstance(v, indexed_slices.IndexedSlices):
         has_indexed_slices = True
 
     if isinstance(reduce_op, six.string_types):
@@ -3418,7 +3419,7 @@ def all_gather(self, value, axis, options=None):
        is the same as `value`.
     """
     for v in nest.flatten(value):
-      if isinstance(v, ops.IndexedSlices):
+      if isinstance(v, indexed_slices.IndexedSlices):
         raise NotImplementedError("all_gather does not support IndexedSlices")
 
     if options is None:

tensorflow/python/distribute/sharded_variable.py

@@ -21,6 +21,7 @@
 from tensorflow.python.framework import composite_tensor
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices as indexed_slices_lib
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor_shape
 from tensorflow.python.framework import type_spec
@@ -596,7 +597,8 @@ def _decompose_indexed_slices(self, indexed_slices):
                                                      len(self._variables))
 
     return [
-        ops.IndexedSlices(values=per_var_values[i], indices=per_var_indices[i])
+        indexed_slices_lib.IndexedSlices(
+            values=per_var_values[i], indices=per_var_indices[i])
         for i in range(len(self._variables))
     ]