aten::arange

src/frontends/pytorch/src/op/arange.cpp

@@ -0,0 +1,96 @@
+// Copyright (C) 2018-2022 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "openvino/frontend/pytorch/node_context.hpp"
+#include "openvino/opsets/opset8.hpp"
+#include "utils.hpp"
+
+namespace ov {
+namespace frontend {
+namespace pytorch {
+namespace op {
+
+OutputVector translate_arange(NodeContext& context) {
+    auto zero = context.mark_node(opset8::Constant::create(element::i32, Shape{}, {0}));
+    auto one = context.mark_node(opset8::Constant::create(element::i32, Shape{}, {1}));
+    auto dtype = element::f32;
+    bool dtype_applied = false;
+    int num_inputs = context.get_input_size();
+    // aten::arange(Scalar end, tensor out)
+    if (num_inputs == 2) {
+        auto end = context.get_input(0);
+        auto range = context.mark_node(std::make_shared<opset8::Range>(zero, end, one, dtype));
+        return {context.mark_node(std::make_shared<opset8::ConvertLike>(range, end))};
+    }
+    // # aten::arange(Scalar start, Scalar end, Scalar step, Tensor out)
+    if (num_inputs == 4) {
+        auto start = context.get_input(0);
+        auto end = context.get_input(1);
+        auto step = context.get_input(2);
+        auto range = context.mark_node(std::make_shared<opset8::Range>(start, end, step, dtype));
+        return {context.mark_node(std::make_shared<opset8::ConvertLike>(range, end))};
+    }
+    // aten::arange(Scalar end, ScalarType dtype, Layout, Device, bool pin_memory)
+    if (num_inputs == 5) {
+        auto end = context.get_input(0);
+        if (!context.input_is_none(1)) {
+            auto pt_type = context.const_input<int64_t>(1);
+            FRONT_END_OP_CONVERSION_CHECK(TORCH_TO_OV_TYPE.count(pt_type), "Unknown type in aten::arange: ", pt_type);
+            dtype = TORCH_TO_OV_TYPE.at(pt_type);
+            end = context.mark_node(std::make_shared<opset8::Convert>(end, dtype));
+            zero = context.mark_node(std::make_shared<opset8::Convert>(zero, dtype));
+            one = context.mark_node(std::make_shared<opset8::Convert>(one, dtype));
+            dtype_applied = true;
+        }
+        auto range = context.mark_node(std::make_shared<opset8::Range>(zero, end, one, dtype));
+        if (!dtype_applied) {
+            return {context.mark_node(std::make_shared<opset8::ConvertLike>(range, end))};
+        }
+        return {range};
+    }
+    // aten::arange(Scalar start, Scalar end, ScalarType dtype, Layout, Device, bool pin_memory)
+    if (num_inputs == 6) {
+        auto start = context.get_input(0);
+        auto end = context.get_input(1);
+        if (!context.input_is_none(2)) {
+            auto pt_type = context.const_input<int64_t>(2);
+            FRONT_END_OP_CONVERSION_CHECK(TORCH_TO_OV_TYPE.count(pt_type), "Unknown type in aten::arange: ", pt_type);
+            dtype = TORCH_TO_OV_TYPE.at(pt_type);
+            dtype_applied = true;
+            end = context.mark_node(std::make_shared<opset8::Convert>(end, dtype));
+            start = context.mark_node(std::make_shared<opset8::Convert>(start, dtype));
+            one = context.mark_node(std::make_shared<opset8::Convert>(one, dtype));
+        }
+        auto range = context.mark_node(std::make_shared<opset8::Range>(start, end, one, dtype));
+        if (!dtype_applied) {
+            return {context.mark_node(std::make_shared<opset8::ConvertLike>(range, end))};
+        }
+        return {range};
+    }
+    // aten::arange(Scalar start, Scalar end, Scalar step, ScalarType dtype, Layout, Device, bool pin_memory)
+    if (num_inputs == 7) {
+        auto start = context.get_input(0);
+        auto end = context.get_input(1);
+        auto step = context.get_input(2);
+        if (!context.input_is_none(3)) {
+            auto pt_type = context.const_input<int64_t>(3);
+            FRONT_END_OP_CONVERSION_CHECK(TORCH_TO_OV_TYPE.count(pt_type), "Unknown type in aten::arange: ", pt_type);
+            dtype = TORCH_TO_OV_TYPE.at(pt_type);
+            end = context.mark_node(std::make_shared<opset8::Convert>(end, dtype));
+            start = context.mark_node(std::make_shared<opset8::Convert>(start, dtype));
+            step = context.mark_node(std::make_shared<opset8::Convert>(step, dtype));
+            dtype_applied = true;
+        }
+        auto range = context.mark_node(std::make_shared<opset8::Range>(start, end, step, dtype));
+        if (!dtype_applied) {
+            return {context.mark_node(std::make_shared<opset8::ConvertLike>(range, end))};
+        }
+        return {range};
+    }
+};
+
+}  // namespace op
+}  // namespace pytorch
+}  // namespace frontend
+}  // namespace ov

src/frontends/pytorch/src/op_table.cpp

@@ -19,6 +19,7 @@ OP_CONVERTER(translate_adaptive_max_pool2d);
 OP_CONVERTER(translate_add);
 OP_CONVERTER(translate_addcmul);
 OP_CONVERTER(translate_addmm);
