mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-24 07:36:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
26,928 Commits 508 Branches 263 Tags
Commit Graph

12 Commits

Author SHA1 Message Date
Sergei Lebedev
0ff234049b Removed trivial docstrings from JAX tests 
These docstrings do not make the tests any more clear and typically just duplicate the test module name.

PiperOrigin-RevId: 737611977
 2025-03-17 07:49:37 -07:00
Jake VanderPlas
f090074d86 Avoid 'from jax import config' imports 
In some environments this appears to import the config module rather than
the config object.
 2024-04-11 13:23:27 -07:00
Yue Sheng
fea26659fe Clean up usage of jax.clear_backends 
PiperOrigin-RevId: 616207814
 2024-03-15 12:17:07 -07:00
Yash Katariya
4ddf6a9a54 Bump minimum_jaxlib_version to 0.4.14. xla_extension_version is 174 and mlir_api_version is 54 
PiperOrigin-RevId: 552816893
 2023-08-01 08:53:28 -07:00
Skye Wanderman-Milne
529822cbcf Fix jax.clear_backends() on Cloud TPU 
PiperOrigin-RevId: 542279072
 2023-06-21 09:29:02 -07:00
Jake VanderPlas
fbe4f10403 Change to simpler import for jax.config 2023-04-21 11:51:22 -07:00
Skye Wanderman-Milne
ef5e4a4035 Remove 'pjrt_c_api_unimplemented' pytest mark. 
Instead, we skip tests that the PJRT C API doesn't support. We had
this tag for feature development so it was easy to broadly disable,
but now we don't expect to need to do that.
 2023-03-24 23:14:54 +00:00
Peter Hawkins
f66f6ec98a [JAX] Move jax._src.lib.xla_bridge to jax._src.xla_bridge. 
Limit jax._src.lib to shims around jaxlib and nothing else.

The goal of this change is to avoid a dependency cycle between the rest of jax and jax._src.lib in a Bazel build. This allows the types for jax._src.lib to be inferred by pytype in isolation without referring to the rest of JAX.

PiperOrigin-RevId: 512922397
 2023-02-28 07:01:57 -08:00
Skye Wanderman-Milne
f90b5eed52 Add pjrt_c_api_unimplemented pytest marker to skip unsupported tests. 
Also adds `test_util.pytest_mark_if_available` helper function.
 2023-01-12 22:17:23 +00:00
Peter Hawkins
ba557d5e1b Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.". 
See https://opensource.google/documentation/reference/releasing/contributions#copyright for more details.

PiperOrigin-RevId: 476167538
 2022-09-22 12:27:19 -07:00
Peter Hawkins
6c59d72c75 Bump the minimum jaxlib version to 0.3.15. 2022-09-08 16:43:46 -04:00
Kuangyuan Chen
c0ec3b33e6 Introduce jax.experimental.clear_backends to delete all JAX runtime backends. 
In cases like unit tests, users may want to clean up all the backends along with the resources used in the end of the test, and reinitialize them in the next test.

PiperOrigin-RevId: 462239974
 2022-07-20 15:10:27 -07:00