[tf.data] Fix broken sparse/ragged iterators on TPU devices.

mrry · tensorflower-gardener · commit 4527fc0cc2ad · 2024-09-10T15:53:43.000-07:00
In a recent change, we added a colocation constraint between iterator ops and the sparse/ragged decoding ops that transform their outputs into structured tensors. This broke some workloads that attempted to prefetch SparseTensor and RaggedTensor objects to TPU memory, because the relevant decoding kernels were not registered for `DEVICE_TPU`. This change adds the missing kernel registrations, using the same host-memory annotations that are present for `DEVICE_GPU` and preserving the previous behavior (when the decoding ops would fall back to running on some `DEVICE_CPU`... although not necessarily in the same process).

PiperOrigin-RevId: 673113849
diff --git a/tensorflow/core/kernels/deserialize_sparse_variant_op.cc b/tensorflow/core/kernels/deserialize_sparse_variant_op.cc
@@ -375,6 +375,14 @@ REGISTER_KERNEL_BUILDER(Name("DeserializeSparse")
                             .HostMemory("sparse_values")
                             .HostMemory("sparse_shape"),
                         DeserializeSparseOp)
+REGISTER_KERNEL_BUILDER(Name("DeserializeSparse")
+                            .Device(DEVICE_TPU)
+                            .TypeConstraint<Variant>("Tserialized")
+                            .HostMemory("serialized_sparse")
+                            .HostMemory("sparse_indices")
+                            .HostMemory("sparse_values")
+                            .HostMemory("sparse_shape"),
+                        DeserializeSparseOp)
 
 }  // namespace
 
diff --git a/tensorflow/core/kernels/ragged_tensor_from_variant_op.cc b/tensorflow/core/kernels/ragged_tensor_from_variant_op.cc
@@ -367,6 +367,14 @@ class RaggedTensorFromVariantOp : public OpKernel {
                               .HostMemory("encoded_ragged")                  \
                               .HostMemory("output_nested_splits")            \
                               .HostMemory("output_dense_values"),            \
+                          RaggedTensorFromVariantOp<value_type, split_type>) \
+  REGISTER_KERNEL_BUILDER(Name("RaggedTensorFromVariant")                    \
+                              .Device(DEVICE_TPU)                            \
+                              .TypeConstraint<value_type>("Tvalues")         \
+                              .TypeConstraint<split_type>("Tsplits")         \
+                              .HostMemory("encoded_ragged")                  \
+                              .HostMemory("output_nested_splits")            \
+                              .HostMemory("output_dense_values"),            \
                           RaggedTensorFromVariantOp<value_type, split_type>);
 #define REGISTER_KERNELS(value_type)                  \
   REGISTER_KERNELS_WITH_SPLIT_TYPE(value_type, int32) \
