rocm_jax

mirror of https://github.com/ROCm/jax.git synced 2025-04-15 19:36:06 +00:00

History

Peter Hawkins 26632fd344 Replace disable_backends with enable_backends on jax_multiplatform_test.

Most users of disable_backends were actually using it to enable only a single backend. So things are simpler if we negate the sense of the option to say that. Change disable_configs to enable_configs, with a default `None` value meaning "everything is enabled".

We change the relationship between enable_backends, disable_configs, enable_configs to be the following:
* `enable_backends` selects a set of initial test configurations to enable, based off backend only.
* `disable_configs` then prunes that set of test configurations, removing elements from the set.
* `enable_configs` then adds additional configurations to the set.

Fix code in jax/experimental/mosaic/gpu/examples not to depend on a Google-internal GPU support target.

PiperOrigin-RevId: 679563155

2024-09-27 06:15:31 -07:00

mosaic

Replace disable_backends with enable_backends on jax_multiplatform_test.

2024-09-27 06:15:31 -07:00

api_benchmark.py

Fixes to api_benchmark.py. Testcases always fail without these fixes.

2024-08-27 23:37:49 -07:00

linalg_benchmark.py

Use a Jacobi SVD solver for unbatched SVDs up to 1024x1024 on NVIDIA GPUs.

2023-11-13 12:04:13 -08:00

math_benchmark.py

Update flags to use the ABSL typed flag API.