add pow module (#4723)

* add pow module

* update

* update

* move pow module into new file

* refine doc

* format

* update pow and its test

* format

Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
Co-authored-by: oneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>

oneflow/python/nn/modules/math_ops.py

@@ -789,3 +789,33 @@ def std_op(tensor, dim, unbiased=True, keepdim=False):
 
     """
     return Std(dim, unbiased, keepdim)(tensor)
+
+
+class Pow(Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self._op = flow.builtin_op("scalar_pow").Input("in").Output("out").Build()
+
+    def forward(self, x, exponent: Union[int, float]):
+        return self._op(x, exponent=float(exponent))[0]
+
+
+@oneflow_export("pow")
+@register_tensor_op("pow")
+def pow_op(tensor, exponent):
+    r"""Takes the power of each element in input with exponent and returns a tensor with the result.
+    exponent can be either a single float number or a single int number.
+    
+    For example:
+
+    .. code-block:: python
+
+        import oneflow as flow
+        import numpy as np
+        
+        x = flow.Tensor(np.array([1, 2, 3, 4, 5, 6]))
+        out = flow.pow(x, 2).numpy()
+        print(out) # [1, 4, 9, 16, 25, 36]
+        
+    """
+    return Pow()(tensor, exponent)

oneflow/python/test/modules/test_math_ops.py

@@ -13,8 +13,8 @@ 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 oneflow as flow
 import unittest
+import oneflow as flow
 import numpy as np
 
 
@@ -148,5 +148,23 @@ class TestSquare(flow.unittest.TestCase):
         )
 
 
+@unittest.skipIf(
+    not flow.unittest.env.eager_execution_enabled(),
+    ".numpy() doesn't work in lazy mode",
+)
+class TestPow(flow.unittest.TestCase):
+    def test_pow(test_case):
+        input = flow.Tensor(np.array([1, 2, 3, 4, 5, 6]), dtype=flow.float32)
+        of_out = flow.pow(input, 2.1)
+        np_out = np.power(input.numpy(), 2.1)
+        test_case.assertTrue(np.allclose(of_out.numpy(), np_out, 1e-5, 1e-5))
+
+    def test_pow_tensor_function(test_case):
+        input = flow.Tensor(np.array([1, 2, 3, 4, 5, 6]), dtype=flow.float32)
+        of_out = input.pow(2.1)
+        np_out = np.power(input.numpy(), 2.1)
+        test_case.assertTrue(np.allclose(of_out.numpy(), np_out, 1e-5, 1e-5))
+
+
 if __name__ == "__main__":
     unittest.main()