Roy Frostig b1b4915c1c remove opaque dtype aval translation to MLIR types ...

We already have a mapping from opaquely-dtyped avals to basic
"physical" avals, and we can map the latter to MLIR types.

2023-01-27 14:27:30 -08:00

..

__init__.py

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

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

ad.py

Move jax.linear_util to jax._src.linear_util

2022-12-20 14:49:27 -08:00

batching.py

Add batch_jaxpr2 which tells the caller where batch dims are.

2023-01-12 21:16:59 -08:00

mlir.py

remove opaque dtype aval translation to MLIR types

2023-01-27 14:27:30 -08:00

partial_eval.py

Add forwarding support to pjit which was introduced as an optimization. The inputs that are forwarded to outputs are pruned from the outputs of a known_jaxpr.

2023-01-20 18:04:26 -08:00

pxla.py

Add axis_substitution_rules rule for pmap so that pjit(pmap) with an axis_index works properly

2023-01-26 07:33:15 -08:00

xla.py

Error on numpy masked array inputs.

2022-12-27 15:42:49 -08:00