-
Notifications
You must be signed in to change notification settings - Fork 2.5k
/
Copy pathtest_addmm.py
100 lines (80 loc) · 4.1 KB
/
test_addmm.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
# Copyright (C) 2018-2025 Intel Corporation
# SPDX-License-Identifier: Apache-2.0
import pytest
from pytorch_layer_test_class import PytorchLayerTest
class TestAddMM(PytorchLayerTest):
def _prepare_input(self, input_shape=(2, 2), matrix1_shape=(2, 2), matrix2_shape=(2, 2)):
import numpy as np
return (
np.random.randn(*input_shape).astype(np.float32),
np.random.randn(*matrix1_shape).astype(np.float32),
np.random.randn(*matrix2_shape).astype(np.float32)
)
def create_model(self, alpha, beta):
import torch
class aten_addmm(torch.nn.Module):
def __init__(self, alpha, beta):
super(aten_addmm, self).__init__()
self.alpha = alpha
self.beta = beta
def forward(self, m0, m1, m2):
return torch.addmm(m0, m1, m2, alpha=self.alpha, beta=self.beta)
ref_net = None
return aten_addmm(alpha, beta), ref_net, 'aten::addmm'
@pytest.mark.parametrize("kwargs_to_prepare_input", [
{"input_shape": (3, 3), 'matrix1_shape': (3, 3), 'matrix2_shape': (3, 3)},
{"input_shape": (2, 2), 'matrix1_shape': (2, 3), 'matrix2_shape': (3, 2)},
{"input_shape": (10, 1), 'matrix1_shape': (10, 5), 'matrix2_shape': (5, 1)},
{"input_shape": (1, 2), 'matrix1_shape': (1, 10), 'matrix2_shape': (10, 2)},
{"input_shape": (1, 1), 'matrix1_shape': (1, 10), 'matrix2_shape': (10, 1)},
])
@pytest.mark.parametrize("alpha,beta",
[(1., 1.), (0., 1.), (1., 0.), (1., 2.), (2., 1.), (-5., -6.), (3., 4.), (0.5, 0.75),
(1, 1)])
@pytest.mark.nightly
@pytest.mark.precommit
@pytest.mark.precommit_torch_export
def test_addmm(self, kwargs_to_prepare_input, alpha, beta, ie_device, precision, ir_version):
self._test(*self.create_model(alpha, beta), ie_device, precision, ir_version,
kwargs_to_prepare_input=kwargs_to_prepare_input)
class TestBAddBMM(PytorchLayerTest):
def _prepare_input(self, input_shape=(2, 2), matrix1_shape=(2, 2), matrix2_shape=(2, 2)):
import numpy as np
return (
np.random.randn(*input_shape).astype(np.float32),
np.random.randn(*matrix1_shape).astype(np.float32),
np.random.randn(*matrix2_shape).astype(np.float32)
)
def create_model(self, alpha, beta):
import torch
class aten_addmm(torch.nn.Module):
def __init__(self, alpha, beta):
super(aten_addmm, self).__init__()
self.alpha = alpha
self.beta = beta
def forward(self, m0, m1, m2):
return torch.baddbmm(m0, m1, m2, alpha=self.alpha, beta=self.beta)
ref_net = None
return aten_addmm(alpha, beta), ref_net, 'aten::baddbmm'
@pytest.mark.parametrize("kwargs_to_prepare_input", [
{"input_shape": (2, 3, 3), 'matrix1_shape': (2, 3, 3), 'matrix2_shape': (2, 3, 3)},
{"input_shape": (2, 2, 2), 'matrix1_shape': (2, 2, 3), 'matrix2_shape': (2, 3, 2)},
{"input_shape": (1, 10, 1), 'matrix1_shape': (1, 10, 5), 'matrix2_shape': (1, 5, 1)},
{"input_shape": (5, 1, 2), 'matrix1_shape': (5, 1, 10), 'matrix2_shape': (5, 10, 2)},
{"input_shape": (1, 1, 1), 'matrix1_shape': (1, 1, 10), 'matrix2_shape': (1, 10, 1)},
])
@pytest.mark.parametrize("alpha,beta",
[ # beta==0 in some cases produce nan in pytorch
(1., 1.),
(0., 1.),
(-5., -6.),
(3., 4.),
(0.5, 0.75),
(1, 1)
])
@pytest.mark.nightly
@pytest.mark.precommit
@pytest.mark.precommit_torch_export
def test_baddbmm(self, kwargs_to_prepare_input, alpha, beta, ie_device, precision, ir_version):
self._test(*self.create_model(alpha, beta), ie_device, precision, ir_version,
kwargs_to_prepare_input=kwargs_to_prepare_input)