Init Tensor with an array of scalar by sharing the memory.

Summary: .

Differential Revision: D77274533

Author: shoumikhin

extension/apple/ExecuTorch/Exported/ExecuTorch+Tensor.swift

@@ -693,6 +693,35 @@ public class Tensor<T: Scalar>: Equatable {
     ))
   }
 
+  /// Initializes a tensor without copying the data from an existing array.
+  ///
+  /// - Parameters:
+  ///   - scalars: An `inout` array of scalar values to share memory with.
+  ///   - shape: An array of integers representing the desired tensor shape. If empty, the shape is inferred as `[scalars.count]`.
+  ///   - strides: An array of integers representing the tensor strides.
+  ///   - dimensionOrder: An array of integers indicating the order of dimensions.
+  ///   - shapeDynamism: A `ShapeDynamism` value indicating the shape dynamism.
+  public convenience init(
+    _ scalars: inout [T],
+    shape: [Int] = [],
+    strides: [Int] = [],
+    dimensionOrder: [Int] = [],
+    shapeDynamism: ShapeDynamism = .dynamicBound
+  ) {
+    let newShape = shape.isEmpty ? [scalars.count] : shape
+    precondition(scalars.count == elementCount(ofShape: newShape))
+    self.init(scalars.withUnsafeMutableBufferPointer {
+      AnyTensor(
+        bytesNoCopy: $0.baseAddress!,
+        shape: newShape,
+        strides: strides,
+        dimensionOrder: dimensionOrder,
+        dataType: T.dataType,
+        shapeDynamism: shapeDynamism
+      )
+    })
+  }
+
   /// Initializes a tensor with an array of scalar values.
   ///
   /// - Parameters:

extension/apple/ExecuTorch/__tests__/TensorTest.swift

