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

..

_src

Merge pull request #5189 from hawkinsp:herbie

2020-12-14 15:04:05 -08:00

experimental

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

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

image

Move implementation of jax.image to jax._src.image.

2020-10-17 11:37:55 -04:00

interpreters

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

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

lax

Merge pull request #4959 from bchetioui:issue4952

2020-12-01 15:17:39 -08:00

lib

Add an experimental xmap lowering via the SPMD partitioner

2020-11-30 14:40:33 +00:00

nn

Delete jax.nn.functions.

2020-11-09 19:48:54 -05:00

numpy

[impl] Add support for setdiff1d

2020-10-21 20:17:12 +03:00

ops

Move jax.ops implementation into jax._src.ops.

2020-10-17 11:45:28 -04:00

scipy

Cleanup/fixup jax.scipy.sparse.linalg.gmres and expose it publicly.

2020-12-03 09:23:00 -08:00

tools

Merge pull request #5173 from jakevdp:conv-elem-type-cleanup

2020-12-11 15:27:40 -08:00

__init__.py

add jax.device_put_replicated

2020-12-04 12:54:07 -08:00

abstract_arrays.py

Add np.intc to the set of valid jaxtypes.

2020-11-20 13:55:12 -05:00

ad_util.py

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

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

api_util.py

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

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

api.py

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

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

BUILD

Merge pull request #5038 from skye:sharded_jit_namespace

2020-12-07 10:18:56 -08:00

config.py

Reenable pytype for previously problematic ABSL library.

2020-11-19 20:13:16 -08:00

core.py

Fix pmap compilation cache regressions from #4904.

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

custom_derivatives.py

Add support for non-zero (but still not-None) out_axes in pmap

2020-11-24 17:11:38 +00:00

dlpack.py

[JAX] Add support for retaining ownership of DLPack tensors.

2020-10-20 13:07:07 -07:00

dtypes.py

Add weak_type argument to convert_element_type_p

2020-12-10 11:10:21 -08:00

flatten_util.py

deflake

2020-06-12 15:41:49 -07:00

jaxpr_util.py

[JAX] Move source_info_util into jax._src.

2020-11-04 11:54:24 -08:00

lax_reference.py

Fix lax_reference population_count implementation to avoid overflowing Python int -> np.uint32 casts.

2020-11-20 14:50:10 -05:00

lazy.py

Use taggedtuple instead of namedtuple when defining ShardingSpecs

2020-11-24 14:48:11 +00:00

linear_util.py

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

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

profiler.py

[JAX] Move implementation of profiler.py into jax._src.

2020-11-04 09:06:40 -08:00

random.py

DOC: fix jax.random documentation

2020-11-20 16:41:32 -08:00

test_util.py

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

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

tree_util.py

Expose is_leaf predicate for pytree.flatten

2020-12-11 11:26:18 +00:00

util.py

Support mapping a single logical axis to multiple mesh axes in xmap

2020-12-11 14:35:31 +00:00

version.py

update version and changelog for pypi

2020-12-04 21:57:47 -08:00