Skip to content

FEAT: configure backend precision - #578

Open
redeboer wants to merge 7 commits into
mainfrom
ENH/configure-precision
Open

FEAT: configure backend precision#578
redeboer wants to merge 7 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() API for independently selecting 32-bit or 64-bit precision in JAX and TensorFlow without eagerly importing optional backends.
  • Use 64-bit backend precision by default while respecting JAX_ENABLE_X64 unless an explicit JAX precision is configured.

🐛 Bug fixes

  • Apply the configured precision consistently when creating backend functions, gradients, lambdified expressions, and TensorFlow random values.

📝 Documentation

  • Document when and how to configure backend precision, including initialization timing and environment-variable precedence.

Squash commit messages

* DOC: document backend precision configuration
* FIX: apply configured precision consistently

@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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

🐛 Bug Something isn't working 📝 Docs Improvements or additions to documentation ✨ 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