TensorWaves currently calls jax.config.update("jax_enable_x64", True) in several places. This duplicates backend configuration and prevents users from running TensorWaves in JAX's 32-bit mode, which can be useful for autodiff workloads.
Centralize JAX initialization and make the precision configurable, with the following behavior:
- TensorWaves defaults to 64-bit precision when nothing is specified.
- Users can explicitly select 32-bit precision through a TensorWaves configuration API or
JAX_ENABLE_X64=false.
- TensorWaves does not repeatedly overwrite JAX configuration while creating functions or gradients.
- Configuration is applied before JAX arrays or compiled functions are created.
- Tests cover both precision modes, and the behavior is documented.
Remove the repeated configuration calls from the library and benchmarks in favor of the centralized mechanism.
TensorWaves currently calls
jax.config.update("jax_enable_x64", True)in several places. This duplicates backend configuration and prevents users from running TensorWaves in JAX's 32-bit mode, which can be useful for autodiff workloads.Centralize JAX initialization and make the precision configurable, with the following behavior:
JAX_ENABLE_X64=false.Remove the repeated configuration calls from the library and benchmarks in favor of the centralized mechanism.