diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm index 3b1c06a5aa0..24d355782c0 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm +++ b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm @@ -131,7 +131,11 @@ - (instancetype)initWithTensor:(ExecuTorchTensor *)otherTensor { auto tensor = make_tensor_ptr( *reinterpret_cast(otherTensor.nativeInstance) ); - return [self initWithNativeInstance:&tensor]; + self = [self initWithNativeInstance:&tensor]; + if (self) { + _data = otherTensor->_data; + } + return self; } - (instancetype)copy { diff --git a/extension/apple/ExecuTorch/__tests__/TensorTest.swift b/extension/apple/ExecuTorch/__tests__/TensorTest.swift index 52cd3421d6b..e2c330ac1c2 100644 --- a/extension/apple/ExecuTorch/__tests__/TensorTest.swift +++ b/extension/apple/ExecuTorch/__tests__/TensorTest.swift @@ -96,6 +96,19 @@ class TensorTest: XCTestCase { XCTAssertEqual(tensor.scalars(), dataArray) } + func testInitDataViewSurvivesSourceScopeEnd() { + let dataArray: [Float] = [1.0, 2.0, 3.0, 4.0] + var view: Tensor! + autoreleasepool { + let data = Data(bytes: dataArray, count: dataArray.count * MemoryLayout.size) + let tensor = Tensor(data: data, shape: [4]) + view = Tensor(tensor) + XCTAssertEqual(view.scalars(), dataArray) + } + XCTAssertEqual(view.count, 4) + XCTAssertEqual(view.scalars(), dataArray) + } + func testWithCustomStridesAndDimensionOrder() { let data: [Float] = [1.0, 2.0, 3.0, 4.0] let tensor = Tensor(