From 6ad2c584524b2987bdafc39b547da2fae8addcb6 Mon Sep 17 00:00:00 2001
From: "Klimenko, Mikhail" <mikhail.klimenko@intel.com>
Date: Thu, 21 Aug 2025 14:58:17 +0200
Subject: [PATCH] Add Tensor__SetElementType to match TensorProto

---
 include/onnxruntime/core/framework/data_types.h            | 4 ++++
 include/onnxruntime/core/framework/tensor.h                | 7 +++++++
 .../core/providers/shared_library/provider_interfaces.h    | 1 +
 .../core/providers/shared_library/provider_wrappedtypes.h  | 1 +
 onnxruntime/core/session/provider_bridge_ort.cc            | 1 +
 5 files changed, 14 insertions(+)

diff --git a/include/onnxruntime/core/framework/data_types.h b/include/onnxruntime/core/framework/data_types.h
index d8822b3e452d5..45061af2d6f40 100644
--- a/include/onnxruntime/core/framework/data_types.h
+++ b/include/onnxruntime/core/framework/data_types.h
@@ -934,6 +934,10 @@ class PrimitiveDataTypeBase : public DataTypeImpl {
     return nullptr;
   }
 
+  void SetDataType(ONNX_NAMESPACE::TensorProto_DataType data_type) {
+    *const_cast<int32_t*>(&data_type_) = data_type;
+  }
+
   int32_t GetDataType() const {
     return data_type_;
   }
diff --git a/include/onnxruntime/core/framework/tensor.h b/include/onnxruntime/core/framework/tensor.h
index dd2603d214f63..8d7fa55b29439 100644
--- a/include/onnxruntime/core/framework/tensor.h
+++ b/include/onnxruntime/core/framework/tensor.h
@@ -161,6 +161,13 @@ class Tensor final {
   */
   MLDataType DataType() const { return dtype_; }
 
+  /**
+     Sets the data type to an enum constant
+  */
+  void SetElementType(ONNX_NAMESPACE::TensorProto_DataType data_type) {
+    const_cast<PrimitiveDataTypeBase*>(dtype_)->SetDataType(data_type);
+  }
+
   /**
      Returns the data type enum constant
      @remarks Use utils::ToTensorProtoElementType<T> for comparison.
diff --git a/onnxruntime/core/providers/shared_library/provider_interfaces.h b/onnxruntime/core/providers/shared_library/provider_interfaces.h
index 9a0bcb53c9ad7..2e7e221229d90 100644
--- a/onnxruntime/core/providers/shared_library/provider_interfaces.h
+++ b/onnxruntime/core/providers/shared_library/provider_interfaces.h
@@ -1290,6 +1290,7 @@ struct ProviderHost {
   virtual ptrdiff_t Tensor__ByteOffset(const Tensor* p) = 0;
   virtual size_t Tensor__SizeInBytes(const Tensor* p) = 0;
   virtual const OrtMemoryInfo& Tensor__Location(const Tensor* p) = 0;
+  virtual void Tensor__SetElementType(Tensor* p, ONNX_NAMESPACE::TensorProto_DataType data_type) = 0;
   virtual int32_t Tensor__GetElementType(const Tensor* p) = 0;
   virtual MLDataType Tensor__DataType(const Tensor* p) = 0;
 #ifdef ENABLE_STRIDED_TENSORS
diff --git a/onnxruntime/core/providers/shared_library/provider_wrappedtypes.h b/onnxruntime/core/providers/shared_library/provider_wrappedtypes.h
index 19b4636c3766d..f328b970b800a 100644
--- a/onnxruntime/core/providers/shared_library/provider_wrappedtypes.h
+++ b/onnxruntime/core/providers/shared_library/provider_wrappedtypes.h
@@ -1448,6 +1448,7 @@ struct Tensor final {
   const OrtMemoryInfo& Location() const { return g_host->Tensor__Location(this); }
 
   int32_t GetElementType() const { return g_host->Tensor__GetElementType(this); }
+  void SetElementType(ONNX_NAMESPACE::TensorProto_DataType data_type) { g_host->Tensor__SetElementType(this, data_type); }
   MLDataType DataType() const { return g_host->Tensor__DataType(this); }
   bool IsDataTypeString() const { return g_host->Tensor__IsDataTypeString(this); }
 
diff --git a/onnxruntime/core/session/provider_bridge_ort.cc b/onnxruntime/core/session/provider_bridge_ort.cc
index 41cf8be1d1412..c663d4b4a9a10 100644
--- a/onnxruntime/core/session/provider_bridge_ort.cc
+++ b/onnxruntime/core/session/provider_bridge_ort.cc
@@ -1651,6 +1651,7 @@ struct ProviderHostImpl : ProviderHost {
   ptrdiff_t Tensor__ByteOffset(const Tensor* p) override { return p->ByteOffset(); }
   size_t Tensor__SizeInBytes(const Tensor* p) override { return p->SizeInBytes(); }
   const OrtMemoryInfo& Tensor__Location(const Tensor* p) override { return p->Location(); }
+  void Tensor__SetElementType(Tensor* p, ONNX_NAMESPACE::TensorProto_DataType data_type) override { p->SetElementType(data_type); }
   int32_t Tensor__GetElementType(const Tensor* p) override { return p->GetElementType(); }
   MLDataType Tensor__DataType(const Tensor* p) override { return p->DataType(); }
 #ifdef ENABLE_STRIDED_TENSORS