def METHOD_NAME(target, dev):
    # The expression computed is sum((2x - 2y) * (y + z))
    # the gradient of x is broadcast_to(2y + 2z, x.shape)
    # the gradient of y is collapse_sum_to((2x - 4y - 2z), y.shape)
    # the gradient of z is collapse_sum_to((2x - 2y), z.shape)
    # the gradient of u is 0
    @I.ir_module
    class Before:
        @R.function
        def main(
            x: R.Tensor((3, 5), "float32"),
            y: R.Tensor((5,), "float32"),
            z: R.Tensor((5,), "float32"),
            u: R.Tensor((5,), "float32"),
        ):
            with R.dataflow():
                lv1 = R.add(x, x)
                lv2 = R.subtract(lv1, y)
                lv3 = R.subtract(lv2, y)
                lv4 = R.add(y, z)
                lv5 = R.multiply(lv3, lv4)
                lv6 = R.sum(lv5)
                R.output(lv6)
            return lv6
    After = relax.transform.Gradient("main")(Before)
    args = [rand("float32", 3, 5), rand("float32", 5), rand("float32", 5), rand("float32", 5)]
    args_np = [x.numpy() for x in args]
    vm = _legalize_and_build(After, target, dev)
    output, grads = vm["main_adjoint"](*args)
    output_np = np.sum((2 * args_np[0] - 2 * args_np[1]) * (args_np[1] + args_np[2]))
    assert_allclose(output.numpy(), output_np, atol=1e-4)
    expected_grads_nd = [
        (2 * args_np[1] + 2 * args_np[2]) * np.ones_like(args_np[0]),
        np.sum((2 * args_np[0] - 4 * args_np[1] - 2 * args_np[2]), axis=0),
        np.sum((2 * args_np[0] - 2 * args_np[1]), axis=0),
        np.zeros_like(args_np[3]),
    ]
    for i, j in zip(grads, expected_grads_nd):
        assert_allclose(i.numpy(), j, atol=1e-4)