From e499e77c2497e3b106d533f7100201222912a5af Mon Sep 17 00:00:00 2001
From: xadupre <xadupre@microsoft.com>
Date: Mon, 6 May 2024 13:41:55 +0000
Subject: [PATCH] fix unit test

---
 operators/cuda/negxplus1.h       |  9 ++++-----
 operators/cuda/negxplus1_impl.cu |  2 +-
 test/cuda/test_cudaops.py        | 15 ++++++++++++---
 3 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/operators/cuda/negxplus1.h b/operators/cuda/negxplus1.h
index 5a0305412..832da6047 100644
--- a/operators/cuda/negxplus1.h
+++ b/operators/cuda/negxplus1.h
@@ -22,11 +22,10 @@ struct NegXPlus1 {
     if (0 == input_length) {
       return nullptr;
     }
-    using TT = typename CudaT<T>::MappedType;
-    LaunchNegXPlus1Kernel<TT>(reinterpret_cast<cudaStream_t>(ctx->GetCudaStream()),
-                              input_length,
-                              reinterpret_cast<const TT*>(input_data),
-                              reinterpret_cast<TT*>(output_data));
+    LaunchNegXPlus1Kernel<T>(reinterpret_cast<cudaStream_t>(ctx->GetCudaStream()),
+                             input_length,
+                             input_data,
+                             output_data);
     return nullptr;
   }
 };
diff --git a/operators/cuda/negxplus1_impl.cu b/operators/cuda/negxplus1_impl.cu
index 89641c9e0..1abb7ebb4 100644
--- a/operators/cuda/negxplus1_impl.cu
+++ b/operators/cuda/negxplus1_impl.cu
@@ -35,7 +35,7 @@ cudaError_t _LaunchNegXPlus1Kernel(cudaStream_t stream, int input_length, const
   constexpr int blockSize = 256;
   const int gridSize = (input_length + blockSize - 1) / blockSize;
   using TT = typename contrib::CudaT<T>::MappedType;
-  NegXPlus1Kernel<TT><<<gridSize, blockSize, 0, stream>>>((TT*)output, (const TT*)input, input_length);
+  NegXPlus1Kernel<TT><<<gridSize, blockSize, 0, stream>>>(reinterpret_cast<TT*>(output), reinterpret_cast<const TT*>(input), input_length);
   return cudaGetLastError();
 }
 
diff --git a/test/cuda/test_cudaops.py b/test/cuda/test_cudaops.py
index be4bedd4a..585c5c6cb 100644
--- a/test/cuda/test_cudaops.py
+++ b/test/cuda/test_cudaops.py
@@ -2,12 +2,21 @@
 import numpy as np
 from numpy.testing import assert_almost_equal
 from onnx import helper, numpy_helper, onnx_pb as onnx_proto, TensorProto
+from onnx.reference import ReferenceEvaluator
+from onnx.reference.op_run import OpRun
 from onnxruntime_extensions import make_onnx_model
 from onnxruntime_extensions import get_library_path as _get_library_path
 
 import onnxruntime as _ort
 
 
+class NegXPlus1(OpRun):
+    op_domain = "ai.onnx.contrib"
+
+    def _run(self, X):
+        return (1 - X,)
+
+
 class TestCudaOps(unittest.TestCase):
     @staticmethod
     def _create_negpos_test_model(domain="ai.onnx.contrib"):
@@ -109,7 +118,7 @@ def _negxplus1_cuda(self, itype):
 
         model2 = helper.make_model(
             helper.make_graph(
-                [helper.make_node("NegXplus1", ["X"], ["Y"], domain="ai.onnx.contrib")],
+                [helper.make_node("NegXPlus1", ["X"], ["Y"], domain="ai.onnx.contrib")],
                 "nd",
                 [helper.make_tensor_value_info("X", itype, [None, None, None])],
                 [helper.make_tensor_value_info("Y", itype, [None, None, None])],
@@ -125,14 +134,14 @@ def _negxplus1_cuda(self, itype):
         x = (np.arange(18) - 4).reshape((3, 2, 3)).astype(dtype)
 
         feeds1 = dict(X=x)
-        ref = CReferenceEvaluator(model1)
+        ref = ReferenceEvaluator(model1, new_ops=[NegXPlus1])
         expected = ref.run(None, feeds1)[0]
 
         opts = onnxruntime.SessionOptions()
         opts.register_custom_ops_library(_get_library_path())
         sess = onnxruntime.InferenceSession(model2.SerializeToString(), opts, providers=["CUDAExecutionProvider"])
         got = sess.run(None, feeds1)[0]
-        self.assertEqualArray(expected, got, atol=1e-5)
+        assert_almost_equal(expected, got, decimal=5)
 
     def test_negxplus1_cuda(self):
         self._negxplus1_cuda(TensorProto.FLOAT)