Skip to content

FEAT: configure backend precision - #578

Merged
redeboer merged 8 commits into
mainfrom
ENH/configure-precision
Aug 7, 2026
Merged

FEAT: configure backend precision#578
redeboer merged 8 commits into
mainfrom
ENH/configure-precision

Conversation

@redeboer

@redeboer redeboer commented Aug 6, 2026

Copy link
Copy Markdown
Member

✨ New features

  • Add a public tensorwaves.configure() entry point for independently selecting 32-bit or 64-bit precision in JAX and TensorFlow. The call never imports a backend itself: a backend that is already imported is reconfigured immediately, and one that is not is configured when TensorWaves first uses it.
  • Apply the configured precision everywhere TensorWaves touches a backend, so that lambdified expressions, backend functions, JAX gradients, and TFUniformRealNumberGenerator all agree on the requested dtype.

❗ Behavioral changes

  • TensorWaves no longer forces 64-bit precision onto JAX unconditionally. 64-bit remains the default, but JAX_ENABLE_X64 is now respected when no explicit jax_precision is configured, so setting JAX_ENABLE_X64=0 results in 32-bit computations where it was previously silently overridden.

🖱️ Developer experience

  • Add a precision benchmark group that evaluates a lambdified expression over a million points for each backend and precision combination, and asserts that the requested precision was indeed applied to both the input arrays and the result.

🔨 Maintenance

  • Annotate raise_missing_module_error() as returning NoReturn. The type checker can now see that the optional-import guards never fall through, which removes the possibly-unresolved-reference ignore comments that were scattered over the package.

📝 Documentation

  • Add a Backend precision section to the installation page that explains when to call tensorwaves.configure(), why precision applies from the moment it is set rather than retroactively, and how it interacts with JAX_ENABLE_X64.

Squash commit messages

* BEHAVIOR: respect JAX_ENABLE_X64 environment variable
* DOC: document backend precision configuration
* DX: add benchmark for backend precision
* MAINT: annotate raise_missing_module_error as NoReturn

@redeboer redeboer added this to the 0.4.17 milestone Aug 6, 2026
@redeboer redeboer self-assigned this Aug 6, 2026
@redeboer redeboer added ⚙️ Enhancement Improvements and optimizations of existing features ✨ Feature New feature added to the package labels Aug 6, 2026
* BEHAVIOR: stop enabling TensorFlow numpy behavior in `TFUniformRealNumberGenerator`
* DOC: describe when backend precision takes effect
* FIX: configure precision for `tuple` and `dict` backends
* FIX: read the value of `JAX_ENABLE_X64` instead of testing for the key
* FIX: validate all `configure()` arguments before applying any of them
* MAINT: annotate `raise_missing_module_error()` as `NoReturn`
* MAINT: confine benchmark precision configuration to the JAX fixtures
* MAINT: move configuration tests to `tests/test_config.py`
@redeboer redeboer changed the title ENH: set JAX precision through configuration FEAT: configure backend precision Aug 6, 2026
@redeboer redeboer added 📝 Docs Improvements or additions to documentation 🐛 Bug Something isn't working and removed ⚙️ Enhancement Improvements and optimizations of existing features labels Aug 6, 2026
@redeboer

redeboer commented Aug 6, 2026

Copy link
Copy Markdown
Member Author

Verified downstream with TensorWaves at 2c97d7c: the default is still 64-bit. JAX initially reports jax_enable_x64=False, but first TensorWaves JAX-backend initialization switches it to True; resulting arrays are float64/complex128. Setting JAX_ENABLE_X64=0 before startup instead produces float32/complex64 as intended.

I also ran six affected jpsi-nstar test fits in both modes. Aggregate wall time was about 612 s (FP64) versus 522 s (FP32), nominally 1.17x faster. That number is misleading, however: the largest speed-up coincided with a worse minimum (-2969.72 vs -2461.82), and xNES landed in a completely different basin (-1480.62 vs +2034.93). Across the four fits with comparable NLLs, FP32 was not faster overall (about 258 s vs 252 s). For these CPU fits, FP32 therefore offers no convincing end-to-end benefit at equivalent fit quality and noticeably reduces optimizer stability; keeping FP64 as the default looks appropriate.

@redeboer redeboer added 🔨 Maintenance Maintenance and upkeep improvements ❗ Behavior Changes that may affect the framework output 🖱️ DX Improvements to the Developer Experience and removed 🐛 Bug Something isn't working 🔨 Maintenance Maintenance and upkeep improvements 📝 Docs Improvements or additions to documentation labels Aug 7, 2026
@redeboer
redeboer merged commit dadc908 into main Aug 7, 2026
33 checks passed
@redeboer
redeboer deleted the ENH/configure-precision branch August 7, 2026 07:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

❗ Behavior Changes that may affect the framework output 🖱️ DX Improvements to the Developer Experience ✨ Feature New feature added to the package

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Make TensorFlow precision configurable while defaulting to 64-bit Make JAX precision configurable while defaulting to 64-bit

1 participant