Function func

imperative/python/test/unit/xla/functional/test_xla_nn.py:37–50

(x, w, b, dy)

            @jit.xla_trace(without_host=True)
            def func(x, w, b, dy):
                gm.attach([x, w, b])
                with gm:
                    y = F.conv2d(
                        x,
                        w,
                        b,
                        stride=stride,
                        padding=padding,
                        groups=groups,
                        compute_mode=cm,
                    )
                    gm.backward(y, dy)
                return [y, x.grad, w.grad, b.grad]

            mge_rsts = func(x, w, b, dy)
            xla_rsts = func(x, w, b, dy)