Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Handle double precision in tests more carefully #288

Open
Michael-T-McCann opened this issue May 4, 2022 · 1 comment
Open

Handle double precision in tests more carefully #288

Michael-T-McCann opened this issue May 4, 2022 · 1 comment
Labels
tests Pertaining to SCICO tests

Comments

@Michael-T-McCann
Copy link
Contributor

Michael-T-McCann commented May 4, 2022

JAX works in single precision by default, and won't even let you create double precision arrays unless an environment variable is set (JAX_ENABLE_X64=True) or a special command is run when jax is imported (config.update("jax_enable_x64", True)). In order to test double precision, these commands are used various places in the tests.

Unfortunately, enabling double precision also makes it the default for new arrays, creating situations where tests have different behavior when run on their own versus in the whole suite (because the config is "sticky" and setting it in one test affects others).

All of this may change in a future JAX release (google/jax#8178), but for now, I propose running all tests with JAX_ENABLE_X64=True JAX_DEFAULT_DTYPE_BITS=32 and removing any config.updates in test files.

@bwohlberg bwohlberg added the tests Pertaining to SCICO tests label May 6, 2022
@bwohlberg
Copy link
Collaborator

bwohlberg commented Apr 28, 2023

While there has been no further discussion in google/jax#8178 for more than a year, it appears that deprecation of the X64 flag is still being considered. Instead of simply running all the tests with X64 enabled and X32 default, perhaps we should configure that state in scico/__init__.py so that it applies across the code base?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
tests Pertaining to SCICO tests
Projects
None yet
Development

No branches or pull requests

2 participants