mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/benchmarks
History
Yash Katariya fd09b35645 Optimize make_array_from_callback for fully replicated shardings by going via batched_device_put 
Before:

```
name                                                      cpu/op
bench_make_array_from_callback_fully_replicated_sharding  467µs ± 3%

name                                                      time/op
bench_make_array_from_callback_fully_replicated_sharding  467µs ± 3%
```

After:

```
name                                                      cpu/op
bench_make_array_from_callback_fully_replicated_sharding  28.1µs ± 2%

name                                                      time/op
bench_make_array_from_callback_fully_replicated_sharding  28.1µs ± 2%
```

PiperOrigin-RevId: 572429822
2023-10-10 19:02:04 -07:00
..
api_benchmark.py
Optimize make_array_from_callback for fully replicated shardings by going via batched_device_put
2023-10-10 19:02:04 -07: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
sparse_benchmark.py
[jax][benchmark] Added clearing caches for benchmarking compilation time in sparse JAX benchmarks
2023-08-02 10:07:54 -07:00