Peter Hawkins 6a45a9236d Remove the _num_buffers attribute from core.AbstractValue. ...

The number of buffers used to represent an abstract value is a property specific to a particular representation of that abstract value. Currently the only representation is an XLA representation, but that may change in the future. Instead, callers who want to know how XLA would represent an aval should ask the XLA module instead. In this case, we call len(xla.aval_to_xla_shapes(...)) instead.

2021-10-13 14:35:07 -04:00

..

__init__.py

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

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

ad.py

Add reverse-mode AD support for pjit

2021-09-29 07:19:55 -07:00

batching.py

Switch internal users of jax.util.partial to use functools.partial.

2021-09-13 21:09:58 -04:00

invertible_ad.py

Move contents of jax.api_util to jax._src.api_util and add a forwarding shim.

2021-09-08 09:00:56 -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

rewrite remat, leave old implementation for now

2021-10-12 21:39:26 -07:00

pxla.py

ahead-of-time lowering and compilation frontend for pjit

2021-10-12 15:01:46 -07:00

sharded_jit.py

[JAX] Fix more pylint errors.

2021-10-04 16:37:15 -07:00

xla.py

Remove the _num_buffers attribute from core.AbstractValue.

2021-10-13 14:35:07 -04:00