diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h index 53193dc1f5c..52ae4277c36 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h +++ b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h @@ -980,4 +980,87 @@ __attribute__((deprecated("This API is experimental."))) @end +#pragma mark - RandomNormal Category + +@interface ExecuTorchTensor (RandomNormal) + +/** + * Creates a tensor with random values drawn from a normal distribution, + * with full specification of 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 ExecuTorchTensor instance filled with values from a normal distribution. + */ ++ (instancetype)randomNormalTensorWithShape:(NSArray *)shape + strides:(NSArray *)strides + dataType:(ExecuTorchDataType)dataType + shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism + NS_SWIFT_NAME(randn(shape:strides:dataType:shapeDynamism:)); + +/** + * Creates a tensor with random values drawn from a normal distribution, + * with the specified shape and data type. + * + * @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 ExecuTorchTensor instance filled with values from a normal distribution. + */ ++ (instancetype)randomNormalTensorWithShape:(NSArray *)shape + dataType:(ExecuTorchDataType)dataType + shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism + NS_SWIFT_NAME(randn(shape:dataType:shapeDynamism:)); + +/** + * Creates a tensor with random values drawn from a normal distribution, + * with the specified shape (using dynamic bound shape) and data type. + * + * @param shape An NSArray of NSNumber objects representing the desired shape. + * @param dataType An ExecuTorchDataType value specifying the element type. + * @return A new ExecuTorchTensor instance filled with values from a normal distribution. + */ ++ (instancetype)randomNormalTensorWithShape:(NSArray *)shape + dataType:(ExecuTorchDataType)dataType + NS_SWIFT_NAME(randn(shape:dataType:)); + +/** + * Creates a tensor with random normal values similar to an existing 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 ExecuTorchTensor instance filled with values from a normal distribution. + */ ++ (instancetype)randomNormalTensorLikeTensor:(ExecuTorchTensor *)tensor + dataType:(ExecuTorchDataType)dataType + shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism + NS_SWIFT_NAME(randn(like:dataType:shapeDynamism:)); + +/** + * Creates a tensor with random normal values similar to an existing 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 ExecuTorchTensor instance filled with values from a normal distribution. + */ ++ (instancetype)randomNormalTensorLikeTensor:(ExecuTorchTensor *)tensor + dataType:(ExecuTorchDataType)dataType + NS_SWIFT_NAME(randn(like:dataType:)); + +/** + * Creates a tensor with random normal values similar to an existing tensor. + * + * @param tensor An existing ExecuTorchTensor instance. + * @return A new ExecuTorchTensor instance filled with values from a normal distribution. + */ ++ (instancetype)randomNormalTensorLikeTensor:(ExecuTorchTensor *)tensor + NS_SWIFT_NAME(randn(like:)); + +@end + NS_ASSUME_NONNULL_END diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm index 494618eb883..634a62c46f2 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm +++ b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm @@ -833,3 +833,61 @@ + (instancetype)randomTensorLikeTensor:(ExecuTorchTensor *)tensor { } @end + +@implementation ExecuTorchTensor (RandomNormal) + ++ (instancetype)randomNormalTensorWithShape:(NSArray *)shape + strides:(NSArray *)strides + dataType:(ExecuTorchDataType)dataType + shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism { + auto tensor = randn_strided( + utils::toVector(shape), + utils::toVector(strides), + static_cast(dataType), + static_cast(shapeDynamism) + ); + return [[self alloc] initWithNativeInstance:&tensor]; +} + ++ (instancetype)randomNormalTensorWithShape:(NSArray *)shape + dataType:(ExecuTorchDataType)dataType + shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism { + return [self randomNormalTensorWithShape:shape + strides:@[] + dataType:dataType + shapeDynamism:shapeDynamism]; +} + ++ (instancetype)randomNormalTensorWithShape:(NSArray *)shape + dataType:(ExecuTorchDataType)dataType { + return [self randomNormalTensorWithShape:shape + strides:@[] + dataType:dataType + shapeDynamism:ExecuTorchShapeDynamismDynamicBound]; +} + ++ (instancetype)randomNormalTensorLikeTensor:(ExecuTorchTensor *)tensor + dataType:(ExecuTorchDataType)dataType + shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism { + return [self randomNormalTensorWithShape:tensor.shape + strides:tensor.strides + dataType:dataType + shapeDynamism:shapeDynamism]; +} + ++ (instancetype)randomNormalTensorLikeTensor:(ExecuTorchTensor *)tensor + dataType:(ExecuTorchDataType)dataType { + return [self randomNormalTensorWithShape:tensor.shape + strides:tensor.strides + dataType:dataType + shapeDynamism:tensor.shapeDynamism]; +} + ++ (instancetype)randomNormalTensorLikeTensor:(ExecuTorchTensor *)tensor { + return [self randomNormalTensorWithShape: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 468ec85d261..27a6cac7b1e 100644 --- a/extension/apple/ExecuTorch/__tests__/TensorTest.swift +++ b/extension/apple/ExecuTorch/__tests__/TensorTest.swift @@ -637,4 +637,22 @@ class TensorTest: XCTestCase { XCTAssertEqual(tensor.shape, other.shape) XCTAssertEqual(tensor.count, other.count) } + + func testRandomNormal() { + let tensor = Tensor.randn(shape: [4], dataType: .double) + XCTAssertEqual(tensor.shape, [4]) + XCTAssertEqual(tensor.count, 4) + tensor.bytes { pointer, count, dataType in + XCTAssertEqual(dataType, .double) + let buffer = UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Double.self), count: count) + XCTAssertEqual(buffer.count, 4) + } + } + + func testRandomNormalLike() { + let other = Tensor.zeros(shape: [4], dataType: .float) + let tensor = Tensor.randn(like: other) + XCTAssertEqual(tensor.shape, other.shape) + XCTAssertEqual(tensor.count, other.count) + } }