From 377ef7361177372d1d3efe536ed39da53d7d4f06 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 14 Oct 2020 09:32:15 -0700
Subject: [PATCH] Support XlaGather and XlaScatter ops for bool. The undelying
XLA ops do work for bool.
PiperOrigin-RevId: 337106941
Change-Id: Ifd920287f1113b9d6140489e8518613547289977
---
tensorflow/compiler/tf2xla/ops/xla_ops.cc | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/tensorflow/compiler/tf2xla/ops/xla_ops.cc b/tensorflow/compiler/tf2xla/ops/xla_ops.cc
index 00cbe7bc9bc..471cc029a59 100644
--- a/tensorflow/compiler/tf2xla/ops/xla_ops.cc
+++ b/tensorflow/compiler/tf2xla/ops/xla_ops.cc
@@ -792,7 +792,7 @@ REGISTER_OP("XlaGather")
.Input("slice_sizes: Tindices")
.Attr("dimension_numbers: string")
.Attr("indices_are_sorted: bool")
- .Attr("T: numbertype")
+ .Attr("T: {numbertype, bool}")
.Attr("Tindices: {int32, int64}")
.Output("output: T")
.SetShapeFn(shape_inference::UnknownShape)
@@ -813,7 +813,7 @@ REGISTER_OP("XlaScatter")
.Attr("update_computation: func")
.Attr("dimension_numbers: string")
.Attr("indices_are_sorted: bool")
- .Attr("T: numbertype")
+ .Attr("T: {numbertype, bool}")
.Attr("Tindices: {int32, int64}")
.Output("output: T")
.SetShapeFn(shape_inference::UnchangedShape)