rocm_jax

mirror of https://github.com/ROCm/jax.git synced 2025-04-13 02:16:06 +00:00

History

Hyeontaek Lim 73b8f6aee2 [JAX] Clean up make_array_from_callback_* API benchmarks and add a partially replicated sharding variant

To prepare for the upcoming `BatchedDevicePut` implementation changes, this
change makes `make_array_from_callback_*` benchmark code to be more
homogeneous. Also it adds a variant that uses a partially replicated sharding.

PiperOrigin-RevId: 736665856

2025-03-13 15:50:46 -07:00

mosaic

Replace disable_backends with enable_backends on jax_multiplatform_test.

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

api_benchmark.py

[JAX] Clean up make_array_from_callback_* API benchmarks and add a partially replicated sharding variant

2025-03-13 15:50:46 -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.