From 886ef520d769bfc1134cfd908b4d943a4653cf5c Mon Sep 17 00:00:00 2001
From: l00591931 <liangzhibo@huawei.com>
Date: Thu, 12 Nov 2020 15:07:53 +0800
Subject: [PATCH] Add Ones and Zeros operators

---
 mindspore/ops/operations/__init__.py  |  2 +-
 mindspore/ops/operations/array_ops.py | 87 +++++++++++++++++++++++++++
 tests/ut/python/ops/test_array_ops.py | 14 +++++
 3 files changed, 102 insertions(+), 1 deletion(-)

diff --git a/mindspore/ops/operations/__init__.py b/mindspore/ops/operations/__init__.py
index 898c6d3b61..057ab8485e 100644
--- a/mindspore/ops/operations/__init__.py
+++ b/mindspore/ops/operations/__init__.py
@@ -22,7 +22,7 @@ A collection of operators to build neural networks or to compute functions.
 from .image_ops import (CropAndResize)
 from .array_ops import (Argmax, Argmin, Cast, Concat, Pack, Unpack,
                         Diag, DiagPart, DType, ExpandDims, Eye,
-                        Fill, GatherNd, GatherV2, SparseGatherV2, InvertPermutation,
+                        Fill, Ones, Zeros, GatherNd, GatherV2, SparseGatherV2, InvertPermutation,
                         IsInstance, IsSubClass, ArgMaxWithValue, OnesLike, ZerosLike,
                         Rank, Reshape, ResizeNearestNeighbor, ArgMinWithValue, Meshgrid,
                         SameTypeShape, ScatterAdd, ScatterSub, ScatterMul, ScatterDiv, ScatterMax, ScatterMin,
diff --git a/mindspore/ops/operations/array_ops.py b/mindspore/ops/operations/array_ops.py
index 956d4a669e..b4ed92a1d0 100644
--- a/mindspore/ops/operations/array_ops.py
+++ b/mindspore/ops/operations/array_ops.py
@@ -998,6 +998,93 @@ class Fill(PrimitiveWithInfer):
         return out
 
 
+class Ones(PrimitiveWithInfer):
+    """
+    Creates a tensor filled with value ones.
+
+    Creates a tensor with shape described by the first argument and
+    fills it with value ones in type of the second argument.
+
+    Inputs:
+        - **shape** (tuple) - The specified shape of output tensor. Only constant value is allowed.
+        - **type** (mindspore.dtype) - The specified type of output tensor. Only constant value is allowed.
+
+    Outputs:
+        Tensor, has the same type and shape as input value.
+
+    Examples:
+        >>> ones = P.Ones()
+        >>> Ones((2, 2), mindspore.float32)
+        [[1.0, 1.0],
+         [1.0, 1.0]]
+    """
+
+    @prim_attr_register
+    def __init__(self):
+        """Initialize Fill"""
+
+    def __infer__(self, dims, dtype):
+        validator.check_value_type("shape", dims['value'], [tuple], self.name)
+        for i, item in enumerate(dims['value']):
+            validator.check_positive_int(item, f'dims[{i}]', self.name)
+        valid_types = [mstype.bool_, mstype.int8, mstype.int16, mstype.int32, mstype.int64,
+                       mstype.uint8, mstype.uint32, mstype.uint64,
+                       mstype.float16, mstype.float32, mstype.float64]
+        validator.check_types_same_and_valid({"value": dtype['value']}, valid_types, self.name)
+        x_nptype = mstype.dtype_to_nptype(dtype['value'])
+        ret = np.ones(dims['value'], x_nptype)
+        out = {
+            'value': Tensor(ret),
+            'shape': dims['value'],
+            'dtype': x_nptype,
+        }
+        return out
+
+
+class Zeros(PrimitiveWithInfer):
+    """
+    Creates a tensor filled with value zeros.
+
+    Creates a tensor with shape described by the first argument and
+    fills it with value zeros in type of the second argument.
+
+    Inputs:
+        - **shape** (tuple) - The specified shape of output tensor. Only constant value is allowed.
+        - **type** (mindspore.dtype) - The specified type of output tensor. Only constant value is allowed.
+
+    Outputs:
+        Tensor, has the same type and shape as input value.
+
+    Examples:
+        >>> zeros = P.Zeros()
+        >>> Zeros((2, 2), mindspore.float32)
+        [[0.0, 0.0],
+         [0.0, 0.0]]
+
+    """
+
+    @prim_attr_register
+    def __init__(self):
+        """Initialize Fill"""
+
+    def __infer__(self, dims, dtype):
+        validator.check_value_type("shape", dims['value'], [tuple], self.name)
+        for i, item in enumerate(dims['value']):
+            validator.check_positive_int(item, f'dims[{i}]', self.name)
+        valid_types = [mstype.bool_, mstype.int8, mstype.int16, mstype.int32, mstype.int64,
+                       mstype.uint8, mstype.uint32, mstype.uint64,
+                       mstype.float16, mstype.float32, mstype.float64]
+        validator.check_types_same_and_valid({"value": dtype['value']}, valid_types, self.name)
+        x_nptype = mstype.dtype_to_nptype(dtype['value'])
+        ret = np.zeros(dims['value'], x_nptype)
+        out = {
+            'value': Tensor(ret),
+            'shape': dims['value'],
+            'dtype': x_nptype,
+        }
+        return out
+
+
 class OnesLike(PrimitiveWithInfer):
     """
     Creates a new tensor. The values of all elements are 1.
diff --git a/tests/ut/python/ops/test_array_ops.py b/tests/ut/python/ops/test_array_ops.py
index 5dbdfe42e8..febf4c6b99 100644
--- a/tests/ut/python/ops/test_array_ops.py
+++ b/tests/ut/python/ops/test_array_ops.py
@@ -52,6 +52,20 @@ def test_cast():
     assert np.all(result.asnumpy() == expect)
 
 
+def test_ones():
+    ones = P.Ones()
+    output = ones((2, 3), mstype.int32)
+    assert output.asnumpy().shape == (2, 3)
+    assert np.sum(output.asnumpy()) == 6
+
+
+def test_zeros():
+    zeros = P.Zeros()
+    output = zeros((2, 3), mstype.int32)
+    assert output.asnumpy().shape == (2, 3)
+    assert np.sum(output.asnumpy()) == 0
+
+
 @non_graph_engine
 def test_reshape():
     input_tensor = Tensor(np.array([[-0.1, 0.3, 3.6], [0.4, 0.5, -3.2]]))