From 5d9047dbabcae3245c8674ffec1604e8c9ace6b5 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <shoumikhin@meta.com>
Date: Wed, 26 Mar 2025 17:12:28 -0700
Subject: [PATCH] Tensor properties to get shape, etc.

Summary: https://github.com/pytorch/executorch/issues/8366

Reviewed By: bsoyluoglu

Differential Revision: D71901794
---
 .../ExecuTorch/Exported/ExecuTorchTensor.h    | 42 +++++++++++++++++++
 .../ExecuTorch/Exported/ExecuTorchTensor.mm   | 37 ++++++++++++++++
 2 files changed, 79 insertions(+)

diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h
index 6bd9854526c..3d54ac5977b 100644
--- a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h
+++ b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h
@@ -78,6 +78,48 @@ __attribute__((deprecated("This API is experimental.")))
  */
 @property(nonatomic, readonly) void *nativeInstance NS_SWIFT_UNAVAILABLE("");
 
+/**
+ * The data type of the tensor.
+ *
+ * @return An ExecuTorchDataType value representing the tensor's element type.
+ */
+@property(nonatomic, readonly) ExecuTorchDataType dataType;
+
+/**
+ * The shape of the tensor.
+ *
+ * @return An NSArray of NSNumber objects representing the size of each dimension.
+ */
+@property(nonatomic, readonly) NSArray<NSNumber *> *shape;
+
+/**
+ * The order of dimensions in the tensor.
+ *
+ * @return An NSArray of NSNumber objects representing the tensor’s dimension order.
+ */
+@property(nonatomic, readonly) NSArray<NSNumber *> *dimensionOrder;
+
+/**
+ * The strides of the tensor.
+ *
+ * @return An NSArray of NSNumber objects representing the step sizes for each dimension.
+ */
+@property(nonatomic, readonly) NSArray<NSNumber *> *strides;
+
+/**
+ * The dynamism of the tensor's shape.
+ *
+ * @return An ExecuTorchShapeDynamism value indicating whether the tensor shape is static or dynamic.
+ */
+@property(nonatomic, readonly) ExecuTorchShapeDynamism shapeDynamism;
+
+/**
+ * The total number of elements in the tensor.
+ *
+ * @return An NSInteger representing the total element count.
+ */
+@property(nonatomic, readonly) NSInteger count;
+
 /**
  * Initializes a tensor with a native TensorPtr instance.
  *
diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm
index ef93b2eb842..ba4fe331cad 100644
--- a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm
+++ b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm
@@ -9,6 +9,7 @@
 #import "ExecuTorchTensor.h"
 
 #import "ExecuTorchError.h"
+#import "ExecuTorchUtils.h"
 
 #import <executorch/extension/tensor/tensor.h>
 
@@ -16,6 +17,9 @@
 
 @implementation ExecuTorchTensor {
   TensorPtr _tensor;
+  NSArray<NSNumber *> *_shape;
+  NSArray<NSNumber *> *_strides;
+  NSArray<NSNumber *> *_dimensionOrder;
 }
 
 - (instancetype)initWithNativeInstance:(void *)nativeInstance {
@@ -31,4 +35,37 @@ - (void *)nativeInstance {
   return &_tensor;
 }
 
+- (ExecuTorchDataType)dataType {
+  return static_cast<ExecuTorchDataType>(_tensor->scalar_type());
+}
+
+- (NSArray<NSNumber *> *)shape {
+  if (!_shape) {
+    _shape = utils::toNSArray(_tensor->sizes());
+  }
+  return _shape;
+}
+
+- (NSArray<NSNumber *> *)dimensionOrder {
+  if (!_dimensionOrder) {
+    _dimensionOrder = utils::toNSArray(_tensor->dim_order());
+  }
+  return _dimensionOrder;
+}
+
+- (NSArray<NSNumber *> *)strides {
+  if (!_strides) {
+    _strides = utils::toNSArray(_tensor->strides());
+  }
+  return _strides;
+}
+
+- (ExecuTorchShapeDynamism)shapeDynamism {
+  return static_cast<ExecuTorchShapeDynamism>(_tensor->shape_dynamism());
+}
+
+- (NSInteger)count {
+  return _tensor->numel();
+}
+
 @end