Helpers to create random integer tensor.

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<NSNumber *> *)shape
+                                   strides:(NSArray<NSNumber *> *)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<NSNumber *> *)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<NSNumber *> *)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

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<NSNumber *> *)shape
+                                   strides:(NSArray<NSNumber *> *)strides
+                                  dataType:(ExecuTorchDataType)dataType
+                             shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism {
+  auto tensor = randint_strided(
+    low,
+    high,
+    utils::toVector<SizesType>(shape),
+    utils::toVector<StridesType>(strides),
+    static_cast<ScalarType>(dataType),
+    static_cast<TensorShapeDynamism>(shapeDynamism)
+  );
+  return [[self alloc] initWithNativeInstance:&tensor];
+}
+
++ (instancetype)randomIntegerTensorWithLow:(NSInteger)low
+                                      high:(NSInteger)high
+                                     shape:(NSArray<NSNumber *> *)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<NSNumber *> *)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

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)
+      }
+    }
+  }
 }