Helpers to create zeros tensor. (pytorch#10643)

Summary: pytorch#8366

Reviewed By: kirklandsign

Differential Revision: D74020937

Author: shoumikhin

extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h

@@ -902,6 +902,69 @@ __attribute__((deprecated("This API is experimental.")))
 
 @end
 
+#pragma mark - Zeros Category
+
+@interface ExecuTorchTensor (Zeros)
+
+/**
+ * Creates a tensor filled with zeros, with the specified shape, data type, and shape dynamism.
+ *
+ * @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 zeros.
+ */
++ (instancetype)zerosTensorWithShape:(NSArray<NSNumber *> *)shape
+                            dataType:(ExecuTorchDataType)dataType
+                       shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism
+    NS_SWIFT_NAME(zeros(shape:dataType:shapeDynamism:));
+
+/**
+ * Creates a tensor filled with zeros, 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.
+ * @return A new ExecuTorchTensor instance filled with zeros.
+ */
++ (instancetype)zerosTensorWithShape:(NSArray<NSNumber *> *)shape
+                            dataType:(ExecuTorchDataType)dataType
+    NS_SWIFT_NAME(zeros(shape:dataType:));
+
+/**
+ * Creates a tensor filled with zeros 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 zeros.
+ */
++ (instancetype)zerosTensorLikeTensor:(ExecuTorchTensor *)tensor
+                             dataType:(ExecuTorchDataType)dataType
+                        shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism
+    NS_SWIFT_NAME(zeros(like:dataType:shapeDynamism:));
+
+/**
+ * Creates a tensor filled with zeros 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 zeros.
+ */
++ (instancetype)zerosTensorLikeTensor:(ExecuTorchTensor *)tensor
+                             dataType:(ExecuTorchDataType)dataType
+    NS_SWIFT_NAME(zeros(like:dataType:));
+
+/**
+ * Creates a tensor filled with zeros similar to an existing tensor.
+ *
+ * @param tensor An existing ExecuTorchTensor instance.
+ * @return A new ExecuTorchTensor instance filled with zeros.
+ */
++ (instancetype)zerosTensorLikeTensor:(ExecuTorchTensor *)tensor
+    NS_SWIFT_NAME(zeros(like:));
+
+@end
+
 #pragma mark - Random Category
 
 @interface ExecuTorchTensor (Random)

extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm

@@ -776,6 +776,56 @@ + (instancetype)onesTensorLikeTensor:(ExecuTorchTensor *)tensor {
 
 @end
 
+@implementation ExecuTorchTensor (Zeros)
+
++ (instancetype)zerosTensorWithShape:(NSArray<NSNumber *> *)shape
+                            dataType:(ExecuTorchDataType)dataType
+                       shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism {
+  return [self fullTensorWithShape:shape
+                            scalar:@(0)
+                           strides:@[]
+                          dataType:dataType
+                     shapeDynamism:shapeDynamism];
+}
+
++ (instancetype)zerosTensorWithShape:(NSArray<NSNumber *> *)shape
+                            dataType:(ExecuTorchDataType)dataType {
+  return [self fullTensorWithShape:shape
+                            scalar:@(0)
+                           strides:@[]
+                          dataType:dataType
+                     shapeDynamism:ExecuTorchShapeDynamismDynamicBound];
+}
+
++ (instancetype)zerosTensorLikeTensor:(ExecuTorchTensor *)tensor
+                             dataType:(ExecuTorchDataType)dataType
+                        shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism {
+  return [self fullTensorWithShape:tensor.shape
+                            scalar:@(0)
+                           strides:tensor.strides
+                          dataType:dataType
+                     shapeDynamism:shapeDynamism];
+}
+
++ (instancetype)zerosTensorLikeTensor:(ExecuTorchTensor *)tensor
+                             dataType:(ExecuTorchDataType)dataType {
+  return [self fullTensorWithShape:tensor.shape
+                            scalar:@(0)
+                           strides:tensor.strides
+                          dataType:dataType
+                     shapeDynamism:tensor.shapeDynamism];
+}
+
++ (instancetype)zerosTensorLikeTensor:(ExecuTorchTensor *)tensor {
+  return [self fullTensorWithShape:tensor.shape
+                            scalar:@(0)
+                           strides:tensor.strides
+                          dataType:tensor.dataType
+                     shapeDynamism:tensor.shapeDynamism];
+}
+
+@end
+
 @implementation ExecuTorchTensor (Random)
 
 + (instancetype)randomTensorWithShape:(NSArray<NSNumber *> *)shape

extension/apple/ExecuTorch/__tests__/TensorTest.swift

@@ -618,6 +618,31 @@ class TensorTest: XCTestCase {
       }
     }
   }
+  
+  func testZeros() {
+    let tensor = Tensor.zeros(shape: [2, 3], dataType: .double)
+    XCTAssertEqual(tensor.shape, [2, 3])
+    XCTAssertEqual(tensor.count, 6)
+    tensor.bytes { pointer, count, dataType in
+      XCTAssertEqual(dataType, .double)
+      let buffer = UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Double.self), count: count)
+      for value in buffer {
+        XCTAssertEqual(value, 0)
+      }
+    }
+  }
+
+  func testZerosLike() {
+    let other = Tensor.full(shape: [3, 2], scalar: 9, dataType: .int)
+    let tensor = Tensor.zeros(like: other)
+    XCTAssertEqual(tensor.shape, other.shape)
+    tensor.bytes { pointer, count, dataType in
+      let buffer = UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Int32.self), count: count)
+      for value in buffer {
+        XCTAssertEqual(value, 0)
+      }
+    }
+  }
 
   func testRandom() {
     let tensor = Tensor.rand(shape: [3, 3], dataType: .float)