feat(dipu): add schema: ones/zeros/zero_ (#742)

* add ones/zeros/zero_

* add tests

* add schema

* update

* opt

* update diopi main

* for python black format

Author: hellozmz

dipu/SupportedDiopiFunctions.txt

@@ -192,6 +192,7 @@ diopiNormalInp
 diopiNormalScalarTensor
 diopiNormalTensor
 diopiNormalTensorScalar
+diopiOnes
 diopiPolar
 diopiPow
 diopiPowInp
@@ -250,3 +251,5 @@ diopiUpsampleLinearBackward
 diopiUpsampleNearest
 diopiUpsampleNearestBackward
 diopiWhere
+diopiZeroInp
+diopiZeros

dipu/scripts/autogen_diopi_wrapper/diopi_functions.yaml

@@ -2279,6 +2279,29 @@
 - schema: "atan.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)"
   interface: diopiAtan(ctx, out, self)
 
+- schema: "ones.out(SymInt[] size, *, Tensor(a!) out) -> Tensor(a!)"
+  interface: diopiOnes(ctx, out, size)
+
+- schema: "ones(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor"
+  custom_code_at_the_beginning: |
+    c10::TensorOptions option;
+    auto shape = c10::asIntArrayRefUnchecked(size);
+    auto out = nodispatch::empty(shape, option.dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory));
+  interface: diopiOnes(ctx, out, size)
+
+- schema: "zeros(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor"
+  custom_code_at_the_beginning: |
+    c10::TensorOptions option;
+    auto shape = c10::asIntArrayRefUnchecked(size);
+    auto out = nodispatch::empty(shape, option.dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory));
+  interface: diopiZeroInp(ctx, out)
+
+- schema: "zeros.out(SymInt[] size, *, Tensor(a!) out) -> Tensor(a!)"
+  interface: diopiZeros(ctx, out, size)
+
+- schema: "zero_(Tensor(a!) self) -> Tensor(a!)"
+  interface: diopiZeroInp(ctx, self)
+
 - schema: "im2col(Tensor self, int[2] kernel_size, int[2] dilation, int[2] padding, int[2] stride) -> Tensor"
   size_attr: [kernel_size, stride, padding, dilation]
   custom_code_at_the_beginning: |

dipu/tests/python/unittests/test_ones.py

@@ -0,0 +1,17 @@
+# Copyright (c) 2024, DeepLink.
+import torch
+import torch_dipu
+from torch_dipu.testing._internal.common_utils import TestCase, run_tests
+
+
+class TestOnes(TestCase):
+    def test_ones(self):
+        device = torch.device("dipu")
+        size = [5, 6]
+        x = torch.ones(size=size)
+        y = torch.ones(size=size, device=device)
+        self.assertEqual(x, y.cpu(), exact_dtype=True)
+
+
+if __name__ == "__main__":
+    run_tests()

dipu/tests/python/unittests/test_zeros.py

@@ -0,0 +1,25 @@
+# Copyright (c) 2024, DeepLink.
+import torch
+import torch_dipu
+from torch_dipu.testing._internal.common_utils import TestCase, run_tests
+
+
+class TestZeros(TestCase):
+    def test_zeros(self):
+        device = torch.device("dipu")
+        size = [5, 6]
+        x = torch.zeros(size=size)
+        y = torch.zeros(size=size, device=device)
+        self.assertEqual(x, y.cpu(), exact_dtype=True)
+
+    def test_zero_(self):
+        size = [3, 5]
+        y = torch.randn(size=size).cuda()
+        x = y.cpu()
+        x.zero_()
+        y.zero_()
+        self.assertEqual(x, y.cpu(), exact_dtype=True)
+
+
+if __name__ == "__main__":
+    run_tests()