Yash Katariya 418c2f9d2a Rename in_axis_resources and out_axis_resources with in_shardings and out_shardings. This is just a simple name replacement. It does not change any of the current pjit semantics and doesn't break any code. ...

This is a safe and trivial name replacement. It does not change any of the semantics. You can still pass in PatitionSpecs to in_shardings and out_shardings.

PiperOrigin-RevId: 510671300

2023-02-18 10:00:36 -08:00

..

api_benchmark.py

Rename in_axis_resources and out_axis_resources with in_shardings and out_shardings. This is just a simple name replacement. It does not change any of the current pjit semantics and doesn't break any code.

2023-02-18 10:00:36 -08:00

benchmark.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

math_benchmark.py

benchmarks: add math unary benchmarks

2023-01-23 08:17:16 -08:00

pmap_benchmark.py

[JAX] Replace uses of jax.interpreters.pxla.ShardedDeviceArray with jax.Array.

2023-02-09 13:39:41 -08:00