Helpers to create full tensors.

shoumikhin · web-flow · commit 9a85b06ff170 · 2025-05-01T17:40:26.000-07:00
Differential Revision: D74020935 Pull Request resolved: #10641
diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h
@@ -748,4 +748,95 @@ __attribute__((deprecated("This API is experimental.")))
 
 @end
 
+#pragma mark - Full Category
+
+@interface ExecuTorchTensor (Full)
+
+/**
+ * Creates a tensor filled with the specified scalar value, with full specification of shape, strides, data type, and shape dynamism.
+ *
+ * @param shape An NSArray of NSNumber objects representing the desired shape.
+ * @param scalar An NSNumber representing the value to fill the tensor.
+ * @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 the scalar value.
+ */
++ (instancetype)fullTensorWithShape:(NSArray<NSNumber *> *)shape
+                             scalar:(NSNumber *)scalar
+                            strides:(NSArray<NSNumber *> *)strides
+                           dataType:(ExecuTorchDataType)dataType
+                      shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism
+    NS_SWIFT_NAME(full(shape:scalar:strides:dataType:shapeDynamism:));
+
+/**
+ * Creates a tensor filled with the specified scalar value, with the given shape, data type, and shape dynamism.
+ *
+ * @param shape An NSArray of NSNumber objects representing the desired shape.
+ * @param scalar An NSNumber representing the value to fill the tensor.
+ * @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 the scalar value.
+ */
++ (instancetype)fullTensorWithShape:(NSArray<NSNumber *> *)shape
+                             scalar:(NSNumber *)scalar
+                           dataType:(ExecuTorchDataType)dataType
+                      shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism
+    NS_SWIFT_NAME(full(shape:scalar:dataType:shapeDynamism:));
+
+/**
+ * Creates a tensor filled with the specified scalar value, with the given shape and data type,
+ * using dynamic bound shape for strides and dimension order.
+ *
+ * @param shape An NSArray of NSNumber objects representing the desired shape.
+ * @param scalar An NSNumber representing the value to fill the tensor.
+ * @param dataType An ExecuTorchDataType value specifying the element type.
+ * @return A new ExecuTorchTensor instance filled with the scalar value.
+ */
++ (instancetype)fullTensorWithShape:(NSArray<NSNumber *> *)shape
+                             scalar:(NSNumber *)scalar
+                           dataType:(ExecuTorchDataType)dataType
+    NS_SWIFT_NAME(full(shape:scalar:dataType:));
+
+/**
+ * Creates a tensor filled with the specified scalar value, similar to an existing tensor, with the given data type and shape dynamism.
+ *
+ * @param tensr An existing ExecuTorchTensor instance whose shape and strides are used.
+ * @param scalar An NSNumber representing the value to fill the tensor.
+ * @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 the scalar value.
+ */
++ (instancetype)fullTensorLikeTensor:(ExecuTorchTensor *)tensr
+                              scalar:(NSNumber *)scalar
+                            dataType:(ExecuTorchDataType)dataType
+                       shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism
+    NS_SWIFT_NAME(full(like:scalar:dataType:shapeDynamism:));
+
+/**
+ * Creates a tensor filled with the specified scalar value, similar to an existing tensor, with the given data type.
+ *
+ * @param tensr An existing ExecuTorchTensor instance whose shape and strides are used.
+ * @param scalar An NSNumber representing the value to fill the tensor.
+ * @param dataType An ExecuTorchDataType value specifying the desired element type.
+ * @return A new ExecuTorchTensor instance filled with the scalar value.
+ */
++ (instancetype)fullTensorLikeTensor:(ExecuTorchTensor *)tensr
+                              scalar:(NSNumber *)scalar
+                            dataType:(ExecuTorchDataType)dataType
+    NS_SWIFT_NAME(full(like:scalar:dataType:));
+
+/**
+ * Creates a tensor filled with the specified scalar value, similar to an existing tensor.
+ *
+ * @param tensr An existing ExecuTorchTensor instance.
+ * @param scalar An NSNumber representing the value to fill the tensor.
+ * @return A new ExecuTorchTensor instance filled with the scalar value.
+ */
++ (instancetype)fullTensorLikeTensor:(ExecuTorchTensor *)tensr
+                              scalar:(NSNumber *)scalar
+    NS_SWIFT_NAME(full(like:scalar:));
+
+@end
+
 NS_ASSUME_NONNULL_END
diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm
@@ -649,3 +649,79 @@ + (instancetype)emptyTensorLikeTensor:(ExecuTorchTensor *)tensor {
 }
 
 @end
+
+@implementation ExecuTorchTensor (Full)
+
++ (instancetype)fullTensorWithShape:(NSArray<NSNumber *> *)shape
+                             scalar:(NSNumber *)scalar
+                            strides:(NSArray<NSNumber *> *)strides
+                           dataType:(ExecuTorchDataType)dataType
+                      shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism {
+  Scalar fillValue;
+  ET_SWITCH_REALHBBF16_AND_UINT_TYPES(
+    static_cast<ScalarType>(dataType), nil, "fullTensor", CTYPE, [&] {
+      fillValue = utils::extractValue<CTYPE>(scalar);
+    }
+  );
+  auto tensor = full_strided(
+    utils::toVector<SizesType>(shape),
+    utils::toVector<StridesType>(strides),
+    fillValue,
+    static_cast<ScalarType>(dataType),
+    static_cast<TensorShapeDynamism>(shapeDynamism)
+  );
+  return [[self alloc] initWithNativeInstance:&tensor];
+}
+
++ (instancetype)fullTensorWithShape:(NSArray<NSNumber *> *)shape
+                             scalar:(NSNumber *)scalar
+                           dataType:(ExecuTorchDataType)dataType
+                      shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism {
+  return [self fullTensorWithShape:shape
+                            scalar:scalar
+                           strides:@[]
+                          dataType:dataType
+                     shapeDynamism:shapeDynamism];
+}
+
++ (instancetype)fullTensorWithShape:(NSArray<NSNumber *> *)shape
+                             scalar:(NSNumber *)scalar
+                           dataType:(ExecuTorchDataType)dataType {
+  return [self fullTensorWithShape:shape
+                            scalar:scalar
+                           strides:@[]
+                          dataType:dataType
+                     shapeDynamism:ExecuTorchShapeDynamismDynamicBound];
+}
+
++ (instancetype)fullTensorLikeTensor:(ExecuTorchTensor *)tensor
+                              scalar:(NSNumber *)scalar
+                            dataType:(ExecuTorchDataType)dataType
+                       shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism {
+  return [self fullTensorWithShape:tensor.shape
+                            scalar:scalar
+                           strides:tensor.strides
+                          dataType:dataType
+                     shapeDynamism:shapeDynamism];
+}
+
++ (instancetype)fullTensorLikeTensor:(ExecuTorchTensor *)tensor
+                              scalar:(NSNumber *)scalar
+                            dataType:(ExecuTorchDataType)dataType {
+  return [self fullTensorWithShape:tensor.shape
+                            scalar:scalar
+                           strides:tensor.strides
+                          dataType:dataType
+                     shapeDynamism:tensor.shapeDynamism];
+}
+
++ (instancetype)fullTensorLikeTensor:(ExecuTorchTensor *)tensor
+                              scalar:(NSNumber *)scalar {
+  return [self fullTensorWithShape:tensor.shape
+                            scalar:scalar
+                           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
@@ -568,4 +568,29 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.dimensionOrder, other.dimensionOrder)
     XCTAssertEqual(tensor.dataType, other.dataType)
   }
+
+  func testFull() {
+    let tensor = Tensor.full(shape: [2, 2], scalar: 7, dataType: .int)
+    XCTAssertEqual(tensor.shape, [2, 2])
+    XCTAssertEqual(tensor.count, 4)
+    tensor.bytes { pointer, count, dataType in
+      XCTAssertEqual(dataType, .int)
+      let buffer = UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Int32.self), count: count)
+      for value in buffer {
+        XCTAssertEqual(value, 7)
+      }
+    }
+  }
+
+  func testFullLike() {
+    let other = Tensor.empty(shape: [2, 2], dataType: .int)
+    let tensor = Tensor.full(like: other, scalar: 42, dataType: .float)
+    XCTAssertEqual(tensor.shape, other.shape)
+    tensor.bytes { pointer, count, dataType in
+      let buffer = UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Float.self), count: count)
+      for value in buffer {
+        XCTAssertEqual(value, 42.0)
+      }
+    }
+  }
 }
