[PT FE]: support aten::linalng_vector_norm (#16109)
* [PT FE]: support aten::linalng_vector_norm * more norm ops * update tests
This commit is contained in:
Ekaterina Aidova authored and 
GitHub committed
1 parent
18df64c135
commit
1ef94ec069
3 files changed
+475
-10
No files matched your search
@@ -5,12 +5,23 @@
|
||||
#include "openvino/frontend/pytorch/node_context.hpp"
|
||||
#include "openvino/op/abs.hpp"
|
||||
#include "openvino/op/constant.hpp"
|
||||
#include "openvino/op/convert.hpp"
|
||||
#include "openvino/op/convert_like.hpp"
|
||||
#include "openvino/op/gather.hpp"
|
||||
#include "openvino/op/multiply.hpp"
|
||||
#include "openvino/op/not_equal.hpp"
|
||||
#include "openvino/op/power.hpp"
|
||||
#include "openvino/op/range.hpp"
|
||||
#include "openvino/op/reduce_l1.hpp"
|
||||
#include "openvino/op/reduce_l2.hpp"
|
||||
#include "openvino/op/reduce_max.hpp"
|
||||
#include "openvino/op/reduce_min.hpp"
|
||||
#include "openvino/op/reduce_sum.hpp"
|
||||
#include "openvino/op/reshape.hpp"
|
||||
#include "openvino/op/shape_of.hpp"
|
||||
#include "openvino/op/sqrt.hpp"
|
||||
#include "openvino/op/squeeze.hpp"
|
||||
#include "pt_framework_node.hpp"
|
||||
#include "utils.hpp"
|
||||
|
||||
namespace ov {
|
||||
@@ -20,13 +31,12 @@ namespace op {
|
||||
|
||||
using namespace ov::op;
|
||||
|
||||
OutputVector translate_norm(const NodeContext& context) {
|
||||
num_inputs_check(context, 4, 4);
|
||||
auto input_tensor = context.get_input(0);
|
||||
auto p = context.const_input<float>(1);
|
||||
auto dim = context.get_input(2);
|
||||
auto keep_dim = context.const_input<bool>(3);
|
||||
|
||||
namespace {
|
||||
Output<Node> norm_vector(const NodeContext& context,
|
||||
Output<Node> input_tensor,
|
||||
Output<Node> dim,
|
||||
float p,
|
||||
bool keep_dim) {
|
||||
Output<Node> res;
|
||||
if (p == 1) {
|
||||
res = context.mark_node(std::make_shared<v4::ReduceL1>(input_tensor, dim, keep_dim));
|
||||
@@ -38,17 +48,253 @@ OutputVector translate_norm(const NodeContext& context) {
|
||||
} else if (p == -std::numeric_limits<float>::infinity()) {
|
||||
auto abs = context.mark_node(std::make_shared<v0::Abs>(input_tensor));
|
||||
res = context.mark_node(std::make_shared<v1::ReduceMin>(abs, dim, keep_dim));
|
||||
} else if (p == 0) {
|
||||
auto input_rank = input_tensor.get_partial_shape().rank();
|
||||
FRONT_END_OP_CONVERSION_CHECK(input_rank.is_dynamic() || input_rank.get_length() == 1,
|
||||
"ord=0 supported only for vector norm");
|
||||
auto zero = context.mark_node(v0::Constant::create(element::f32, Shape{}, {0}));
|
||||
zero = context.mark_node(std::make_shared<v1::ConvertLike>(zero, input_tensor));
|
||||
auto cond = context.mark_node(std::make_shared<v1::NotEqual>(input_tensor, zero));
|
||||
cond = context.mark_node(std::make_shared<v1::ConvertLike>(cond, input_tensor));
|
||||
res = context.mark_node(std::make_shared<v1::ReduceSum>(cond, dim, keep_dim));
|
||||
} else {
|
||||
auto const_p = context.mark_node(v0::Constant::create(element::f64, Shape{1}, {p}));
|
||||
auto const_p_inv = context.mark_node(v0::Constant::create(element::f64, Shape{1}, {1.0 / p}));
|
||||
auto const_p = context.mark_node(v0::Constant::create(element::f32, Shape{1}, {p}));
|
||||
const_p = context.mark_node(std::make_shared<v1::ConvertLike>(const_p, input_tensor));
|
||||
auto const_p_inv = context.mark_node(v0::Constant::create(element::f32, Shape{1}, {1.0 / p}));
|
||||
const_p_inv = context.mark_node(std::make_shared<v1::ConvertLike>(const_p_inv, input_tensor));
|
||||
auto abs = context.mark_node(std::make_shared<v0::Abs>(input_tensor));
|
||||
auto pow = context.mark_node(std::make_shared<v1::Power>(abs, const_p));
|
||||
auto sum = context.mark_node(std::make_shared<v1::ReduceSum>(pow, dim, keep_dim));
|
||||
res = context.mark_node(std::make_shared<v1::Power>(sum, const_p_inv));
|
||||
}
|
||||
return res;
|
||||
};
|
||||
|
||||
Output<Node> norm_matrix(const NodeContext& context,
|
||||
Output<Node> input_tensor,
|
||||
Output<Node> dim,
|
||||
float p,
|
||||
bool keep_dim) {
|
||||
Output<Node> res;
|
||||
auto one = context.mark_node(v0::Constant::create(element::i32, Shape{}, {1}));
|
||||
auto zero = context.mark_node(v0::Constant::create(element::i32, Shape{}, {0}));
|
||||
auto first_dim = context.mark_node(std::make_shared<v8::Gather>(dim, zero, zero));
|
||||
auto second_dim = context.mark_node(std::make_shared<v8::Gather>(dim, one, zero));
|
||||
if (p == 1) {
|
||||
auto abs = context.mark_node(std::make_shared<v0::Abs>(input_tensor));
|
||||
auto sum = context.mark_node(std::make_shared<v1::ReduceSum>(abs, first_dim, true));
|
||||
res = context.mark_node(std::make_shared<v1::ReduceMax>(sum, second_dim, keep_dim));
|
||||
} else if (p == std::numeric_limits<float>::infinity()) {
|
||||
auto abs = context.mark_node(std::make_shared<v0::Abs>(input_tensor));
|
||||
auto sum = context.mark_node(std::make_shared<v1::ReduceSum>(abs, second_dim, true));
|
||||
res = context.mark_node(std::make_shared<v1::ReduceMax>(sum, first_dim, keep_dim));
|
||||
} else if (p == -std::numeric_limits<float>::infinity()) {
|
||||
auto abs = context.mark_node(std::make_shared<v0::Abs>(input_tensor));
|
||||
auto sum = context.mark_node(std::make_shared<v1::ReduceSum>(abs, second_dim, true));
|
||||
res = context.mark_node(std::make_shared<v1::ReduceMin>(sum, first_dim, keep_dim));
|
||||
} else if (p == -1) {
|
||||
auto abs = context.mark_node(std::make_shared<v0::Abs>(input_tensor));
|
||||
auto sum = context.mark_node(std::make_shared<v1::ReduceSum>(abs, first_dim, true));
|
||||
res = context.mark_node(std::make_shared<v1::ReduceMin>(sum, second_dim, keep_dim));
|
||||
} else {
|
||||
FRONT_END_OP_CONVERSION_CHECK(false, "Unsupported ord ", p, " for matrix norm");
|
||||
}
|
||||
|
||||
return res;
|
||||
};
|
||||
|
||||
Output<Node> frobenius_norm(const NodeContext& context, Output<Node> x, Output<Node> dim, bool keep_dim) {
|
||||
auto sqr = context.mark_node(std::make_shared<v1::Multiply>(x, x));
|
||||
auto sumsqr = context.mark_node(std::make_shared<v1::ReduceSum>(sqr, dim, keep_dim));
|
||||
return context.mark_node(std::make_shared<v0::Sqrt>(sumsqr));
|
||||
}
|
||||
}; // namespace
|
||||
|
||||
OutputVector translate_norm(const NodeContext& context) {
|
||||
num_inputs_check(context, 4, 6);
|
||||
auto input_tensor = context.get_input(0);
|
||||
auto p_node_type = context.get_input_type(1);
|
||||
Output<Node> dim;
|
||||
if (context.input_is_none(2)) {
|
||||
auto zero = context.mark_node(v0::Constant::create(element::i32, Shape{}, {0}));
|
||||
auto one = context.mark_node(v0::Constant::create(element::i32, Shape{}, {1}));
|
||||
auto input_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input_tensor, element::i32));
|
||||
auto rank = context.mark_node(std::make_shared<v3::ShapeOf>(input_shape, element::i32));
|
||||
rank = context.mark_node(std::make_shared<v0::Squeeze>(rank, zero));
|
||||
dim = context.mark_node(std::make_shared<v0::Range>(zero, rank, one));
|
||||
} else {
|
||||
dim = context.get_input(2);
|
||||
}
|
||||
auto keep_dim = context.const_input<bool>(3);
|
||||
if (!context.input_is_none(4)) {
|
||||
input_tensor = apply_dtype(context, 4, input_tensor);
|
||||
}
|
||||
Output<Node> res;
|
||||
if (p_node_type.is<type::Str>()) {
|
||||
auto p_str = context.const_input<std::string>(1);
|
||||
if (p_str == "fro") {
|
||||
res = frobenius_norm(context, input_tensor, dim, keep_dim);
|
||||
} else {
|
||||
FRONT_END_OP_CONVERSION_CHECK(false, "Umsupported ord ", p_str);
|
||||
}
|
||||
} else {
|
||||
auto p = context.const_input<float>(1);
|
||||
res = norm_vector(context, input_tensor, dim, p, keep_dim);
|
||||
}
|
||||
// output tensor
|
||||
if (!context.input_is_none(5)) {
|
||||
context.mutate_input(5, res);
|
||||
}
|
||||
return {res};
|
||||
};
|
||||
|
||||
OutputVector translate_linalg_vector_norm(const NodeContext& context) {
|
||||
// aten::linalg_vector_norm(Tensor self, Scalar ord=2, int[1]? dim=None, bool keepdim=False, *, ScalarType?
|
||||
// dtype=None) -> Tensor
|
||||
// aten::linalg_vector_norm.out(Tensor self, Scalar ord=2, int[1]? dim=None, bool
|
||||
// keepdim=False, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!):
|
||||
num_inputs_check(context, 5, 6);
|
||||
auto x = context.get_input(0);
|
||||
// ord defines the vector norm that is computed.
|
||||
auto ord = context.const_input<float>(1);
|
||||
bool keep_dim = context.const_input<bool>(3);
|
||||
Output<Node> dim;
|
||||
Output<Node> result;
|
||||
// If dim= None, x will be flattened before the norm is computed.
|
||||
if (context.input_is_none(2)) {
|
||||
keep_dim = false;
|
||||
auto minus_one = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {-1}));
|
||||
x = context.mark_node(std::make_shared<v1::Reshape>(x, minus_one, false));
|
||||
dim = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {0}));
|
||||
} else {
|
||||
dim = context.get_input(2);
|
||||
}
|
||||
// dtype may be used to perform the computation in a more precise dtype. It is semantically equivalent to calling
|
||||
// linalg.vector_norm(x.to(dtype))
|
||||
if (!context.input_is_none(4)) {
|
||||
x = apply_dtype(context, 4, x);
|
||||
}
|
||||
result = norm_vector(context, x, dim, ord, keep_dim);
|
||||
// output tensor
|
||||
if (!context.input_is_none(5)) {
|
||||
context.mutate_input(5, result);
|
||||
}
|
||||
return {result};
|
||||
};
|
||||
|
||||
OutputVector translate_linalg_matrix_norm(const NodeContext& context) {
|
||||
// aten::linalg_matrix_norm.out(Tensor self, Scalar ord, int[] dim=[-2, -1], bool keepdim=False, *, ScalarType?
|
||||
// dtype=None, Tensor(a!) out) -> Tensor(a!) aten::linalg_matrix_norm(Tensor self, Scalar ord, int[] dim=[-2, -1],
|
||||
// bool keepdim=False, *, ScalarType? dtype=None) aten::linalg_matrix_norm.str_ord(Tensor self, str ord="fro", int[]
|
||||
// dim=[-2, -1], bool keepdim=False, *, ScalarType? dtype=None)
|
||||
num_inputs_check(context, 5, 6);
|
||||
auto x = context.get_input(0);
|
||||
// ord defines the vector norm that is computed can be string or number
|
||||
auto ord_type = context.get_input_type(1);
|
||||
auto dim = context.get_input(2);
|
||||
bool keep_dim = context.const_input<bool>(3);
|
||||
Output<Node> result;
|
||||
|
||||
// dtype may be used to perform the computation in a more precise dtype. It is semantically equivalent to calling
|
||||
// linalg.mtrix_norm(x.to(dtype))
|
||||
if (!context.input_is_none(4)) {
|
||||
x = apply_dtype(context, 4, x);
|
||||
}
|
||||
if (ord_type.is<type::Str>()) {
|
||||
auto p_str = context.const_input<std::string>(1);
|
||||
if (p_str == "fro") {
|
||||
result = frobenius_norm(context, x, dim, keep_dim);
|
||||
} else {
|
||||
FRONT_END_OP_CONVERSION_CHECK(false, "Umsupported ord ", p_str);
|
||||
}
|
||||
} else {
|
||||
auto p = context.const_input<float>(1);
|
||||
result = norm_matrix(context, x, dim, p, keep_dim);
|
||||
}
|
||||
// output tensor
|
||||
if (!context.input_is_none(5)) {
|
||||
context.mutate_input(5, result);
|
||||
}
|
||||
return {result};
|
||||
};
|
||||
|
||||
OutputVector translate_linalg_norm(const NodeContext& context) {
|
||||
// aten::linalg_norm(Tensor self, Scalar? ord=None, int[1]? dim=None, bool keepdim=False, *, ScalarType? dtype=None)
|
||||
// aten::linalg_norm.ord_str(Tensor self, str ord, int[1]? dim=None, bool keepdim=False, *, ScalarType? dtype=None)
|
||||
// aten::linalg_norm.ord_str_out(Tensor self, str ord, int[1]? dim=None, bool keepdim=False, *, ScalarType?
|
||||
// dtype=None, Tensor(a!) out) -> Tensor(a!)
|
||||
num_inputs_check(context, 5, 6);
|
||||
auto x = context.get_input(0);
|
||||
bool keep_dim = context.const_input<bool>(3);
|
||||
Output<Node> result;
|
||||
Output<Node> dim;
|
||||
// dtype may be used to perform the computation in a more precise dtype. It is semantically equivalent to calling
|
||||
// linalg.norm(x.to(dtype))
|
||||
if (!context.input_is_none(4)) {
|
||||
x = apply_dtype(context, 4, x);
|
||||
}
|
||||
// If dim= None and ord= None, A will be flattened to 1D and the 2-norm of the resulting vector will be computed.
|
||||
if (context.input_is_none(2) && context.input_is_none(1)) {
|
||||
auto minus_one = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {-1}));
|
||||
x = context.mark_node(std::make_shared<v1::Reshape>(x, minus_one, false));
|
||||
dim = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {0}));
|
||||
result = norm_vector(context, x, dim, 2, false);
|
||||
} else {
|
||||
// If dim=None apply for all dimesions
|
||||
if (context.input_is_none(2)) {
|
||||
auto zero = context.mark_node(v0::Constant::create(element::i32, Shape{}, {0}));
|
||||
auto one = context.mark_node(v0::Constant::create(element::i32, Shape{}, {1}));
|
||||
auto input_shape = context.mark_node(std::make_shared<v3::ShapeOf>(x, element::i32));
|
||||
auto rank = context.mark_node(std::make_shared<v3::ShapeOf>(input_shape, element::i32));
|
||||
rank = context.mark_node(std::make_shared<v0::Squeeze>(rank, zero));
|
||||
dim = context.mark_node(std::make_shared<v0::Range>(zero, rank, one));
|
||||
} else {
|
||||
dim = context.get_input(2);
|
||||
}
|
||||
// default norm for matrix is frobenius norm, for vector - L2, for other ranks are not detrmined
|
||||
if (context.input_is_none(1)) {
|
||||
auto input_rank = x.get_partial_shape().rank();
|
||||
if (input_rank.is_static() && input_rank.get_length() == 2) {
|
||||
result = frobenius_norm(context, x, dim, keep_dim);
|
||||
} else if (input_rank.is_static() && input_rank.get_length() == 1) {
|
||||
result = norm_vector(context, x, dim, 2, keep_dim);
|
||||
} else {
|
||||
FRONT_END_OP_CONVERSION_CHECK(false,
|
||||
"linalg norm for tensor rank > 2 without ord specification unsupported");
|
||||
}
|
||||
} else {
|
||||
// ord defines the norm that is computed can be string or number
|
||||
auto ord_type = context.get_input_type(1);
|
||||
if (ord_type.is<type::Str>()) {
|
||||
auto p_str = context.const_input<std::string>(1);
|
||||
if (p_str == "fro") {
|
||||
result = frobenius_norm(context, x, dim, keep_dim);
|
||||
} else {
|
||||
FRONT_END_OP_CONVERSION_CHECK(false, "Umsupported ord ", p_str);
|
||||
}
|
||||
} else {
|
||||
auto p = context.const_input<float>(1);
|
||||
if (!context.input_is_none(2)) {
|
||||
auto const_dim = context.const_input<std::vector<int64_t>>(2);
|
||||
if (const_dim.size() == 2) {
|
||||
result = norm_matrix(context, x, dim, p, keep_dim);
|
||||
} else {
|
||||
result = norm_vector(context, x, dim, p, keep_dim);
|
||||
}
|
||||
} else {
|
||||
result = norm_vector(context, x, dim, p, keep_dim);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// output tensor
|
||||
if (!context.input_is_none(5)) {
|
||||
context.mutate_input(5, result);
|
||||
}
|
||||
return {result};
|
||||
};
|
||||
|
||||
} // namespace op
|
||||
} // namespace pytorch
|
||||
} // namespace frontend
|
||||
|
||||
@@ -62,6 +62,9 @@ OP_CONVERTER(translate_instance_norm);
|
||||
OP_CONVERTER(translate_int);
|
||||
OP_CONVERTER(translate_layer_norm);
|
||||
OP_CONVERTER(translate_len);
|
||||
OP_CONVERTER(translate_linalg_norm);
|
||||
OP_CONVERTER(translate_linalg_matrix_norm);
|
||||
OP_CONVERTER(translate_linalg_vector_norm);
|
||||
OP_CONVERTER(translate_linear);
|
||||
OP_CONVERTER(translate_list_construct);
|
||||
OP_CONVERTER(translate_log);
|
||||
@@ -235,6 +238,9 @@ const std::map<std::string, CreatorFunction> get_supported_ops() {
|
||||
{"aten::leaky_relu", op::translate_1to1_match_2_inputs<opset10::PRelu>},
|
||||
{"aten::leaky_relu_", op::inplace_op<op::translate_1to1_match_2_inputs<opset10::PRelu>>},
|
||||
{"aten::len", op::translate_len},
|
||||
{"aten::linalg_norm", op::translate_linalg_norm},
|
||||
{"aten::linalg_matrix_norm", op::translate_linalg_matrix_norm},
|
||||
{"aten::linalg_vector_norm", op::translate_linalg_vector_norm},
|
||||
{"aten::linear", op::translate_linear},
|
||||
{"aten::log", op::translate_log},
|
||||
{"aten::log_", op::inplace_op<op::translate_log>},
|
||||
|
||||
@@ -14,7 +14,7 @@ from pytorch_layer_test_class import PytorchLayerTest
|
||||
class TestNorm(PytorchLayerTest):
|
||||
|
||||
def _prepare_input(self):
|
||||
return (np.random.randn(2, 3, 4, 5),)
|
||||
return (np.random.randn(1, 2, 3).astype(np.float32),)
|
||||
|
||||
def create_model(self, p, dim, keepdim):
|
||||
class aten_norm(torch.nn.Module):
|
||||
@@ -37,3 +37,216 @@ class TestNorm(PytorchLayerTest):
|
||||
def test_norm(self, ie_device, precision, ir_version, p, dim, keepdim):
|
||||
self._test(*self.create_model(p, dim, keepdim),
|
||||
ie_device, precision, ir_version)
|
||||
|
||||
|
||||
class TestLinalgVectorNorm(PytorchLayerTest):
|
||||
|
||||
def _prepare_input(self, out=False, out_dtype=None):
|
||||
if not out:
|
||||
return (np.random.randn(1, 2, 3).astype(np.float32),)
|
||||
x = np.random.randn(1, 2, 3).astype(np.float32)
|
||||
y = np.random.randn(1, 2, 3).astype(
|
||||
out_dtype if out_dtype is not None else np.float32)
|
||||
return (x, y)
|
||||
|
||||
def create_model(self, p, dim, keepdim, dtype_str, out, out_as_dtype):
|
||||
dtypes = {
|
||||
"float32": torch.float32,
|
||||
"float64": torch.float64
|
||||
}
|
||||
dtype = dtypes.get(dtype_str)
|
||||
|
||||
class aten_linalg_vector_norm(torch.nn.Module):
|
||||
|
||||
def __init__(self, p, dim, keepdim, dtype, out, out_as_dtype) -> None:
|
||||
super().__init__()
|
||||
self.ord = p
|
||||
self.dim = dim
|
||||
self.keepdim = keepdim
|
||||
self.dtype = dtype
|
||||
if self.dtype is not None:
|
||||
self.forward = self.forward_dtype
|
||||
if out:
|
||||
self.forward = self.forward_out
|
||||
if out_as_dtype:
|
||||
self.forward = self.forward_prim_dtype
|
||||
|
||||
def forward(self, x):
|
||||
return torch.linalg.vector_norm(
|
||||
x, ord=self.ord, dim=self.dim, keepdim=self.keepdim
|
||||
)
|
||||
|
||||
def forward_dtype(self, x):
|
||||
return torch.linalg.vector_norm(
|
||||
x, ord=self.ord, dim=self.dim, keepdim=self.keepdim, dtype=self.dtype
|
||||
)
|
||||
|
||||
def forward_prim_dtype(self, x, y):
|
||||
return torch.linalg.vector_norm(
|
||||
x, ord=self.ord, dim=self.dim, keepdim=self.keepdim, dtype=y.dtype
|
||||
)
|
||||
|
||||
def forward_out(self, x, y):
|
||||
return y, torch.linalg.vector_norm(
|
||||
x, ord=self.ord, dim=self.dim, keepdim=self.keepdim, out=y
|
||||
)
|
||||
|
||||
ref_net = None
|
||||
|
||||
return aten_linalg_vector_norm(p, dim, keepdim, dtype, out, out_as_dtype), ref_net, "aten::linalg_vector_norm"
|
||||
|
||||
@pytest.mark.nightly
|
||||
@pytest.mark.precommit
|
||||
@pytest.mark.parametrize('p', [-2, -1, 0, 1, 2, 2.5, float('inf'), float('-inf')])
|
||||
@pytest.mark.parametrize('dim', [0, [0, 1], None])
|
||||
@pytest.mark.parametrize('keepdim', [True, False])
|
||||
@pytest.mark.parametrize("dtype", ["float32", "float64", None])
|
||||
@pytest.mark.parametrize("out", [True, False])
|
||||
@pytest.mark.parametrize("prim_dtype", [True, False])
|
||||
def test_linalg_vector_norm(self, p, dim, keepdim, dtype, out, prim_dtype, ie_device, precision, ir_version):
|
||||
self._test(*self.create_model(p, dim, keepdim, dtype, out, prim_dtype),
|
||||
ie_device, precision, ir_version,
|
||||
kwargs_to_prepare_input={"out": out or prim_dtype, "out_dtype": dtype if prim_dtype else None})
|
||||
|
||||
|
||||
class TestLinalgMatrixNorm(PytorchLayerTest):
|
||||
|
||||
def _prepare_input(self, out=False, out_dtype=None):
|
||||
if not out:
|
||||
return (np.random.randn(3, 3).astype(np.float32),)
|
||||
x = np.random.randn(1, 3, 3).astype(np.float32)
|
||||
y = np.random.randn(1, 3, 3).astype(
|
||||
out_dtype if out_dtype is not None else np.float32)
|
||||
return (x, y)
|
||||
|
||||
def create_model(self, p, dim, keepdim, dtype_str, out, out_as_dtype):
|
||||
dtypes = {
|
||||
"float32": torch.float32,
|
||||
"float64": torch.float64
|
||||
}
|
||||
dtype = dtypes.get(dtype_str)
|
||||
|
||||
class aten_linalg_matrix_norm(torch.nn.Module):
|
||||
|
||||
def __init__(self, p, dim, keepdim, dtype, out, out_as_dtype) -> None:
|
||||
super().__init__()
|
||||
self.ord = p
|
||||
self.dim = dim
|
||||
self.keepdim = keepdim
|
||||
self.dtype = dtype
|
||||
if self.dtype is not None:
|
||||
self.forward = self.forward_dtype
|
||||
if out:
|
||||
self.forward = self.forward_out
|
||||
if out_as_dtype:
|
||||
self.forward = self.forward_prim_dtype
|
||||
|
||||
def forward(self, x):
|
||||
return torch.linalg.matrix_norm(
|
||||
x, ord=self.ord, dim=self.dim, keepdim=self.keepdim
|
||||
)
|
||||
|
||||
def forward_dtype(self, x):
|
||||
return torch.linalg.matrix_norm(
|
||||
x, ord=self.ord, dim=self.dim, keepdim=self.keepdim, dtype=self.dtype
|
||||
)
|
||||
|
||||
def forward_prim_dtype(self, x, y):
|
||||
return torch.linalg.matrix_norm(
|
||||
x, ord=self.ord, dim=self.dim, keepdim=self.keepdim, dtype=y.dtype
|
||||
)
|
||||
|
||||
def forward_out(self, x, y):
|
||||
return y, torch.linalg.matrix_norm(
|
||||
x, ord=self.ord, dim=self.dim, keepdim=self.keepdim, out=y
|
||||
)
|
||||
|
||||
ref_net = None
|
||||
|
||||
return aten_linalg_matrix_norm(p, dim, keepdim, dtype, out, out_as_dtype), ref_net, "aten::linalg_matrix_norm"
|
||||
|
||||
@pytest.mark.nightly
|
||||
@pytest.mark.precommit
|
||||
@pytest.mark.parametrize('p', [-1, 1, float('inf'), float('-inf'), "fro"])
|
||||
@pytest.mark.parametrize('dim', [[0, 1], [-1, -2]])
|
||||
@pytest.mark.parametrize('keepdim', [True, False])
|
||||
@pytest.mark.parametrize("dtype", ["float32", "float64", None])
|
||||
@pytest.mark.parametrize("out", [True, False])
|
||||
@pytest.mark.parametrize("prim_dtype", [True, False])
|
||||
def test_linalg_matrix_norm(self, p, dim, keepdim, dtype, out, prim_dtype, ie_device, precision, ir_version):
|
||||
self._test(*self.create_model(p, dim, keepdim, dtype, out, prim_dtype),
|
||||
ie_device, precision, ir_version,
|
||||
kwargs_to_prepare_input={"out": out or prim_dtype, "out_dtype": dtype if prim_dtype else None})
|
||||
|
||||
|
||||
class TestLinalgNorm(PytorchLayerTest):
|
||||
|
||||
def _prepare_input(self, out=False, out_dtype=None):
|
||||
if not out:
|
||||
return (np.random.randn(3, 3).astype(np.float32),)
|
||||
x = np.random.randn(3, 3).astype(np.float32)
|
||||
y = np.random.randn(3, 3).astype(
|
||||
out_dtype if out_dtype is not None else np.float32)
|
||||
return (x, y)
|
||||
|
||||
def create_model(self, p, dim, keepdim, dtype_str, out, out_as_dtype):
|
||||
dtypes = {
|
||||
"float32": torch.float32,
|
||||
"float64": torch.float64
|
||||
}
|
||||
dtype = dtypes.get(dtype_str)
|
||||
|
||||
class aten_linalg_matrix_norm(torch.nn.Module):
|
||||
|
||||
def __init__(self, p, dim, keepdim, dtype, out, out_as_dtype) -> None:
|
||||
super().__init__()
|
||||
self.ord = p
|
||||
self.dim = dim
|
||||
self.keepdim = keepdim
|
||||
self.dtype = dtype
|
||||
if self.dtype is not None:
|
||||
self.forward = self.forward_dtype
|
||||
if out:
|
||||
self.forward = self.forward_out
|
||||
if out_as_dtype:
|
||||
self.forward = self.forward_prim_dtype
|
||||
|
||||
def forward(self, x):
|
||||
return torch.linalg.norm(
|
||||
x, ord=self.ord, dim=self.dim, keepdim=self.keepdim
|
||||
)
|
||||
|
||||
def forward_dtype(self, x):
|
||||
return torch.linalg.norm(
|
||||
x, ord=self.ord, dim=self.dim, keepdim=self.keepdim, dtype=self.dtype
|
||||
)
|
||||
|
||||
def forward_prim_dtype(self, x, y):
|
||||
return torch.linalg.norm(
|
||||
x, ord=self.ord, dim=self.dim, keepdim=self.keepdim, dtype=y.dtype
|
||||
)
|
||||
|
||||
def forward_out(self, x, y):
|
||||
return y, torch.linalg.norm(
|
||||
x, ord=self.ord, dim=self.dim, keepdim=self.keepdim, out=y
|
||||
)
|
||||
|
||||
ref_net = None
|
||||
|
||||
return aten_linalg_matrix_norm(p, dim, keepdim, dtype, out, out_as_dtype), ref_net, "aten::linalg_norm"
|
||||
|
||||
@pytest.mark.nightly
|
||||
@pytest.mark.precommit
|
||||
@pytest.mark.parametrize('p,dim', [
|
||||
(-1, [0, 1]), (1, [-1, -2]), (float('inf'), [1, 0]),
|
||||
(float('-inf'), [-2, -1]), (0, 1), (1, -1),
|
||||
(None, None), (2.5, 0), (-1, 1), (2, 0),
|
||||
(float('inf'), 1), (float('-inf'), 1), ("fro", (0, 1))])
|
||||
@pytest.mark.parametrize('keepdim', [True, False])
|
||||
@pytest.mark.parametrize("dtype", ["float32", "float64", None])
|
||||
@pytest.mark.parametrize("out", [True, False])
|
||||
@pytest.mark.parametrize("prim_dtype", [True, False])
|
||||
def test_linalg_norm(self, p, dim, keepdim, dtype, out, prim_dtype, ie_device, precision, ir_version):
|
||||
self._test(*self.create_model(p, dim, keepdim, dtype, out, prim_dtype),
|
||||
ie_device, precision, ir_version,
|
||||
kwargs_to_prepare_input={"out": out or prim_dtype, "out_dtype": dtype if prim_dtype else None})
|
||||
Reference in new issue
Block a user