Annotate ReduceDataset ops run on TPU _xla_compile_device_type='TPU' attribute.
Adding this attribute will denote to TF to decompose the ReduceDataset op into dataset iteration and a call to the reduce_fn.  The reduce_fn will be compiled for TPU.
@@ -117,6 +117,15 @@ const AcceleratorDevices& GetAcceleratorDevices() {
void SetDevice(absl::string_view device, tensorflow::GraphDef* graph_def,
const char* device_type) {
for (tensorflow::NodeDef& node_pb : *graph_def->mutable_node()) {
+ // Annotating ReduceDataset with _xla_compile_device_type will denote to
+ // Tensorflow to decompose the op into dataset iteration and reduce_fn and
+ // compile the reduce_fn.
+ // TODO(b/233627338): Currently only supported for TPU. Enable for CPU/GPU
+ // once MLIR-based TF2XLA bridge supports CPU/GPU.
+ if (node_pb.op() == "ReduceDataset" &&
+ !(strcmp(device_type, tensorflow::DEVICE_TPU))) {
+ (*node_pb.mutable_attr())["_xla_compile_device_type"].set_s(device_type);
+ }
if (!node_pb.device().empty()) {
VLOG(5) << "Skipping already placed node [" << node_pb.name() << "] ("
<< node_pb.op() << ") on " << node_pb.device();
