Yash Katariya 6a7a34603d Move PartitionSpec from sharded_jit.py to pxla.py. The public endpoint is via jax.experimental so that should be used (no changes to the public endpoint). ...

This move is because sharded_jit is being deprecated.

PiperOrigin-RevId: 439948391

2022-04-06 15:19:19 -07:00

..

api_benchmark.py

Add small and big matmul to api_benchmarks.

2022-03-17 14:51:17 -07:00

benchmark.py

[JAX] Fix pylint errors.

2021-10-04 17:54:46 -07:00

gda_benchmark.py

Move PartitionSpec from sharded_jit.py to pxla.py. The public endpoint is via jax.experimental so that should be used (no changes to the public endpoint).

2022-04-06 15:19:19 -07:00

pmap_benchmark.py

Revert import path changes to examples/ and benchmarks/

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