def METHOD_NAME(m):
    a = torch.tensor([1.0, 2.0, 3.0])
    b = torch.tensor([4.0, 5.0, 6.0])
    if dr.is_cuda_v(m.Float):
        a = a.cuda()
        b = b.cuda()
    a.requires_grad = True
    b.requires_grad = True
    @dr.wrap_ad(source='torch', target='drjit')
    def func(a, b):
        return a * 4, b * 3
    c, d = func(a, b)
    e = (c + d).sum()
    e.backward()
    assert dr.allclose(m.Float(c), [4, 8, 12])
    assert dr.allclose(m.Float(d), [12, 15, 18])
    assert dr.allclose(m.Float(a.grad), [4, 4, 4])
    assert dr.allclose(m.Float(b.grad), [3, 3, 3])