Roy Frostig 6071a8f875 roll-forward #11952, take 2 ...

Now with:
* resetting the `random.PRNGKeyArray` type during Python typechecks
* zeroing JVP rules for random primitives
* temporarily skipping vmap-of-pmap test with keys under `config.jax_array`

PiperOrigin-RevId: 469276609

2022-08-22 13:57:31 -07:00

..

__init__.py

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

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

ad.py

Merge pull request #11900 from sharadmv:for-loop

2022-08-18 20:14:17 -07:00

batching.py

Add spmd_axis_name to vmap to allow constraining mapped PartitionSpecs.

2022-08-08 19:41:42 -07:00

mlir.py

roll-forward #11952, take 2

2022-08-22 13:57:31 -07:00

partial_eval.py

rolling forward deletion of custom_jvp_call_jaxpr_p yet again...

2022-08-18 14:02:40 -07:00

pxla.py

roll-forward #11952, take 2

2022-08-22 13:57:31 -07:00

xla.py

BUG: ensure that boolean scalars are never marked weak

2022-07-07 15:41:23 -07:00