Junwhan Ahn 6617a0d1ed Expand device_put benchmarks to run with different numbers of arrays and input types ...

For the upcoming batching changes for `device_put`, it is useful to benchmark `device_put` with varying numbers of arrays.

PiperOrigin-RevId: 641716268

2024-06-09 13:01:51 -07:00

..

mosaic

[Mosaic GPU] Add f32 benchmarks for matmul.

2024-06-06 02:31:40 -07:00

api_benchmark.py

Expand device_put benchmarks to run with different numbers of arrays and input types

2024-06-09 13:01:51 -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.

2023-07-27 12:15:58 -07:00

random_benchmark.py

fast dispatch for functions over typed PRNG key arrays

2023-09-13 09:43:58 -07:00

shape_poly_benchmark.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

sparse_benchmark.py

[jax][benchmark] Added clearing caches for benchmarking compilation time in sparse JAX benchmarks

2023-08-02 10:07:54 -07:00