From 3a9a1a4ef198dc81eed2e9378ccd75994941545a Mon Sep 17 00:00:00 2001
From: Sergii Dymchenko <sedymche@microsoft.com>
Date: Fri, 9 Oct 2020 11:57:50 -0700
Subject: [PATCH] Fix registration for GatherGrad (#5382)

* Fix registration for GatherGrad to fix GatherGradOpTest.GatherGrad_axis0_indices2d_half.

* Fix GatherGrad registration for CUDA also.
---
 orttraining/orttraining/core/graph/training_op_defs.cc         | 2 +-
 orttraining/orttraining/training_ops/cpu/tensor/gather_grad.cc | 3 ++-
 .../orttraining/training_ops/cuda/tensor/gather_grad.cc        | 3 ++-
 3 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/orttraining/orttraining/core/graph/training_op_defs.cc b/orttraining/orttraining/core/graph/training_op_defs.cc
index 7af242a848..6ec1b4d3e4 100644
--- a/orttraining/orttraining/core/graph/training_op_defs.cc
+++ b/orttraining/orttraining/core/graph/training_op_defs.cc
@@ -446,7 +446,7 @@ void RegisterTrainingOpSchemas() {
           "Constrain input shape to integer tensors.")
       .TypeConstraint(
           "T",
-          OpSchema::all_tensor_types_with_bfloat(),
+          {"tensor(float16)", "tensor(float)", "tensor(double)"},
           "Constrain input and output types to float tensors.")
       .TypeConstraint(
           "Tind",
diff --git a/orttraining/orttraining/training_ops/cpu/tensor/gather_grad.cc b/orttraining/orttraining/training_ops/cpu/tensor/gather_grad.cc
index 5f37bf35ce..7607b13788 100644
--- a/orttraining/orttraining/training_ops/cpu/tensor/gather_grad.cc
+++ b/orttraining/orttraining/training_ops/cpu/tensor/gather_grad.cc
@@ -15,7 +15,8 @@ ONNX_OPERATOR_KERNEL_EX(
     kCpuExecutionProvider,
     KernelDefBuilder()
         .TypeConstraint("I", DataTypeImpl::GetTensorType<int64_t>())
-        .TypeConstraint("T", DataTypeImpl::AllTensorTypes())
+        .TypeConstraint("T", {DataTypeImpl::GetTensorType<float>(),
+                              DataTypeImpl::GetTensorType<double>()})
         .TypeConstraint("Tind", std::vector<MLDataType>{
                                     DataTypeImpl::GetTensorType<int32_t>(),
                                     DataTypeImpl::GetTensorType<int64_t>()}),
diff --git a/orttraining/orttraining/training_ops/cuda/tensor/gather_grad.cc b/orttraining/orttraining/training_ops/cuda/tensor/gather_grad.cc
index 1c2a6a662d..ac2049f9ba 100644
--- a/orttraining/orttraining/training_ops/cuda/tensor/gather_grad.cc
+++ b/orttraining/orttraining/training_ops/cuda/tensor/gather_grad.cc
@@ -16,7 +16,8 @@ ONNX_OPERATOR_KERNEL_EX(
     KernelDefBuilder()
         .InputMemoryType<OrtMemTypeCPUInput>(0)
         .TypeConstraint("I", DataTypeImpl::GetTensorType<int64_t>())
-        .TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes())
+        .TypeConstraint("T", {DataTypeImpl::GetTensorType<float>(),
+                              DataTypeImpl::GetTensorType<MLFloat16>()})
         .TypeConstraint("Tind", std::vector<MLDataType>{
                                     DataTypeImpl::GetTensorType<int32_t>(),
                                     DataTypeImpl::GetTensorType<int64_t>()}),