From e5617617fc8ae6c6c7ccf76b71be488c9a8d4574 Mon Sep 17 00:00:00 2001
From: Aditya Goel <48102515+adityagoel4512@users.noreply.github.com>
Date: Fri, 7 Apr 2023 00:05:36 +0100
Subject: [PATCH] Float to float label encoder (#15400)

---
 .../providers/cpu/cpu_execution_provider.cc   |  4 +++-
 .../core/providers/cpu/ml/label_encoder.cc    | 18 ++++++++++++++++
 .../providers/cpu/ml/label_encoder_test.cc    | 21 +++++++++++++++++++
 3 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc
index 3aeb6f651d..8af6b9cb30 100644
--- a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc
+++ b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc
@@ -2338,7 +2338,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMLDomain, 2,
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMLDomain, 2, string_int64, LabelEncoder);
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMLDomain, 2, int64_int64, LabelEncoder);
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMLDomain, 2, string_string, LabelEncoder);
-
+class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMLDomain, 2, float_float, LabelEncoder);
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMLDomain, 3, float, TreeEnsembleClassifier);
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMLDomain, 3, double, TreeEnsembleClassifier);
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMLDomain, 3, int64_t, TreeEnsembleClassifier);
@@ -2434,6 +2434,8 @@ Status RegisterOnnxMLOperatorKernels(KernelRegistry& kernel_registry) {
                                                                   LabelEncoder)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMLDomain, 2, string_string,
                                                                   LabelEncoder)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMLDomain, 2, float_float,
+                                                                  LabelEncoder)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMLDomain, 3, float,
                                                                   TreeEnsembleClassifier)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMLDomain, 3, double,
diff --git a/onnxruntime/core/providers/cpu/ml/label_encoder.cc b/onnxruntime/core/providers/cpu/ml/label_encoder.cc
index 655c60e3b2..55dbb62753 100644
--- a/onnxruntime/core/providers/cpu/ml/label_encoder.cc
+++ b/onnxruntime/core/providers/cpu/ml/label_encoder.cc
@@ -151,6 +151,24 @@ void LabelEncoder_2<std::string, std::string>::InitializeSomeFields(const OpKern
   info.GetAttrOrDefault<std::string>("default_string", &_default_value, std::string("_Unused"));
 };
 
+ONNX_CPU_OPERATOR_TYPED_ML_KERNEL(
+    LabelEncoder,
+    2,
+    float_float,
+    KernelDefBuilder().TypeConstraint("T1",
+                                      std::vector<MLDataType>{DataTypeImpl::GetTensorType<float>()})
+        .TypeConstraint("T2",
+                        std::vector<MLDataType>{DataTypeImpl::GetTensorType<float>()}),
+    LabelEncoder_2<float, float>)
+
+template <>
+void LabelEncoder_2<float, float>::InitializeSomeFields(const OpKernelInfo& info) {
+  _key_field_name = "keys_floats";
+  _value_field_name = "values_floats";
+  info.GetAttrOrDefault<float>("default_float", &_default_value, -0.0f);
+};
+
+
 ONNX_CPU_OPERATOR_TYPED_ML_KERNEL(
     LabelEncoder,
     2,
diff --git a/onnxruntime/test/providers/cpu/ml/label_encoder_test.cc b/onnxruntime/test/providers/cpu/ml/label_encoder_test.cc
index 4f04cd1d45..f1269a497c 100644
--- a/onnxruntime/test/providers/cpu/ml/label_encoder_test.cc
+++ b/onnxruntime/test/providers/cpu/ml/label_encoder_test.cc
@@ -211,5 +211,26 @@ TEST(LabelEncoder, StringToStringOpset2) {
   test.Run();
 }
 
+TEST(LabelEncoder, FloatToFloatOpset2) {
+  std::vector<std::int64_t> dims{1, 4};
+
+  std::vector<float> input{-1.0f, 0.0f, 3.1427f, 7.25f};
+  std::vector<float> output{1.0f, 0.0f, 2.718f, NAN};
+
+  OpTester test("LabelEncoder", 2, onnxruntime::kMLDomain);
+
+  const std::vector<float> keys{-1.0f, 0.0f, 7.25f};
+  const std::vector<float> values{1.0f, 0.0f, NAN};
+
+  test.AddAttribute("keys_floats", keys);
+  test.AddAttribute("values_floats", values);
+  test.AddAttribute("default_float", 2.718f);
+
+  test.AddInput<float>("X", dims, input);
+  test.AddOutput<float>("Y", dims, output);
+
+  test.Run();
+}
+
 }  // namespace test
 }  // namespace onnxruntime