Function func

imperative/python/test/unit/xla/functional/test_xla_math.py:38–43

(lhs, rhs, dout)

        @jit.xla_trace(without_host=True)
        def func(lhs, rhs, dout):
            gm.attach([lhs, rhs])
            with gm:
                out = F.matmul(lhs, rhs, lhs_transpose, rhs_transpose)
                gm.backward(out, dout)
            return out, lhs.grad, rhs.grad

        mge_rsts = func(lhs, rhs, dout)
        xla_rsts = func(lhs, rhs, dout)