From e10c84e0b38b7a5aeda2a85e479e6b32330b3efb Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <shoumikhin@meta.com>
Date: Wed, 26 Mar 2025 17:42:00 -0700
Subject: [PATCH] Tensor constructor to create with a raw pointer by copying
 the data.

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

Reviewed By: mergennachin

Differential Revision: D71904351
---
 .../ExecuTorch/Exported/ExecuTorchTensor.h    | 24 ++++++++++++++++++
 .../ExecuTorch/Exported/ExecuTorchTensor.mm   | 25 +++++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h
index 16b1ec392ec..fadbf137fad 100644
--- a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h
+++ b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h
@@ -222,4 +222,28 @@ __attribute__((deprecated("This API is experimental.")))
 
 @end
 
+#pragma mark - Bytes Category
+
+@interface ExecuTorchTensor (Bytes)
+
+/**
+ * Initializes a tensor by copying bytes from the provided pointer.
+ *
+ * @param pointer A pointer to the source data buffer.
+ * @param shape An NSArray of NSNumber objects representing the tensor's shape.
+ * @param strides An NSArray of NSNumber objects representing the tensor's strides.
+ * @param dimensionOrder An NSArray of NSNumber objects indicating the order of dimensions.
+ * @param dataType An ExecuTorchDataType value specifying the element type.
+ * @param shapeDynamism An ExecuTorchShapeDynamism value indicating the shape dynamism.
+ * @return An initialized ExecuTorchTensor instance with its own copy of the data.
+ */
+- (instancetype)initWithBytes:(const void *)pointer
+                        shape:(NSArray<NSNumber *> *)shape
+                      strides:(NSArray<NSNumber *> *)strides
+               dimensionOrder:(NSArray<NSNumber *> *)dimensionOrder
+                     dataType:(ExecuTorchDataType)dataType
+                shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism;
+
+@end
+
 NS_ASSUME_NONNULL_END
diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm
index 19851172b12..7f02320d2db 100644
--- a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm
+++ b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm
@@ -140,3 +140,28 @@ - (instancetype)initWithBytesNoCopy:(void *)pointer
 }
 
 @end
+
+@implementation ExecuTorchTensor (Bytes)
+
+- (instancetype)initWithBytes:(const void *)pointer
+                        shape:(NSArray<NSNumber *> *)shape
+                      strides:(NSArray<NSNumber *> *)strides
+               dimensionOrder:(NSArray<NSNumber *> *)dimensionOrder
+                     dataType:(ExecuTorchDataType)dataType
+                shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism {
+  ET_CHECK(pointer);
+  const auto size = ExecuTorchElementCountOfShape(shape) * ExecuTorchSizeOfDataType(dataType);
+  std::vector<uint8_t> data(static_cast<const uint8_t *>(pointer),
+                            static_cast<const uint8_t *>(pointer) + size);
+  auto tensor = make_tensor_ptr(
+    utils::toVector<SizesType>(shape),
+    std::move(data),
+    utils::toVector<DimOrderType>(dimensionOrder),
+    utils::toVector<StridesType>(strides),
+    static_cast<ScalarType>(dataType),
+    static_cast<TensorShapeDynamism>(shapeDynamism)
+  );
+  return [self initWithNativeInstance:&tensor];
+}
+
+@end