[PyOV] support Type as arg in convert op (#21603)

* [PyOV] support Type as arg in convert op

* add test

* fix linter
This commit is contained in:
Anastasia Kuporosova authored and GitHub committed 2023-12-12 18:21:47 +01:00
1 parent a2cc6a036f
commit 3e42ddbde5
2 files changed
+10 -5

No files matched your search

@@ -3,7 +3,7 @@
# SPDX-License-Identifier: Apache-2.0
"""Factory functions for all openvino ops."""
from typing import List, Optional, Union
from typing import List, Optional, Union, get_args
import numpy as np
from functools import partial
@@ -352,7 +352,7 @@ def constant(
@nameable_op
def convert(
data: NodeInput,
destination_type: Union[str, NumericType],
destination_type: Union[str, NumericType, Type],
name: Optional[str] = None,
) -> Node:
"""Return node which casts input node values to specified type.
@@ -362,12 +362,15 @@ def convert(
:param name: Optional name for the output node.
:return: New node performing the conversion operation.
"""
if not isinstance(destination_type, str):
destination_type = get_element_type_str(destination_type)
_destination_type = None # type: Union[str, Type]
if isinstance(destination_type, get_args(NumericType)):
_destination_type = get_element_type_str(destination_type).lower()
else:
_destination_type = destination_type
return _get_node_factory_opset1().create(
"Convert",
[as_node(data)],
{"destination_type": destination_type.lower()},
{"destination_type": _destination_type},
)
@@ -144,6 +144,8 @@ def test_convert_to_bool(destination_type, input_data):
pytest.param(np.float64, (-16383, 16383), np.int64, np.float64),
pytest.param("f32", (-8, 8), np.int32, np.float32),
pytest.param("f64", (-16383, 16383), np.int64, np.float64),
pytest.param(Type.f32, (-8, 8), np.int32, np.float32),
pytest.param(Type.f64, (-16383, 16383), np.int64, np.float64),
],
)
def test_convert_to_float(destination_type, rand_range, in_dtype, expected_type):