From b3caff096f8f747cf29f1007de09ffb3786066b9 Mon Sep 17 00:00:00 2001
From: David Majnemer <majnemer@google.com>
Date: Thu, 6 Aug 2020 10:55:33 -0700
Subject: [PATCH] Add half support to CrossReplicaSum

PiperOrigin-RevId: 325264977
Change-Id: I74d8c5667f8c89fc0c78641ab5be3576f5855c3f
---
 tensorflow/core/ops/tpu_cross_replica_ops.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/tpu_cross_replica_ops.cc b/tensorflow/core/ops/tpu_cross_replica_ops.cc
index adce0b51a05..1f10fe3136d 100644
--- a/tensorflow/core/ops/tpu_cross_replica_ops.cc
+++ b/tensorflow/core/ops/tpu_cross_replica_ops.cc
@@ -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")