@@ -56,8 +56,8 @@ class TensorTest: XCTestCase {
 
   func testInitBytesNoCopy() {
     var data: [Float] = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
-    let tensor: Tensor<Float> = data.withUnsafeMutableBytes {
-      Tensor(bytesNoCopy: $0.baseAddress!, shape: [2, 3])
+    let tensor = data.withUnsafeMutableBytes {
+      Tensor<Float>(bytesNoCopy: $0.baseAddress!, shape: [2, 3])
     }
     // Modify the original data to make sure the tensor does not copy the data.
     data.indices.forEach { data[$0] += 1 }
@@ -73,8 +73,8 @@ class TensorTest: XCTestCase {
 
   func testInitBytes() {
     var data: [Double] = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
-    let tensor: Tensor<Double> = data.withUnsafeMutableBytes {
-      Tensor(bytes: $0.baseAddress!, shape: [2, 3])
+    let tensor = data.withUnsafeMutableBytes {
+      Tensor<Double>(bytes: $0.baseAddress!, shape: [2, 3])
     }
     // Modify the original data to make sure the tensor copies the data.
     data.indices.forEach { data[$0] += 1 }
@@ -91,14 +91,14 @@ class TensorTest: XCTestCase {
   func testInitData() {
     let dataArray: [Float] = [1.0, 2.0, 3.0, 4.0]
     let data = Data(bytes: dataArray, count: dataArray.count * MemoryLayout<Float>.size)
-    let tensor: Tensor<Float> = Tensor(data: data, shape: [4])
+    let tensor = Tensor<Float>(data: data, shape: [4])
     XCTAssertEqual(tensor.count, 4)
     XCTAssertEqual(try tensor.scalars(), dataArray)
   }
 
   func testWithCustomStridesAndDimensionOrder() {
     let data: [Float] = [1.0, 2.0, 3.0, 4.0]
-    let tensor: Tensor<Float> = Tensor(
+    let tensor = Tensor<Float>(
       bytes: data.withUnsafeBytes { $0.baseAddress! },
       shape: [2, 2],
       strides: [1, 2],
@@ -113,8 +113,8 @@ class TensorTest: XCTestCase {
 
   func testMutableBytes() {
     var data: [Int32] = [1, 2, 3, 4]
-    let tensor: Tensor<Int32> = data.withUnsafeMutableBytes {
-      Tensor(bytes: $0.baseAddress!, shape: [4])
+    let tensor = data.withUnsafeMutableBytes {
+      Tensor<Int32>(bytes: $0.baseAddress!, shape: [4])
     }
     XCTAssertNoThrow(try tensor.withUnsafeMutableBytes { buffer in
       for i in buffer.indices {
@@ -126,8 +126,8 @@ class TensorTest: XCTestCase {
 
   func testInitWithTensor() throws {
     var data: [Int] = [10, 20, 30, 40]
-    let tensor1: Tensor<Int> = data.withUnsafeMutableBytes {
-      Tensor(bytesNoCopy: $0.baseAddress!, shape: [2, 2])
+    let tensor1 = data.withUnsafeMutableBytes {
+      Tensor<Int>(bytesNoCopy: $0.baseAddress!, shape: [2, 2])
     }
     let tensor2 = Tensor(tensor1)
 
@@ -157,8 +157,8 @@ class TensorTest: XCTestCase {
 
   func testCopy() {
     var data: [Double] = [10.0, 20.0, 30.0, 40.0]
-    let tensor1: Tensor<Double> = data.withUnsafeMutableBytes {
-      Tensor(bytesNoCopy: $0.baseAddress!, shape: [2, 2])
+    let tensor1 = data.withUnsafeMutableBytes {
+      Tensor<Double>(bytesNoCopy: $0.baseAddress!, shape: [2, 2])
     }
     let tensor2 = tensor1.copy()
 
@@ -171,8 +171,8 @@ class TensorTest: XCTestCase {
 
   func testResize() {
     var data: [Int] = [1, 2, 3, 4]
-    let tensor: Tensor<Int> = data.withUnsafeMutableBytes {
-      Tensor(bytesNoCopy: $0.baseAddress!, shape: [4, 1])
+    let tensor = data.withUnsafeMutableBytes {
+      Tensor<Int>(bytesNoCopy: $0.baseAddress!, shape: [4, 1])
     }
     XCTAssertNoThrow(try tensor.resize(to: [2, 2]))
     XCTAssertEqual(tensor.dataType, .long)
@@ -185,38 +185,84 @@ class TensorTest: XCTestCase {
 
   func testResizeError() {
     var data: [Int] = [1, 2, 3, 4]
-    let tensor: Tensor<Int> = data.withUnsafeMutableBytes {
-      Tensor(bytesNoCopy: $0.baseAddress!, shape: [4, 1])
+    let tensor = data.withUnsafeMutableBytes {
+      Tensor<Int>(bytesNoCopy: $0.baseAddress!, shape: [4, 1])
     }
     XCTAssertThrowsError(try tensor.resize(to: [2, 3]))
   }
 
   func testIsEqual() {
     var data: [Float] = [1.0, 2.0, 3.0, 4.0]
-    let tensor1: Tensor<Float> = data.withUnsafeMutableBytes {
-      Tensor(bytesNoCopy: $0.baseAddress!, shape: [2, 2])
+    let tensor1 = data.withUnsafeMutableBytes {
+      Tensor<Float>(bytesNoCopy: $0.baseAddress!, shape: [2, 2])
     }
     let tensor2 = Tensor(tensor1)
     XCTAssertEqual(tensor1, tensor2)
     XCTAssertEqual(tensor2, tensor1)
 
     var dataModified: [Float] = [1.0, 2.0, 3.0, 5.0]
-    let tensor3: Tensor<Float> = dataModified.withUnsafeMutableBytes {
-      Tensor(bytesNoCopy: $0.baseAddress!, shape: [2, 2])
+    let tensor3 = dataModified.withUnsafeMutableBytes {
+      Tensor<Float>(bytesNoCopy: $0.baseAddress!, shape: [2, 2])
     }
     XCTAssertNotEqual(tensor1, tensor3)
-    let tensor4: Tensor<Float> = data.withUnsafeMutableBytes {
-      Tensor(bytesNoCopy: $0.baseAddress!, shape: [4, 1])
+    let tensor4 = data.withUnsafeMutableBytes {
+      Tensor<Float>(bytesNoCopy: $0.baseAddress!, shape: [4, 1])
     }
     XCTAssertNotEqual(tensor1, tensor4)
     XCTAssertEqual(tensor1, tensor1)
     XCTAssertNotEqual(tensor4, tensor2)
-    let tensor5: Tensor<Float> = data.withUnsafeMutableBytes {
-      Tensor(bytesNoCopy: $0.baseAddress!, shape: [2, 2], shapeDynamism: .static)
+    let tensor5 = data.withUnsafeMutableBytes {
+      Tensor<Float>(bytesNoCopy: $0.baseAddress!, shape: [2, 2], shapeDynamism: .static)
     }
     XCTAssertEqual(tensor1, tensor5)
   }
 
+  func testInitScalarsNoCopyDefaultShape() throws {
+    var data: [Float] = [1.0, 2.0, 3.0, 4.0]
+    let tensor = Tensor(&data)
+
+    XCTAssertEqual(tensor.dataType, .float)
+    XCTAssertEqual(tensor.shape, [4])
+    XCTAssertEqual(tensor.strides, [1])
+    XCTAssertEqual(tensor.dimensionOrder, [0])
+    XCTAssertEqual(tensor.shapeDynamism, .dynamicBound)
+    XCTAssertEqual(tensor.count, 4)
+    data[2] = 42.0
+    XCTAssertEqual(try tensor.scalars(), data)
+  }
+
+  func testInitScalarsNoCopyWithExplicitParams() throws {
+    var data: [Int] = [10, 20, 30, 40]
+    let tensor = Tensor(
+      &data,
+      shape: [2, 2],
+      strides: [2, 1],
+      dimensionOrder: [1, 0],
+      shapeDynamism: .static
+    )
+    XCTAssertEqual(tensor.dataType, .long)
+    XCTAssertEqual(tensor.shape, [2, 2])
+    XCTAssertEqual(tensor.strides, [2, 1])
+    XCTAssertEqual(tensor.dimensionOrder, [1, 0])
+    XCTAssertEqual(tensor.shapeDynamism, .static)
+    XCTAssertEqual(tensor.count, 4)
+    data = data.map { -$0 }
+    XCTAssertEqual(try tensor.scalars(), data)
+  }
+
+  func testInitScalarsBoolNoCopy() throws {
+    var data: [Bool] = [true, false, true]
+    let tensor = Tensor(&data)
+
+    XCTAssertEqual(tensor.dataType, .bool)
+    XCTAssertEqual(tensor.shape, [3])
+    XCTAssertEqual(tensor.strides, [1])
+    XCTAssertEqual(tensor.dimensionOrder, [0])
+    XCTAssertEqual(tensor.count, 3)
+    data[1].toggle()
+    XCTAssertEqual(try tensor.scalars(), data)
+  }
+
   func testInitScalarsUInt8() {
     let data: [UInt8] = [1, 2, 3, 4, 5, 6]
     let tensor = Tensor(data)