(dim)
| 169 | |
| 170 | @pytest.mark.parametrize("dim", [0, 1, 2]) |
| 171 | def test_stack(dim): |
| 172 | if is_torch_1_9_plus: |
| 173 | wrappers = [ComplexTensor, torch.complex] |
| 174 | modules = [FC, torch] |
| 175 | else: |
| 176 | wrappers = [ComplexTensor] |
| 177 | modules = [FC] |
| 178 | |
| 179 | for complex_wrapper, complex_module in zip(wrappers, modules): |
| 180 | print(complex_wrapper, complex_module) |
| 181 | mat1 = complex_wrapper(torch.rand(2, 3, 4), torch.rand(2, 3, 4)) |
| 182 | mat2 = complex_wrapper(torch.rand(2, 3, 4), torch.rand(2, 3, 4)) |
| 183 | ret = stack([mat1, mat2], dim=dim) |
| 184 | ret2 = complex_module.stack([mat1, mat2], dim=dim) |
| 185 | assert complex_module.allclose(ret, ret2) |
| 186 | |
| 187 | |
| 188 | def test_complex_impl_consistency(): |
nothing calls this directly
no test coverage detected
searching dependent graphs…