FEAT: configure backend precision - #578
Conversation
* 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`
|
Verified downstream with TensorWaves at 2c97d7c: the default is still 64-bit. JAX initially reports I also ran six affected |
✨ New features
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.TFUniformRealNumberGeneratorall agree on the requested dtype.❗ Behavioral changes
JAX_ENABLE_X64is now respected when no explicitjax_precisionis configured, so settingJAX_ENABLE_X64=0results in 32-bit computations where it was previously silently overridden.🖱️ Developer experience
precisionbenchmark 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
raise_missing_module_error()as returningNoReturn. The type checker can now see that the optional-import guards never fall through, which removes thepossibly-unresolved-referenceignore comments that were scattered over the package.📝 Documentation
tensorwaves.configure(), why precision applies from the moment it is set rather than retroactively, and how it interacts withJAX_ENABLE_X64.Squash commit messages