diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h index 52ae4277c36..67904e97b31 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h +++ b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h @@ -1063,4 +1063,111 @@ __attribute__((deprecated("This API is experimental."))) @end +#pragma mark - RandomInteger Category + +@interface ExecuTorchTensor (RandomInteger) + +/** + * Creates a tensor with random integer values in the specified range, + * with full specification of shape, strides, data type, and shape dynamism. + * + * @param low An NSInteger specifying the inclusive lower bound of random values. + * @param high An NSInteger specifying the exclusive upper bound of random values. + * @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 random integer values. + */ ++ (instancetype)randomIntegerTensorWithLow:(NSInteger)low + high:(NSInteger)high + shape:(NSArray *)shape + strides:(NSArray *)strides + dataType:(ExecuTorchDataType)dataType + shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism + NS_SWIFT_NAME(randint(low:high:shape:strides:dataType:shapeDynamism:)); + +/** + * Creates a tensor with random integer values in the specified range, + * with the given shape and data type. + * + * @param low An NSInteger specifying the inclusive lower bound of random values. + * @param high An NSInteger specifying the exclusive upper bound of random values. + * @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 random integer values. + */ ++ (instancetype)randomIntegerTensorWithLow:(NSInteger)low + high:(NSInteger)high + shape:(NSArray *)shape + dataType:(ExecuTorchDataType)dataType + shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism + NS_SWIFT_NAME(randint(low:high:shape:dataType:shapeDynamism:)); + +/** + * Creates a tensor with random integer values in the specified range, + * with the given shape (using dynamic bound shape) and data type. + * + * @param low An NSInteger specifying the inclusive lower bound of random values. + * @param high An NSInteger specifying the exclusive upper bound of random values. + * @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 random integer values. + */ ++ (instancetype)randomIntegerTensorWithLow:(NSInteger)low + high:(NSInteger)high + shape:(NSArray *)shape + dataType:(ExecuTorchDataType)dataType + NS_SWIFT_NAME(randint(low:high:shape:dataType:)); + +/** + * Creates a tensor with random integer values in the specified range, similar to an existing tensor, + * with the given data type and shape dynamism. + * + * @param tensor An existing ExecuTorchTensor instance whose shape and strides are used. + * @param low An NSInteger specifying the inclusive lower bound of random values. + * @param high An NSInteger specifying the exclusive upper bound of random values. + * @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 random integer values. + */ ++ (instancetype)randomIntegerTensorLikeTensor:(ExecuTorchTensor *)tensor + low:(NSInteger)low + high:(NSInteger)high + dataType:(ExecuTorchDataType)dataType + shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism + NS_SWIFT_NAME(randint(like:low:high:dataType:shapeDynamism:)); + +/** + * Creates a tensor with random integer values in the specified range, similar to an existing tensor, + * with the given data type. + * + * @param tensor An existing ExecuTorchTensor instance whose shape and strides are used. + * @param low An NSInteger specifying the inclusive lower bound of random values. + * @param high An NSInteger specifying the exclusive upper bound of random values. + * @param dataType An ExecuTorchDataType value specifying the element type. + * @return A new ExecuTorchTensor instance filled with random integer values. + */ ++ (instancetype)randomIntegerTensorLikeTensor:(ExecuTorchTensor *)tensor + low:(NSInteger)low + high:(NSInteger)high + dataType:(ExecuTorchDataType)dataType + NS_SWIFT_NAME(randint(like:low:high:dataType:)); + +/** + * Creates a tensor with random integer values in the specified range, similar to an existing tensor. + * + * @param tensor An existing ExecuTorchTensor instance. + * @param low An NSInteger specifying the inclusive lower bound of random values. + * @param high An NSInteger specifying the exclusive upper bound of random values. + * @return A new ExecuTorchTensor instance filled with random integer values. + */ ++ (instancetype)randomIntegerTensorLikeTensor:(ExecuTorchTensor *)tensor + low:(NSInteger)low + high:(NSInteger)high + NS_SWIFT_NAME(randint(like:low:high:)); + +@end + NS_ASSUME_NONNULL_END diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm index 634a62c46f2..663ade54457 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm +++ b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm @@ -891,3 +891,85 @@ + (instancetype)randomNormalTensorLikeTensor:(ExecuTorchTensor *)tensor { } @end + +@implementation ExecuTorchTensor (RandomInteger) + ++ (instancetype)randomIntegerTensorWithLow:(NSInteger)low + high:(NSInteger)high + shape:(NSArray *)shape + strides:(NSArray *)strides + dataType:(ExecuTorchDataType)dataType + shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism { + auto tensor = randint_strided( + low, + high, + utils::toVector(shape), + utils::toVector(strides), + static_cast(dataType), + static_cast(shapeDynamism) + ); + return [[self alloc] initWithNativeInstance:&tensor]; +} + ++ (instancetype)randomIntegerTensorWithLow:(NSInteger)low + high:(NSInteger)high + shape:(NSArray *)shape + dataType:(ExecuTorchDataType)dataType + shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism { + return [self randomIntegerTensorWithLow:low + high:high + shape:shape + strides:@[] + dataType:dataType + shapeDynamism:shapeDynamism]; +} + ++ (instancetype)randomIntegerTensorWithLow:(NSInteger)low + high:(NSInteger)high + shape:(NSArray *)shape + dataType:(ExecuTorchDataType)dataType { + return [self randomIntegerTensorWithLow:low + high:high + shape:shape + strides:@[] + dataType:dataType + shapeDynamism:ExecuTorchShapeDynamismDynamicBound]; +} + ++ (instancetype)randomIntegerTensorLikeTensor:(ExecuTorchTensor *)tensor + low:(NSInteger)low + high:(NSInteger)high + dataType:(ExecuTorchDataType)dataType + shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism { + return [self randomIntegerTensorWithLow:low + high:high + shape:tensor.shape + strides:tensor.strides + dataType:dataType + shapeDynamism:shapeDynamism]; +} + ++ (instancetype)randomIntegerTensorLikeTensor:(ExecuTorchTensor *)tensor + low:(NSInteger)low + high:(NSInteger)high + dataType:(ExecuTorchDataType)dataType { + return [self randomIntegerTensorWithLow:low + high:high + shape:tensor.shape + strides:tensor.strides + dataType:dataType + shapeDynamism:tensor.shapeDynamism]; +} + ++ (instancetype)randomIntegerTensorLikeTensor:(ExecuTorchTensor *)tensor + low:(NSInteger)low + high:(NSInteger)high { + return [self randomIntegerTensorWithLow:low + high:high + shape: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 27a6cac7b1e..6554ac4d2d6 100644 --- a/extension/apple/ExecuTorch/__tests__/TensorTest.swift +++ b/extension/apple/ExecuTorch/__tests__/TensorTest.swift @@ -655,4 +655,28 @@ class TensorTest: XCTestCase { XCTAssertEqual(tensor.shape, other.shape) XCTAssertEqual(tensor.count, other.count) } + + func testRandomInteger() { + let tensor = Tensor.randint(low: 10, high: 20, shape: [5], dataType: .int) + XCTAssertEqual(tensor.shape, [5]) + XCTAssertEqual(tensor.count, 5) + tensor.bytes { pointer, count, dataType in + XCTAssertEqual(dataType, .int) + let buffer = UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Int32.self), count: count) + for value in buffer { + XCTAssertTrue(value >= 10 && value < 20) + } + } + } + + func testRandomIntegerLike() { + let other = Tensor.ones(shape: [5], dataType: .int) + let tensor = Tensor.randint(like: other, low: 100, high: 200) + tensor.bytes { pointer, count, dataType in + let buffer = UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Int32.self), count: count) + for value in buffer { + XCTAssertTrue(value >= 100 && value < 200) + } + } + } }