|
2 | 2 |
|
3 | 3 | import torch |
4 | 4 | import torch.nn as nn |
| 5 | +from .harness import DispatchTestCase |
5 | 6 | from parameterized import parameterized |
6 | 7 | from torch.testing._internal.common_utils import run_tests |
7 | 8 | from torch_tensorrt import Input |
8 | 9 |
|
9 | | -from .harness import DispatchTestCase |
10 | | - |
11 | 10 | NEED_TEST_BOTH_CONSTANTS_CASE = True |
12 | 11 |
|
13 | 12 | elementwise_ops = [ |
@@ -228,6 +227,28 @@ def forward(self, x, y): |
228 | 227 | ] |
229 | 228 | self.run_test_with_dynamic_shape(Op(), input_specs) |
230 | 229 |
|
| 230 | + @parameterized.expand( |
| 231 | + [ |
| 232 | + (f"bf16_{op[0].__name__}_one_constant", op[0]) |
| 233 | + for op in elementwise_ops |
| 234 | + if op[0].__name__ not in ["pow.Tensor_Tensor", "fmod.Tensor"] |
| 235 | + ] |
| 236 | + ) |
| 237 | + def test_elementwise_ops_bf16(self, _, orig_op): |
| 238 | + class TestModule(nn.Module): |
| 239 | + def __init__(self, orig_op): |
| 240 | + super().__init__() |
| 241 | + self.constant = torch.randn(1) |
| 242 | + self.orig_op = orig_op |
| 243 | + |
| 244 | + def forward(self, x): |
| 245 | + x = self.orig_op(x, self.constant) |
| 246 | + return self.orig_op(x, -2) |
| 247 | + |
| 248 | + m = TestModule(orig_op) |
| 249 | + inputs = [torch.randn(2, 2, dtype=torch.bfloat16)] |
| 250 | + self.run_test(m, inputs) |
| 251 | + |
231 | 252 |
|
232 | 253 | if __name__ == "__main__": |
233 | 254 | run_tests() |
0 commit comments