(TestMod, ishape, dtype=None, atol=1e-5, **kwargs)
| 17 | @pytest.mark.skipif(not is_cuda_available(), reason="only support cuda now") |
| 18 | def test_elemwise_activation(): |
| 19 | def tester(TestMod, ishape, dtype=None, atol=1e-5, **kwargs): |
| 20 | dtype = dtype or np.float32 |
| 21 | inp = tensor(0.1 * np.random.randn(*ishape), dtype=dtype) |
| 22 | doup = tensor(0.1 * np.random.randn(*ishape), dtype=dtype) |
| 23 | |
| 24 | gm = GradManager() |
| 25 | mod = TestMod(**kwargs) |
| 26 | |
| 27 | @jit.xla_trace(without_host=True) |
| 28 | def func(mod, inp, doup): |
| 29 | gm.attach(inp) |
| 30 | with gm: |
| 31 | oup = mod(inp) |
| 32 | gm.backward(oup, doup) |
| 33 | return oup, inp.grad |
| 34 | |
| 35 | mge_rsts = func(mod, inp, doup) |
| 36 | xla_rsts = func(mod, inp, doup) |
| 37 | for mge_rst, xla_rst in zip(mge_rsts, xla_rsts): |
| 38 | np.testing.assert_allclose(mge_rst.numpy(), xla_rst.numpy(), atol=atol) |
| 39 | |
| 40 | tester(M.Sigmoid, (2, 3, 4, 5)) |
| 41 | tester(M.ReLU, (2, 3,)) |
no test coverage detected