mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-15 19:36:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
25,441 Commits 500 Branches 263 Tags
Commit Graph

9 Commits

Author SHA1 Message Date
Bill Varcho
0abd9538ce [JAX] disable flaky parameter permutations for sparse_bcoo_bcsr test. 
PiperOrigin-RevId: 722832212
 2025-02-03 16:02:06 -08:00
Peter Hawkins
70f91db853 Set PYTHONWARNINGS=error in bazel tests. 
The goal of this change is to catch PRs that introduce new warnings sooner.

To help pass the environment variable more easily, rename the jax_test Bazel test macro to jax_multiplatform_test, and introduce a new jax_py_test macro that wraps py_test. Add code to both to set the environment variable.

Add code to suppress some new warnings uncovered in CI.

PiperOrigin-RevId: 678352286
 2024-09-24 12:30:11 -07:00
Michael Hudgins
d4d1518c3d Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax 
PiperOrigin-RevId: 676843138
 2024-09-20 07:52:33 -07:00
Ruturaj4
637370baa0 [ROCM] Fix version checks after rocm pjrt integration 2024-07-12 08:07:52 -05:00
Jake VanderPlas
c099f714cc Fix bug in bcoo_spdot_general abstract_eval 2024-06-25 10:44:55 -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
Meekail Zain
a7737ca618 Clean up sparse test run conditions 2024-04-09 23:16:12 +00:00
Sergei Lebedev
cbcaac2756 MAINT Migrate remaining internal/test modules to use state objects 
The motivation here is to gradually replace all dynamic lookups on `jax.config`
with statically-typed state objects, which are more type checker/IDE friendly.

This is a follow up to #18008.
 2023-10-12 17:32:15 +01:00
jax authors
1c37f5091c sparse_test: Split into two so that each target is small enough to fit within a medium timeout. 
PiperOrigin-RevId: 570882867
 2023-10-04 19:59:03 -07:00