def same(list1, list2):
    allclose = functools.partial(jnp.allclose, atol=0.001, rtol=0.001)
    elements_close = list(map(allclose, list1, list2))
    return jnp.all(jnp.array(elements_close))
def same(list1, list2):
    allclose = functools.partial(jnp.allclose, atol=0.001, rtol=0.001)
    elements_close = list(map(allclose, list1, list2))
    return jnp.all(jnp.array(elements_close))