+OP_CONVERTER(translate_arange);
 OP_CONVERTER(translate_as_tensor);
 OP_CONVERTER(translate_avg_pool2d);
 OP_CONVERTER(translate_batch_norm);
@@ -98,6 +99,7 @@ const std::map<std::string, CreatorFunction> get_supported_ops() {
         {"aten::add_", op::inplace_op<op::translate_add>},
         {"aten::addcmul", op::translate_addcmul},
         {"aten::addmm", op::translate_addmm},
+        {"aten::arange", op::translate_arange},
         {"aten::as_tensor", op::translate_as_tensor},
         {"aten::avg_pool2d", op::translate_avg_pool2d},
         {"aten::batch_norm", op::translate_batch_norm},

tests/layer_tests/pytorch_tests/pytorch_layer_test_class.py

@@ -102,7 +102,7 @@ def _test(self, model, ref_net, kind, ie_device, precision, ir_version, infer_ti
                 else:
                     assert type(fw_tensor) == type(ov_tensor)
                 continue
-            assert torch.tensor(np.array(ov_tensor)).dtype == fw_tensor.dtype
+            assert torch.tensor(np.array(ov_tensor)).dtype == fw_tensor.dtype, f"dtype validation failed: {torch.tensor(np.array(ov_tensor)).dtype} != {fw_tensor.dtype}"
 
         if 'custom_eps' in kwargs and kwargs['custom_eps'] is not None:
             custom_eps = kwargs['custom_eps']

tests/layer_tests/pytorch_tests/test_arange.py

@@ -0,0 +1,80 @@
+# Copyright (C) 2018-2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+from typing import Union
+from pytorch_layer_test_class import PytorchLayerTest
+
+
+class TestExp(PytorchLayerTest):
+    def _prepare_input(self, end, start=None, step=None, dtype="int64"):
+        import numpy as np
+        if start is None and step is None:
+            return (np.array(end).astype(dtype), )
+        if step is None:
+            return (np.array(start).astype(dtype), np.array(end).astype(dtype))
+        return (np.array(start).astype(dtype), np.array(end).astype(dtype), np.array(step).astype(dtype))
+
+    def create_model(self, dtype, num_inputs):
+        import torch
+
+        dtype_map = {
+            "float32": torch.float32,
+            "float64": torch.float64,
+            "int64": torch.int64,
+            "int32": torch.int32,
+            "uint8": torch.uint8,
+            "int8": torch.int8
+        }
+        class aten_arange_end(torch.nn.Module):
+            def __init__(self, dtype) -> None:
+                super(aten_arange_end, self).__init__()
+                self.dtype = dtype
+
+            def forward(self, x:int):
+                return torch.arange(x, dtype=self.dtype)
+
+        class aten_arange_start_end(torch.nn.Module):
+            def __init__(self, dtype) -> None:
+                super(aten_arange_start_end, self).__init__()
+                self.dtype = dtype
+
+            def forward(self, x:float, y:float):
+                return torch.arange(start=x, end=y, dtype=self.dtype)
+
+        class aten_arange_start_end_step(torch.nn.Module):
+            def __init__(self, dtype) -> None:
+                super(aten_arange_start_end_step, self).__init__()
+                self.dtype = dtype
+
+            def forward(self, x:float, y:float, z:float):
+                return torch.arange(start=x, end=y, step=z, dtype=self.dtype)
+        model_classes = {
+            1: aten_arange_end,
+            2: aten_arange_start_end,
+            3: aten_arange_start_end_step
+        }
+        dtype = dtype_map.get(dtype)
+        model_class = model_classes[num_inputs]
+
+        ref_net = None
+
+        return model_class(dtype), ref_net, "aten::arange"
+
+    @pytest.mark.nightly
+    @pytest.mark.parametrize("dtype", [None, "float32", "float64", "int32", "int64", "int8", "uin8"])
+    @pytest.mark.parametrize("end", [1, 2, 3])
+    def test_arange_end_only(self, dtype, end, ie_device, precision, ir_version):
+        self._test(*self.create_model(dtype, 1), ie_device, precision, ir_version, kwargs_to_prepare_input={"end": end})
+
+    @pytest.mark.nightly
+    @pytest.mark.parametrize("dtype", [None, "float32", "float64", "int32", "int64", "int8"])
+    @pytest.mark.parametrize("start,end", [(0, 1), (-1, 1), (1, 5), (0.5, 2.5)])
+    def test_arange_start_end(self, dtype, end, start, ie_device, precision, ir_version):
+        self._test(*self.create_model(dtype, 2), ie_device, precision, ir_version, kwargs_to_prepare_input={"end": end, "start": start, "dtype": "float32"})
+
+    @pytest.mark.nightly
+    @pytest.mark.parametrize("dtype", [None, "float32", "float64", "int32", "int64", "int8"])
+    @pytest.mark.parametrize("start,end,step", [(0, 1, 1), (-2, 1, 1.25), (1, -5, -1), (1, 10, 2), (-1, -5, -2)])
+    def test_arange_start_end_step(self, dtype, end, start, step, ie_device, precision, ir_version):
+        self._test(*self.create_model(dtype, 3), ie_device, precision, ir_version, kwargs_to_prepare_input={"end": end, "start": start, "step": step, "dtype": "float32"})