Add polyval op (#3541)

* add python interface for polyval and test, currently only work for float

* add test case for double type

* fix polyval op

* Add License

* format polyval

* update polyval op test case

* Add polyval bw test

* Fix 1n2d test case

* rm tensorflow in test case
Co-authored-by: Niamyf <yangf@zhejianglab.com>
Co-authored-by: Noneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>

oneflow/python/ops/math_ops.py

@@ -16,7 +16,7 @@ limitations under the License.
 from __future__ import absolute_import
 
 import os
-from typing import Union, Optional, Sequence
+from typing import Union, Optional, Sequence, List, Tuple
 
 import oneflow as flow
 import oneflow.core.operator.op_conf_pb2 as op_conf_util
@@ -1909,6 +1909,55 @@ def tril(
     )
 
 
+@oneflow_export("math.polyval")
+def polyval(
+    coeffs: Union[List, Tuple], x: remote_blob_util.BlobDef, name: Optional[str] = None
+) -> remote_blob_util.BlobDef:
+    r"""Computes the elementwise value of a polynomial.
+
+    Args:
+        coeffs (Union[List, Tuple]): The coefficients of the polynomial.
+        x (remote_blob_util.BlobDef): A Blob.
+        name (Optional[str], optional): The name for the operation. Defaults to None.
+
+    Returns:
+        remote_blob_util.BlobDef: A Blob, has the same data type of x.
+    
+    For example:
+
+    .. code-block:: python
+
+        import oneflow as flow
+        import numpy as np
+        import oneflow.typing as tp
+        
+        @flow.global_function()
+        def polyval_Job(
+            x: tp.Numpy.Placeholder((3,), dtype=flow.float32)
+        ) -> tp.Numpy:
+            coeffs = [1.0, 3.0, -2.0]
+            return flow.math.polyval(coeffs, x)
+
+        x = np.array([1.0, 2.0, 3.0]).astype(np.float32)
+        out = polyval_Job(x)
+        
+        # output [ 2. 8. 16.]
+
+    """
+    if name is None:
+        name = id_util.UniqueStr("Polyval_")
+    if not isinstance(coeffs, (list, tuple)):
+        raise ValueError(
+            "Argument coeffs must be list type " "found {}".format(type(coeffs))
+        )
+    if len(coeffs) < 1:
+        return flow.zeros_like(x, name=name)
+    p = flow.zeros_like(x, name=name)
+    for c in coeffs:
+        p = flow.math.add(c, flow.math.multiply(p, x))
+    return p
+
+
 @oneflow_export("range")
 def range(
     start, limit=None, delta=1, dtype=None, name="range"

oneflow/python/test/ops/test_polyval.py

+"""
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+import unittest
+import numpy as np
+import oneflow as flow
+import oneflow.typing as tp
+from collections import OrderedDict
+
+from test_util import GenArgList, type_name_to_flow_type, type_name_to_np_type
+
+
+def compare_with_tensorflow(device_type, device_num, in_shape, data_type, coeffs):
+    assert device_type in ["cpu", "gpu"]
+    assert data_type in ["float32", "double"]
+    flow_data_type = type_name_to_flow_type[data_type]
+    flow.clear_default_session()
+    if device_type == "cpu":
+        flow.config.cpu_device_num(device_num)
+    else:
+        flow.config.gpu_device_num(device_num)
+    func_config = flow.FunctionConfig()
+    func_config.default_data_type(flow_data_type)
+    func_config.default_placement_scope(
+        flow.scope.placement(device_type, "0:0-{}".format(device_num - 1))
+    )
+    func_config.default_logical_view(flow.scope.consistent_view())
+
+    x = (np.random.random(in_shape) * 100).astype(type_name_to_np_type[data_type])
+
+    def np_polyval_grad(coeffs, x):
+        coeffs_len = len(coeffs)
+        coeffs_diff = [(coeffs_len - i - 1) * coeffs[i] for i in range(coeffs_len - 1)]
+        np_x_diff = np.polyval(coeffs_diff, x)
+        return np_x_diff
+
+    def assert_prediction_grad(blob: tp.Numpy):
+        np_x_diff = np_polyval_grad(coeffs, x)
+        assert np.allclose(blob, np_x_diff, rtol=1e-5, atol=1e-5)
+
+    @flow.global_function(type="train", function_config=func_config)
+    def PolyValJob(x: tp.Numpy.Placeholder(shape=in_shape)):
+        with flow.scope.placement(device_type, "0:0"):
+            x += flow.get_variable(
+                name="x",
+                shape=in_shape,
+                dtype=flow_data_type,
+                initializer=flow.zeros_initializer(),
+                trainable=True,
+            )
+        flow.watch_diff(x, assert_prediction_grad)
+        out = flow.math.polyval(coeffs, x)
+        with flow.scope.placement(device_type, "0:0"):
+            flow.optimizer.SGD(
+                flow.optimizer.PiecewiseConstantScheduler([], [1e-4]), momentum=0
+            ).minimize(out)
+
+        return out
+
+    check_point = flow.train.CheckPoint()
+    check_point.init()
+    of_out = PolyValJob(x).get().numpy()
+    np_out = np.polyval(coeffs, x)
+    assert np.allclose(of_out, np_out, rtol=1e-5, atol=1e-5)
+
+
+def gen_arg_list(type):
+    arg_dict = OrderedDict()
+    if type == "1n2d":
+        arg_dict["device_type"] = ["gpu"]
+        arg_dict["device_num"] = [2]
+    else:
+        arg_dict["device_type"] = ["cpu", "gpu"]
+        arg_dict["device_num"] = [1]
+    arg_dict["in_shape"] = [(2, 3)]
+    arg_dict["data_type"] = ["float32"]
+    arg_dict["coeffs"] = [[1.0, 2.0], [1.0, 2.0, 3.0]]
+    return GenArgList(arg_dict)
+
+
+@flow.unittest.skip_unless_1n1d()
+class TestPolyval1n1d(flow.unittest.TestCase):
+    def test_polyval(test_case):
+        for arg in gen_arg_list("1n1d"):
+            compare_with_tensorflow(*arg)
+
+
+@flow.unittest.skip_unless_1n2d()
+class TestPolyval1n2d(flow.unittest.TestCase):
+    def test_polyval(test_case):
+        for arg in gen_arg_list("1n2d"):
+            compare_with_tensorflow(*arg)
+
+
+if __name__ == "__main__":
+    unittest.main()