(ishape, tgtshape)
| 15 | @pytest.mark.skipif(not is_cuda_available(), reason="only support cuda now") |
| 16 | def test_broadcast_to(): |
| 17 | def tester(ishape, tgtshape): |
| 18 | dtype = None |
| 19 | dtype = dtype or np.float32 |
| 20 | inp = tensor(np.random.randn(*ishape), dtype=dtype) |
| 21 | dout = tensor(np.random.randn(*tgtshape), dtype=dtype) |
| 22 | |
| 23 | gm = GradManager() |
| 24 | |
| 25 | @jit.xla_trace(without_host=True) |
| 26 | def func(inp, dout): |
| 27 | gm.attach([inp]) |
| 28 | with gm: |
| 29 | out = F.broadcast_to(inp, tgtshape) |
| 30 | gm.backward(out, dout) |
| 31 | return [out, inp.grad] |
| 32 | |
| 33 | mge_rsts = func(inp, dout) |
| 34 | xla_rsts = func(inp, dout) |
| 35 | for mge_rst, xla_rst in zip(mge_rsts, xla_rsts): |
| 36 | np.testing.assert_allclose(mge_rst.numpy(), xla_rst.numpy(), atol=1e-5) |
| 37 | |
| 38 | tester((1, 1, 1), (1, 1, 1, 1)) |
| 39 | tester((1, 1, 1, 1), (1, 1, 1, 1)) |
no test coverage detected