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

Add half support to CrossReplicaSum

Browse Source 
PiperOrigin-RevId: 325264977
Change-Id: I74d8c5667f8c89fc0c78641ab5be3576f5855c3f
This commit is contained in:
David Majnemer 2020-08-06 10:55:33 -07:00 committed by TensorFlower Gardener
parent 916e5e54be
commit b3caff096f
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tensorflow/core/ops/tpu_cross_replica_ops.cc
View File

@ -78,7 +78,7 @@ REGISTER_OP("CrossReplicaSum")
.Input("input: T")
.Input("group_assignment: int32")
.Output("output: T")
.Attr("T: {bfloat16, float, int32, uint32}")
.Attr("T: {half, bfloat16, float, int32, uint32}")
.SetShapeFn(shape_inference::UnchangedShape);
REGISTER_OP("CollectivePermute")