Peter Hawkins d6c67c97db Remove redundant dtype canonicalization from jax.device_put(). ...

Gives a small improvement to the included jax.device_put() benchmark on my VM:

```
name        old cpu/op  new cpu/op  delta
device_put  91.3µs ± 5%  80.1µs ± 3%  -12.29%  (p=0.008 n=5+5)

name        old time/op             new time/op             delta
device_put  91.4µs ± 5%             80.1µs ± 3%  -12.29%          (p=0.008 n=5+5)
```

jax.device_put() has not been optimized that much yet and there is plenty of room for further improvement.

PiperOrigin-RevId: 491727173

2022-11-29 13:47:36 -08:00

..

api_benchmark.py

Remove redundant dtype canonicalization from jax.device_put().

2022-11-29 13:47: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

gda_benchmark.py

Make MeshPspecSharding an alias for NamedSharding (it was the other way around before this CL).

2022-11-14 14:44:00 -08:00

pmap_benchmark.py

Replace jax.xla.DeviceArray private type with the new public type jax.Array.

2022-09-28 16:34:10 -07:00