Function func

imperative/python/test/unit/xla/module/test_elemwise.py:28–33

(mod, inp, doup)

        @jit.xla_trace(without_host=True)
        def func(mod, inp, doup):
            gm.attach(inp)
            with gm:
                oup = mod(inp)
                gm.backward(oup, doup)
            return oup, inp.grad

        mge_rsts = func(mod, inp, doup)
        xla_rsts = func(mod, inp, doup)