return anp.prod(
                anp.stack(anp.meshgrid(*([weights] * dim), like=backend)), axis=0
            ).ravel()

import jax
import jax.numpy as jnp
from torchquad import set_up_backend, MonteCarlo, Gaussian, GaussLegendre

set_up_backend(backend="jax")

@jax.jit
def some_function(x):
    return jnp.power(x[:, 0] - x[:, 1], 2)

g = Gaussian()
# It also fails with GaussLegendre
# g = GaussLegendre()

integral_value = g.integrate(
    lambda x: some_function(x),
    dim=2,
    N=10000,
    integration_domain=jnp.asarray([[-1.0, 1.0], [-1.0, 1.0]]),
)