George Necula 20be478a6e [host_callback] Add support for pmap and for passing the device to tap ...

* Adds support for jit of pmap and pmap of pmap.
* Also adds a `tap_with_device` optional argument to `id_print` and
  `id_tap`, to have the tap function invoked with a device keyword argument.
* Added multiple tests involving pmap

Issue: #5134
Fixes: #5169

2020-12-15 10:46:23 +02:00

..

__init__.py

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

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

ad.py

Fix pmap compilation cache regressions from #4904.

2020-12-02 14:40:45 +00:00

batching.py

Fix pmap compilation cache regressions from #4904.

2020-12-02 14:40:45 +00:00

invertible_ad.py

Pass around the flattened args and outputs in invertible_ad

2020-11-08 19:41:51 +01:00

masking.py

Short-circuit references to jax.core via jax.abstract_arrays.

2020-11-19 14:15:28 -05:00

partial_eval.py

Fix pmap compilation cache regressions from #4904.

2020-12-02 14:40:45 +00:00

pxla.py

[host_callback] Add support for pmap and for passing the device to tap

2020-12-15 10:46:23 +02:00

sharded_jit.py

Fix pmap compilation cache regressions from #4904.

2020-12-02 14:40:45 +00:00

xla.py

[host_callback] Add support for pmap and for passing the device to tap

2020-12-15 10:46:23 +02:00