From 8031030fb0a36cef7b59f9c506609c179a2bebe8 Mon Sep 17 00:00:00 2001
From: Rishi Dave <rishipdave@gmail.com>
Date: Sun, 3 May 2026 12:01:30 +0000
Subject: [PATCH] fix: support Float16/BFloat16/Float8 types in TensorArray
 custom op

The TensorArray (Variadic) constructor's switch in
onnxruntime_lite_custom_op.h was missing case arms for Float16,
BFloat16, and the four Float8 variants, so any custom op accepting
const Ort::Custom::TensorArray& with those input types threw
"unknown input type" at construction.

Add the six missing case arms, mirroring the type-name parity already
established in CREATE_TUPLE and PARSE_ARGS within the same header.

Fixes #23373
---
 .../core/session/onnxruntime_lite_custom_op.h  | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/include/onnxruntime/core/session/onnxruntime_lite_custom_op.h b/include/onnxruntime/core/session/onnxruntime_lite_custom_op.h
index 81c20768c3120..2eb995d864ccd 100644
--- a/include/onnxruntime/core/session/onnxruntime_lite_custom_op.h
+++ b/include/onnxruntime/core/session/onnxruntime_lite_custom_op.h
@@ -330,6 +330,12 @@ struct TensorArray : public ArgBase {
           case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT:
             tensor = std::make_unique<Custom::Tensor<float>>(ctx, ith_input, true);
             break;
+          case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16:
+            tensor = std::make_unique<Custom::Tensor<Ort::Float16_t>>(ctx, ith_input, true);
+            break;
+          case ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16:
+            tensor = std::make_unique<Custom::Tensor<Ort::BFloat16_t>>(ctx, ith_input, true);
+            break;
           case ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE:
             tensor = std::make_unique<Custom::Tensor<double>>(ctx, ith_input, true);
             break;
@@ -360,6 +366,18 @@ struct TensorArray : public ArgBase {
           case ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING:
             tensor = std::make_unique<Custom::Tensor<std::string>>(ctx, ith_input, true);
             break;
+          case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT8E4M3FN:
+            tensor = std::make_unique<Custom::Tensor<Ort::Float8E4M3FN_t>>(ctx, ith_input, true);
+            break;
+          case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT8E4M3FNUZ:
+            tensor = std::make_unique<Custom::Tensor<Ort::Float8E4M3FNUZ_t>>(ctx, ith_input, true);
+            break;
+          case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT8E5M2:
+            tensor = std::make_unique<Custom::Tensor<Ort::Float8E5M2_t>>(ctx, ith_input, true);
+            break;
+          case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT8E5M2FNUZ:
+            tensor = std::make_unique<Custom::Tensor<Ort::Float8E5M2FNUZ_t>>(ctx, ith_input, true);
+            break;
           default:
             ORT_CXX_API_THROW("unknown input type", ORT_RUNTIME_EXCEPTION);
             break;