[jax] set_precision throws exception on latest jax version

[javier-garcia-tilburg]

Issue

Problem Description

from torchquad import set_up_backend, MonteCarlo

set_up_backend("jax", data_type="float32")

The code above throws the exception below on latest jax version. The reason is that the jax.config submodule has been removed (previously it was also deprecated)

File <project>\env\Lib\site-packages\torchquad\utils\set_precision.py:58, in set_precision(data_type, backend)
     [56](file://<project>/env/Lib/site-packages/torchquad/utils/set_precision.py:56)     torch.set_default_tensor_type(tensor_dtype)
     [57](file://<project>/env/Lib/site-packages/torchquad/utils/set_precision.py:57) elif backend == "jax":
---> [58](file://<project>/env/Lib/site-packages/torchquad/utils/set_precision.py:58)     from jax.config import config
     [60](file://<project>/env/Lib/site-packages/torchquad/utils/set_precision.py:60)     config.update("jax_enable_x64", data_type == "float64")
     [61](file://<project>/env/Lib/site-packages/torchquad/utils/set_precision.py:61)     logger.info(f"JAX data type set to {data_type}")

ModuleNotFoundError: No module named 'jax.config'

Expected Behavior

It should not throw an exception

What Needs to be Done

Replacing from jax.config import config with from jax import config

How Can It Be Tested or Reproduced

Install latest jax and torchquad and run example above. In my case I used jax-0.4.35 and torchquad 0.4.0.

[javier-garcia-tilburg]

Related to pull request #199

Workaround is to use set_up_backend("jax") and avoid set_precision. Numerical precision can still be configured manually within jax, e.g. jax.config.update("jax_enable_x64", True).