Dev type as module

docs/source/experimental.rst

@@ -220,3 +220,5 @@ Experimental features
 .. autofunction:: oneflow.experimental.nn.ZeroPad2d
 .. autofunction:: oneflow.experimental.tensor_buffer_to_tensor
 .. autofunction:: oneflow.experimental.tensor_to_tensor_buffer
+.. autofunction:: oneflow.experimental.Tensor.type_as
+.. autofunction:: oneflow.experimental.Tensor.long

oneflow/python/nn/modules/tensor_ops.py

@@ -0,0 +1,94 @@
+"""
+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 oneflow as flow
+from oneflow.python.oneflow_export import experimental_api
+from oneflow.python.nn.module import Module
+from oneflow.python.framework.tensor import register_tensor_op
+
+
+class TypeAs(Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, input, target):
+        return input.to(dtype=target.dtype)
+
+
+@register_tensor_op("type_as")
+@experimental_api
+def type_as_op(input, target):
+    r"""Returns this tensor cast to the type of the given tensor.
+        This is a no-op if the tensor is already of the correct type.
+
+    Args:
+        input  (Tensor): the input tensor.
+        target (Tensor): the tensor which has the desired type.
+
+    For example:
+
+    .. code-block:: python
+
+        >>> import oneflow.experimental as flow
+        >>> import numpy as np
+        >>> flow.enable_eager_execution()
+
+        >>> input = flow.Tensor(np.random.randn(1, 2, 3), dtype=flow.float32)
+        >>> target = flow.Tensor(np.random.randn(4, 5, 6), dtype = flow.int32)
+        >>> input = input.type_as(target)
+        >>> input.dtype
+        oneflow.int32
+
+    """
+    return TypeAs()(input, target)
+
+
+class Long(Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, input):
+        return input.to(dtype=flow.int64)
+
+
+@register_tensor_op("long")
+@experimental_api
+def long_op(input):
+    r"""`Tensor.long()` is equivalent to `Tensor.to(flow.int64)`. See to().
+
+    Args:
+        input  (Tensor): the input tensor.
+
+    For example:
+
+    .. code-block:: python
+
+        >>> import oneflow.experimental as flow
+        >>> import numpy as np
+        >>> flow.enable_eager_execution()
+
+        >>> input = flow.Tensor(np.random.randn(1, 2, 3), dtype=flow.float32)
+        >>> input = input.long()
+        >>> input.dtype
+        oneflow.int64
+
+    """
+    return Long()(input)
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod(raise_on_error=True)

oneflow/python/test/modules/test_tensor_ops.py

@@ -0,0 +1,64 @@
+"""
+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
+from collections import OrderedDict
+
+import numpy as np
+
+import oneflow.experimental as flow
+from test_util import GenArgList
+
+
+def _test_type_as(test_case, shape, device, src_dtype, tgt_dtype):
+    np_input = np.random.rand(*shape)
+    input = flow.tensor(np_input, dtype=src_dtype, device=device)
+    target = flow.tensor(np_input, dtype=tgt_dtype, device=device)
+    input = input.type_as(target)
+    test_case.assertEqual(input.dtype, target.dtype)
+
+
+def _test_long(test_case, shape, device, dtype):
+    np_input = np.random.rand(*shape)
+    input = flow.tensor(np_input, dtype=dtype, device=device)
+    input = input.long()
+    test_case.assertEqual(input.dtype, flow.int64)
+
+
+@unittest.skipIf(
+    not flow.unittest.env.eager_execution_enabled(),
+    ".numpy() doesn't work in lazy mode",
+)
+class TestTensorOps(flow.unittest.TestCase):
+    def test_type_as(test_case):
+        arg_dict = OrderedDict()
+        arg_dict["shape"] = [(1, 2), (3, 4, 5), (2, 3, 4, 5)]
+        arg_dict["device"] = ["cpu", "cuda"]
+        arg_dict["src_dtype"] = [flow.int64, flow.int32, flow.float32, flow.float64]
+        arg_dict["tgt_dtype"] = [flow.int64, flow.int32, flow.float32, flow.float64]
+        for arg in GenArgList(arg_dict):
+            _test_type_as(test_case, *arg)
+
+    def test_long(test_case):
+        arg_dict = OrderedDict()
+        arg_dict["shape"] = [(1, 2), (3, 4, 5), (2, 3, 4, 5)]
+        arg_dict["device"] = ["cpu", "cuda"]
+        arg_dict["dtype"] = [flow.int64, flow.int32, flow.float32, flow.float64]
+        for arg in GenArgList(arg_dict):
+            _test_long(test_case, *arg)
+
+
+if __name__ == "__main__":
+    unittest.main()