experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Support XlaGather and XlaScatter ops for bool.

Browse Source 
The undelying XLA ops do work for bool.

PiperOrigin-RevId: 337106941
Change-Id: Ifd920287f1113b9d6140489e8518613547289977
This commit is contained in:
A. Unique TensorFlower 2020-10-14 09:32:15 -07:00 committed by TensorFlower Gardener
parent 668c686e6d
commit 377ef73611
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tensorflow/compiler/tf2xla/ops/xla_ops.cc
View File

@ -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)