Qiao Zhang 619f2b93f5 Set ShardedArray device_buffer aval. ...

PiperOrigin-RevId: 393281615

2021-08-26 22:37:35 -07:00

..

__init__.py

[JAX] Move traceback_util.py into jax._src.

2020-11-04 09:02:59 -08:00

ad.py

custom_jvp/vjp tweaks and fixes

2021-08-17 17:51:35 -07:00

batching.py

custom_jvp/vjp tweaks and fixes

2021-08-17 17:51:35 -07:00

invertible_ad.py

Move contents of jax.custom_derivatives to jax._src.custom_derivatives.

2021-04-19 17:51:49 -07:00

masking.py

[masking] Remove references to masking.Poly from the lax.py and lax_numpy.py

2021-04-08 17:45:14 +03:00

partial_eval.py

Most of JAX now uses concrete types for things defined in jaxlib.xla_client

2021-08-16 20:33:36 +01:00

pxla.py

Set ShardedArray device_buffer aval.

2021-08-26 22:37:35 -07:00

sharded_jit.py

Lazy eval for vlogging on pmap/pjit critical path.

2021-06-25 19:35:16 -07:00

xla.py

Specify weak_type in DeviceArray repr

2021-08-23 13:19:33 -07:00