Kuangyuan Chen b764aadbcf Add pmap benchmarks using jax.Array ...

PiperOrigin-RevId: 473325120

2022-09-09 13:10:42 -07:00

..

api_benchmark.py

Add pmap benchmarks using jax.Array

2022-09-09 13:10:42 -07:00

benchmark.py

target_total_secs has type int but used as type None

2022-08-22 15:46:38 +02:00

gda_benchmark.py

Fix the gda_xla_sharding_match benchmark which was regressing. This was happening because that function was executed from top to bottom a couple of times and each time a new mesh object was created violating the already created cache which doesn't happen in real life.

2022-07-27 23:08:55 -07:00

pmap_benchmark.py

Revert import path changes to examples/ and benchmarks/

2021-01-20 17:35:55 -08:00