From 874385bc288c3e6c638adfee063c942daa9c3016 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 1 May 2025 12:44:53 -0700 Subject: [PATCH] Helpers to create empty tensors. Summary: https://github.com/pytorch/executorch/issues/8366 Differential Revision: D74020941 --- .../ExecuTorch/Exported/ExecuTorchTensor.h | 78 +++++++++++++++++++ .../ExecuTorch/Exported/ExecuTorchTensor.mm | 58 ++++++++++++++ .../ExecuTorch/__tests__/TensorTest.swift | 20 +++++ 3 files changed, 156 insertions(+) diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h index 45b6bd768cb..ebdfb8e2567 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h +++ b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h @@ -670,4 +670,82 @@ __attribute__((deprecated("This API is experimental."))) @end +#pragma mark - Empty Category + +@interface ExecuTorchTensor (Empty) + +/** + * Creates an empty tensor with the specified shape, strides, data type, and shape dynamism. + * + * @param shape An NSArray of NSNumber objects representing the desired shape. + * @param strides An NSArray of NSNumber objects representing the desired strides. + * @param dataType An ExecuTorchDataType value specifying the element type. + * @param shapeDynamism An ExecuTorchShapeDynamism value specifying whether the shape is static or dynamic. + * @return A new, empty ExecuTorchTensor instance. + */ ++ (instancetype)emptyTensorWithShape:(NSArray *)shape + strides:(NSArray *)strides + dataType:(ExecuTorchDataType)dataType + shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism + NS_SWIFT_NAME(empty(shape:strides:dataType:shapeDynamism:)); + +/** + * Creates an empty tensor with the specified shape, data type, and shape dynamism. + * + * @param shape An NSArray of NSNumber objects representing the desired shape. + * @param dataType An ExecuTorchDataType value specifying the element type. + * @param shapeDynamism An ExecuTorchShapeDynamism value specifying whether the shape is static or dynamic. + * @return A new, empty ExecuTorchTensor instance. + */ ++ (instancetype)emptyTensorWithShape:(NSArray *)shape + dataType:(ExecuTorchDataType)dataType + shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism + NS_SWIFT_NAME(empty(shape:dataType:shapeDynamism:)); + +/** + * Creates an empty tensor with the specified shape and data type, using dynamic bound shape. + * + * @param shape An NSArray of NSNumber objects representing the desired shape. + * @param dataType An ExecuTorchDataType value specifying the element type. + * @return A new, empty ExecuTorchTensor instance. + */ ++ (instancetype)emptyTensorWithShape:(NSArray *)shape + dataType:(ExecuTorchDataType)dataType + NS_SWIFT_NAME(empty(shape:dataType:)); + +/** + * Creates an empty tensor similar to the given tensor, with the specified data type and shape dynamism. + * + * @param tensor An existing ExecuTorchTensor instance whose shape and strides are used. + * @param dataType An ExecuTorchDataType value specifying the desired element type. + * @param shapeDynamism An ExecuTorchShapeDynamism value specifying whether the shape is static or dynamic. + * @return A new, empty ExecuTorchTensor instance with the same shape as the provided tensor. + */ ++ (instancetype)emptyTensorLikeTensor:(ExecuTorchTensor *)tensor + dataType:(ExecuTorchDataType)dataType + shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism + NS_SWIFT_NAME(empty(like:dataType:shapeDynamism:)); + +/** + * Creates an empty tensor similar to the given tensor, with the specified data type. + * + * @param tensor An existing ExecuTorchTensor instance whose shape and strides are used. + * @param dataType An ExecuTorchDataType value specifying the desired element type. + * @return A new, empty ExecuTorchTensor instance with the same shape as the provided tensor. + */ ++ (instancetype)emptyTensorLikeTensor:(ExecuTorchTensor *)tensor + dataType:(ExecuTorchDataType)dataType + NS_SWIFT_NAME(empty(like:dataType:)); + +/** + * Creates an empty tensor similar to the given tensor. + * + * @param tensor An existing ExecuTorchTensor instance. + * @return A new, empty ExecuTorchTensor instance with the same properties as the provided tensor. + */ ++ (instancetype)emptyTensorLikeTensor:(ExecuTorchTensor *)tensor + NS_SWIFT_NAME(empty(like:)); + +@end + NS_ASSUME_NONNULL_END diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm index acb88e2a8cc..d577f37a894 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm +++ b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm @@ -591,3 +591,61 @@ - (instancetype)initWithUnsignedInteger:(NSUInteger)scalar { } @end + +@implementation ExecuTorchTensor (Empty) + ++ (instancetype)emptyTensorWithShape:(NSArray *)shape + strides:(NSArray *)strides + dataType:(ExecuTorchDataType)dataType + shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism { + auto tensor = empty_strided( + utils::toVector(shape), + utils::toVector(strides), + static_cast(dataType), + static_cast(shapeDynamism) + ); + return [[self alloc] initWithNativeInstance:&tensor]; +} + ++ (instancetype)emptyTensorWithShape:(NSArray *)shape + dataType:(ExecuTorchDataType)dataType + shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism { + return [self emptyTensorWithShape:shape + strides:@[] + dataType:dataType + shapeDynamism:shapeDynamism]; +} + ++ (instancetype)emptyTensorWithShape:(NSArray *)shape + dataType:(ExecuTorchDataType)dataType { + return [self emptyTensorWithShape:shape + strides:@[] + dataType:dataType + shapeDynamism:ExecuTorchShapeDynamismDynamicBound]; +} + ++ (instancetype)emptyTensorLikeTensor:(ExecuTorchTensor *)tensor + dataType:(ExecuTorchDataType)dataType + shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism { + return [self emptyTensorWithShape:tensor.shape + strides:tensor.strides + dataType:dataType + shapeDynamism:shapeDynamism]; +} + ++ (instancetype)emptyTensorLikeTensor:(ExecuTorchTensor *)tensor + dataType:(ExecuTorchDataType)dataType { + return [self emptyTensorWithShape:tensor.shape + strides:tensor.strides + dataType:dataType + shapeDynamism:tensor.shapeDynamism]; +} + ++ (instancetype)emptyTensorLikeTensor:(ExecuTorchTensor *)tensor { + return [self emptyTensorWithShape:tensor.shape + strides:tensor.strides + dataType:tensor.dataType + shapeDynamism:tensor.shapeDynamism]; +} + +@end diff --git a/extension/apple/ExecuTorch/__tests__/TensorTest.swift b/extension/apple/ExecuTorch/__tests__/TensorTest.swift index c6bb1e2461f..3ee2c0f29be 100644 --- a/extension/apple/ExecuTorch/__tests__/TensorTest.swift +++ b/extension/apple/ExecuTorch/__tests__/TensorTest.swift @@ -548,4 +548,24 @@ class TensorTest: XCTestCase { XCTAssertEqual(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: UInt.self), count: count).first, 42) } } + + func testEmpty() { + let tensor = Tensor.empty(shape: [3, 4], dataType: .float) + XCTAssertEqual(tensor.shape, [3, 4]) + XCTAssertEqual(tensor.count, 12) + tensor.bytes { pointer, count, dataType in + XCTAssertNotNil(pointer) + XCTAssertEqual(count, 12) + XCTAssertEqual(dataType, .float) + } + } + + func testEmptyLike() { + let other = Tensor.empty(shape: [2, 2], dataType: .int) + let tensor = Tensor.empty(like: other) + XCTAssertEqual(tensor.shape, other.shape) + XCTAssertEqual(tensor.strides, other.strides) + XCTAssertEqual(tensor.dimensionOrder, other.dimensionOrder) + XCTAssertEqual(tensor.dataType, other.dataType) + } }