Skip to content

Commit 74dbf15

Browse files
shoumikhinshoumikhin
authored
Helpers to create random tensor.
Differential Revision: D74020936 Pull Request resolved: pytorch#10645
1 parent 5202b43 commit 74dbf15

File tree

3 files changed

+155
-0
lines changed

3 files changed

+155
-0
lines changed

‎extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -902,4 +902,82 @@ __attribute__((deprecated("This API is experimental.")))
902902

903903
@end
904904

905+
#pragma mark - Random Category
906+
907+
@interface ExecuTorchTensor (Random)
908+
909+
/**
910+
* Creates a tensor with random values, with full specification of shape, strides, data type, and shape dynamism.
911+
*
912+
* @param shape An NSArray of NSNumber objects representing the desired shape.
913+
* @param strides An NSArray of NSNumber objects representing the desired strides.
914+
* @param dataType An ExecuTorchDataType value specifying the element type.
915+
* @param shapeDynamism An ExecuTorchShapeDynamism value specifying whether the shape is static or dynamic.
916+
* @return A new ExecuTorchTensor instance filled with random values.
917+
*/
918+
+ (instancetype)randomTensorWithShape:(NSArray<NSNumber *> *)shape
919+
strides:(NSArray<NSNumber *> *)strides
920+
dataType:(ExecuTorchDataType)dataType
921+
shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism
922+
NS_SWIFT_NAME(rand(shape:strides:dataType:shapeDynamism:));
923+
924+
/**
925+
* Creates a tensor with random values, with the specified shape and data type.
926+
*
927+
* @param shape An NSArray of NSNumber objects representing the desired shape.
928+
* @param dataType An ExecuTorchDataType value specifying the element type.
929+
* @param shapeDynamism An ExecuTorchShapeDynamism value specifying whether the shape is static or dynamic.
930+
* @return A new ExecuTorchTensor instance filled with random values.
931+
*/
932+
+ (instancetype)randomTensorWithShape:(NSArray<NSNumber *> *)shape
933+
dataType:(ExecuTorchDataType)dataType
934+
shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism
935+
NS_SWIFT_NAME(rand(shape:dataType:shapeDynamism:));
936+
937+
/**
938+
* Creates a tensor with random values, with the specified shape (using dynamic bound shape) and data type.
939+
*
940+
* @param shape An NSArray of NSNumber objects representing the desired shape.
941+
* @param dataType An ExecuTorchDataType value specifying the element type.
942+
* @return A new ExecuTorchTensor instance filled with random values.
943+
*/
944+
+ (instancetype)randomTensorWithShape:(NSArray<NSNumber *> *)shape
945+
dataType:(ExecuTorchDataType)dataType
946+
NS_SWIFT_NAME(rand(shape:dataType:));
947+
948+
/**
949+
* Creates a tensor with random values similar to an existing tensor, with the specified data type and shape dynamism.
950+
*
951+
* @param tensor An existing ExecuTorchTensor instance whose shape and strides are used.
952+
* @param dataType An ExecuTorchDataType value specifying the desired element type.
953+
* @param shapeDynamism An ExecuTorchShapeDynamism value specifying whether the shape is static or dynamic.
954+
* @return A new ExecuTorchTensor instance filled with random values.
955+
*/
956+
+ (instancetype)randomTensorLikeTensor:(ExecuTorchTensor *)tensor
957+
dataType:(ExecuTorchDataType)dataType
958+
shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism
959+
NS_SWIFT_NAME(rand(like:dataType:shapeDynamism:));
960+
961+
/**
962+
* Creates a tensor with random values similar to an existing tensor, with the specified data type.
963+
*
964+
* @param tensor An existing ExecuTorchTensor instance whose shape and strides are used.
965+
* @param dataType An ExecuTorchDataType value specifying the desired element type.
966+
* @return A new ExecuTorchTensor instance filled with random values.
967+
*/
968+
+ (instancetype)randomTensorLikeTensor:(ExecuTorchTensor *)tensor
969+
dataType:(ExecuTorchDataType)dataType
970+
NS_SWIFT_NAME(rand(like:dataType:));
971+
972+
/**
973+
* Creates a tensor with random values similar to an existing tensor.
974+
*
975+
* @param tensor An existing ExecuTorchTensor instance.
976+
* @return A new ExecuTorchTensor instance filled with random values.
977+
*/
978+
+ (instancetype)randomTensorLikeTensor:(ExecuTorchTensor *)tensor
979+
NS_SWIFT_NAME(rand(like:));
980+
981+
@end
982+
905983
NS_ASSUME_NONNULL_END

