mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 03:46:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
26,443 Commits 500 Branches 263 Tags
Commit Graph

15 Commits

Author SHA1 Message Date
Dan Foreman-Mackey
d42e3650d0 Handle effects in lax.custom_linear_solve. 2025-02-03 11:14:48 -05: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
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
Jake VanderPlas
cddee4654c tests: access tree utilities via jax.tree.* 2024-02-26 14:17:18 -08:00
Peter Hawkins
1885c4933c Add a new internal test utility test_device_matches() and use it instead of equality tests on device_under_test(). 
This change prepares for allowing more flexible tag matching. For example, we may want to write "gpu" in a test and have it match both "cuda" and "rocm" devices, which we cannot do under the current API but can easily do under this design.

Replace uses of device_under_test() in a context that performs an equality test with a call to test_device_matches().
Replace uses of if_device_under_test() with test_device_matches() and delete if_device_under_test().

PiperOrigin-RevId: 568923117
 2023-09-27 12:10:43 -07:00
Jake VanderPlas
fbe4f10403 Change to simpler import for jax.config 2023-04-21 11:51:22 -07:00
Nicholas Junge
27b26515fe Add regression test 2023-03-07 09:30:03 +01:00
Filippo Vicentini
793fb9b22c Fix issue in check_tree, so that custom_linear_solve supports 
hax_aux=True when the vector and the aux are both pytrees.
 2022-11-07 21:29:43 +01:00
jax authors
363cc124e3 Merge pull request #12197 from ROCmSoftwarePlatform:fixedRocmUnitTestsSkip 
PiperOrigin-RevId: 479566021
 2022-10-07 06:36:11 -07: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
Chao Chen
1c69f594fb testSphHarmOrderZeroDegreeOne and test_custom_linear_solve_cholesky have been fixed in ROCm, no need to skip 2022-09-01 13:27:23 +00:00
Matthew Johnson
aa043a60b6 add test for custom_linear_solve + new remat 2022-07-28 21:37:23 -07:00
Peter Hawkins
62a392a7e2 Relax test tolerances. 
These tests current fail on M1 Mac.
 2022-07-03 15:13:28 -04:00
Reza Rahimi
a0d9d81f92 Update JAX to use new math libraries in ROCm-5.0. 2022-03-01 20:02:15 +00:00
Peter Hawkins
b509aae2a2 Split lax_control_flow_test into three separate tests. 
Split the custom root and custom linear solve tests into separate test files.

Disable two slow custom linear solve tests.
Add a few jit decorators to slow tests in lax_control_flow_test.

PiperOrigin-RevId: 421901487
 2022-01-14 13:36:34 -08:00