Jake VanderPlas a10f0377db Avoid top-level aliases of jax.tree_util.*

2022-07-07 11:41:02 -07:00

..

api_benchmark.py

Avoid top-level aliases of jax.tree_util.*

2022-07-07 11:41:02 -07:00

benchmark.py

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01: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