‎extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -775,3 +775,61 @@ + (instancetype)onesTensorLikeTensor:(ExecuTorchTensor *)tensor {
775775
}
776776

777777
@end
778+
779+
@implementation ExecuTorchTensor (Random)
780+
781+
+ (instancetype)randomTensorWithShape:(NSArray<NSNumber *> *)shape
782+
strides:(NSArray<NSNumber *> *)strides
783+
dataType:(ExecuTorchDataType)dataType
784+
shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism {
785+
auto tensor = rand_strided(
786+
utils::toVector<SizesType>(shape),
787+
utils::toVector<StridesType>(strides),
788+
static_cast<ScalarType>(dataType),
789+
static_cast<TensorShapeDynamism>(shapeDynamism)
790+
);
791+
return [[self alloc] initWithNativeInstance:&tensor];
792+
}
793+
794+
+ (instancetype)randomTensorWithShape:(NSArray<NSNumber *> *)shape
795+
dataType:(ExecuTorchDataType)dataType
796+
shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism {
797+
return [self randomTensorWithShape:shape
798+
strides:@[]
799+
dataType:dataType
800+
shapeDynamism:shapeDynamism];
801+
}
802+
803+
+ (instancetype)randomTensorWithShape:(NSArray<NSNumber *> *)shape
804+
dataType:(ExecuTorchDataType)dataType {
805+
return [self randomTensorWithShape:shape
806+
strides:@[]
807+
dataType:dataType
808+
shapeDynamism:ExecuTorchShapeDynamismDynamicBound];
809+
}
810+
811+
+ (instancetype)randomTensorLikeTensor:(ExecuTorchTensor *)tensor
812+
dataType:(ExecuTorchDataType)dataType
813+
shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism {
814+
return [self randomTensorWithShape:tensor.shape
815+
strides:tensor.strides
816+
dataType:dataType
817+
shapeDynamism:shapeDynamism];
818+
}
819+
820+
+ (instancetype)randomTensorLikeTensor:(ExecuTorchTensor *)tensor
821+
dataType:(ExecuTorchDataType)dataType {
822+
return [self randomTensorWithShape:tensor.shape
823+
strides:tensor.strides
824+
dataType:dataType
825+
shapeDynamism:tensor.shapeDynamism];
826+
}
827+
828+
+ (instancetype)randomTensorLikeTensor:(ExecuTorchTensor *)tensor {
829+
return [self randomTensorWithShape:tensor.shape
830+
strides:tensor.strides
831+
dataType:tensor.dataType
832+
shapeDynamism:tensor.shapeDynamism];
833+
}
834+
835+
@end

‎extension/apple/ExecuTorch/__tests__/TensorTest.swift‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -618,4 +618,23 @@ class TensorTest: XCTestCase {
618618
}
619619
}
620620
}
621+
622+
func testRandom() {
623+
let tensor = Tensor.rand(shape: [3, 3], dataType: .float)
624+
XCTAssertEqual(tensor.shape, [3, 3])
625+
XCTAssertEqual(tensor.count, 9)
626+
tensor.bytes { pointer, count, dataType in
627+
XCTAssertEqual(dataType, .float)
628+
let buffer = UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Float.self), count: count)
629+
let uniqueValues = Set(buffer.map { $0 })
630+
XCTAssertTrue(uniqueValues.count > 1)
631+
}
632+
}
633+
634+
func testRandomLike() {
635+
let other = Tensor.full(shape: [3, 3], scalar: 9, dataType: .int)
636+
let tensor = Tensor.rand(like: other)
637+
XCTAssertEqual(tensor.shape, other.shape)
638+
XCTAssertEqual(tensor.count, other.count)
639+
}
621640
}

0 commit comments

Comments
 (0)