You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
[tf.data] Fix broken sparse/ragged iterators on TPU devices.
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
0 commit comments