Tensor constructor to create with an array of scalars. (#9692)

Summary: #8366

Reviewed By: bsoyluoglu

Differential Revision: D71929714

Co-authored-by: Anthony Shoumikhin <[email protected]>

Author: kirklandsign

extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h

@@ -411,4 +411,29 @@ __attribute__((deprecated("This API is experimental.")))
 
 @end
 
+#pragma mark - Scalars Category
+
+@interface ExecuTorchTensor (Scalars)
+
+/**
+ * Initializes a tensor with an array of scalar values and full tensor properties.
+ *
+ * @param scalars An NSArray of NSNumber objects representing the scalar values.
+ * @param shape An NSArray of NSNumber objects representing the desired tensor shape.
+ * @param strides An NSArray of NSNumber objects representing the tensor 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 containing the provided scalar values.
+ */
+- (instancetype)initWithScalars:(NSArray<NSNumber *> *)scalars
+                          shape:(NSArray<NSNumber *> *)shape
+                        strides:(NSArray<NSNumber *> *)strides
+                 dimensionOrder:(NSArray<NSNumber *> *)dimensionOrder
+                       dataType:(ExecuTorchDataType)dataType
+                  shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism
+    NS_SWIFT_NAME(init(_:shape:strides:dimensionOrder:dataType:shapeDynamism:));
+
+@end
+
 NS_ASSUME_NONNULL_END

extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm

@@ -335,3 +335,36 @@ - (instancetype)initWithData:(NSData *)data
 }
 
 @end
+
+@implementation ExecuTorchTensor (Scalars)
+
+- (instancetype)initWithScalars:(NSArray<NSNumber *> *)scalars
+                          shape:(NSArray<NSNumber *> *)shape
+                        strides:(NSArray<NSNumber *> *)strides
+                 dimensionOrder:(NSArray<NSNumber *> *)dimensionOrder
+                       dataType:(ExecuTorchDataType)dataType
+                  shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism {
+  const NSInteger count = scalars.count;
+  ET_CHECK_MSG(count == ExecuTorchElementCountOfShape(shape),
+               "Number of scalars does not match the shape");
+  std::vector<uint8_t> data;
+  data.resize(count * ExecuTorchSizeOfDataType(dataType));
+  for (NSUInteger index = 0; index < count; ++index) {
+    ET_SWITCH_REALHBBF16_AND_UINT_TYPES(
+      static_cast<ScalarType>(dataType), nil, "initWithScalars", CTYPE, [&] {
+        reinterpret_cast<CTYPE *>(data.data())[index] = utils::extractValue<CTYPE>(scalars[index]);
+      }
+    );
+  }
+  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

extension/apple/ExecuTorch/__tests__/TensorTest.swift

@@ -223,4 +223,17 @@ class TensorTest: XCTestCase {
     XCTAssertFalse(tensor1.isEqual(NSString(string: "Not a tensor")))
     XCTAssertFalse(tensor4.isEqual(tensor2.copy()))
   }
+
+  func testInitScalarsFloat() {
+    let data: [Float] = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
+    let tensor = Tensor(data.map(NSNumber.init), shape: [2, 3], strides: [3, 1], dimensionOrder: [0, 1], dataType: .float, shapeDynamism: .dynamicBound)
+    XCTAssertEqual(tensor.dataType, .float)
+    XCTAssertEqual(tensor.shape, [2, 3])
+    XCTAssertEqual(tensor.strides, [3, 1])
+    XCTAssertEqual(tensor.dimensionOrder, [0, 1])
+    XCTAssertEqual(tensor.count, 6)
+    tensor.bytes { pointer, count, dataType in
+      XCTAssertEqual(Array(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Float.self), count: count)), data)
+    }
+  }
 }