From d8f8974e2d0797114e0205a3f17ed5511ad024cb Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <shoumikhin@meta.com>
Date: Tue, 25 Mar 2025 15:47:59 -0700
Subject: [PATCH] Enum for Tensor scalar type and shape dynamism.

Summary: .

Reviewed By: bsoyluoglu

Differential Revision: D71752748
---
 .../ExecuTorch/Exported/ExecuTorchTensor.h    | 51 +++++++++++++++++++
 1 file changed, 51 insertions(+)

diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h
index 220e377b60d..b840cd2faac 100644
--- a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h
+++ b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h
@@ -10,6 +10,57 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
+/**
+ * Enum to define the data type of a Tensor.
+ * Values can be a subset, but must numerically match exactly those defined in
+ * runtime/core/portable_type/scalar_type.h
+ */
+typedef NS_ENUM(int8_t, ExecuTorchDataType) {
+  ExecuTorchDataTypeByte,
+  ExecuTorchDataTypeChar,
+  ExecuTorchDataTypeShort,
+  ExecuTorchDataTypeInt,
+  ExecuTorchDataTypeLong,
+  ExecuTorchDataTypeHalf,
+  ExecuTorchDataTypeFloat,
+  ExecuTorchDataTypeDouble,
+  ExecuTorchDataTypeComplexHalf,
+  ExecuTorchDataTypeComplexFloat,
+  ExecuTorchDataTypeComplexDouble,
+  ExecuTorchDataTypeBool,
+  ExecuTorchDataTypeQInt8,
+  ExecuTorchDataTypeQUInt8,
+  ExecuTorchDataTypeQInt32,
+  ExecuTorchDataTypeBFloat16,
+  ExecuTorchDataTypeQUInt4x2,
+  ExecuTorchDataTypeQUInt2x4,
+  ExecuTorchDataTypeBits1x8,
+  ExecuTorchDataTypeBits2x4,
+  ExecuTorchDataTypeBits4x2,
+  ExecuTorchDataTypeBits8,
+  ExecuTorchDataTypeBits16,
+  ExecuTorchDataTypeFloat8_e5m2,
+  ExecuTorchDataTypeFloat8_e4m3fn,
+  ExecuTorchDataTypeFloat8_e5m2fnuz,
+  ExecuTorchDataTypeFloat8_e4m3fnuz,
+  ExecuTorchDataTypeUInt16,
+  ExecuTorchDataTypeUInt32,
+  ExecuTorchDataTypeUInt64,
+  ExecuTorchDataTypeUndefined,
+  ExecuTorchDataTypeNumOptions,
+} NS_SWIFT_NAME(DataType);
+
+/**
+ * Enum to define the shape dynamism of a Tensor.
+ * Values can be a subset, but must numerically match exactly those defined in
+ * runtime/core/tensor_shape_dynamism.h
+ */
+typedef NS_ENUM(uint8_t, ExecuTorchShapeDynamism) {
+  ExecuTorchShapeDynamismStatic,
+  ExecuTorchShapeDynamismDynamicBound,
+  ExecuTorchShapeDynamismDynamicUnbound,
+} NS_SWIFT_NAME(ShapeDynamism);
+
 NS_SWIFT_NAME(Tensor)
 __attribute__((deprecated("This API is experimental.")))
 @interface ExecuTorchTensor : NSObject