Peter Hawkins d262bae88b Split jax.interpreters.xla up into three pieces: ...

* jax._src.device_array, which contains the definition of DeviceArray.
* jax.interpreters.xla, which contains code for lowering jaxprs into XLA computations.
* jax._src.dispatch, which contains code for executing primitives and jit-compiled functions (xla_call_p's impl logic).

The purpose of splitting up this file is that I would like to treat jax.interpreters.mlir lowering as an alternative to jax.interpreters.xla, but we wish to share the device_array and computation dispatch pieces. Currently jax.interpreters.mlir duplicates most of the dispatch logic. (That refactoring is for a future change; this change just moves the existing code around.)

PiperOrigin-RevId: 411565432

2021-11-22 08:22:43 -08:00

..

image

Added workaround for bug in XLA

2021-11-18 11:01:50 +02:00

lax

Split jax.interpreters.xla up into three pieces:

2021-11-22 08:22:43 -08:00

lib

Add a prototype IREE backend for JAX.

2021-11-15 07:57:04 -08:00

nn

Tighten up dtypes across the package

2021-10-29 13:50:30 -07:00

numpy

Split jax.interpreters.xla up into three pieces:

2021-11-22 08:22:43 -08:00

ops

[JAX] Update users of jax.ops.index... functions, which are deprecated.

2021-10-19 16:39:09 -07:00

scipy

Document jax.lax.Precision

2021-11-08 14:15:31 -08:00

third_party

Move jax.third_party to jax._src.third_party.

2020-10-17 11:43:07 -07:00

__init__.py

Move most of the implementation of jax.scipy into jax._src.scipy.

2020-10-16 17:04:25 -04:00

abstract_arrays.py

Ensure zeros from AD are generated on device.

2021-07-26 20:40:13 +03:00

ad_checkpoint.py

Port remaining translation rules inside JAX to new style.

2021-10-19 09:48:37 -07:00

ad_util.py

Move jax.ad_util to jax._src.ad_util.

2021-06-07 14:51:34 -07:00

api_util.py

Disallow non-hashable static arguments in pmap().

2021-09-30 15:50:07 -04:00

api.py

Split jax.interpreters.xla up into three pieces:

2021-11-22 08:22:43 -08:00

cloud_tpu_init.py

Changes to make jax[tpu] work better in a docker container.

2021-07-12 17:42:46 -07:00

config.py

Internal change

2021-11-12 22:42:15 -08:00

custom_derivatives.py

Port remaining translation rules inside JAX to new style.

2021-10-19 09:48:37 -07:00

device_array.py

Split jax.interpreters.xla up into three pieces:

2021-11-22 08:22:43 -08:00

dispatch.py

Split jax.interpreters.xla up into three pieces:

2021-11-22 08:22:43 -08:00

distributed.py

Add jax.distributed.initialize for multi-host GPU.

2021-10-26 14:37:54 -07:00

dlpack.py

Split jax.interpreters.xla up into three pieces:

2021-11-22 08:22:43 -08:00

dtypes.py

[x64] set scalar defaults to 64-bit types

2021-11-12 15:09:17 -08:00

errors.py

[JAX] Update users of jax.ops.index... functions, which are deprecated.

2021-10-28 09:54:26 -07:00

flatten_util.py

Move implementation of jax.flatten_util to jax._src.flatten_util. Add a jax.flatten_util shim.

2021-09-08 13:54:25 -07:00

iree.py

Add a prototype IREE backend for JAX.

2021-11-15 07:57:04 -08:00

lax_reference.py

Improve implementation of cbrt() in JAX.

2021-07-22 14:01:28 -07:00

pretty_printer.py

Adds a Wadler-Lindig pretty printer.

2021-09-27 21:09:24 -04:00

prng.py

Overload jnp ops which are polymorphic to an array's value and support PRNGKeys.

2021-11-16 23:00:32 +00:00

profiler.py

Move contents of jax.lib to jax._src.lib.

2021-09-23 06:33:55 -07:00

random.py

jax.random.permutation: add independent keyword

2021-11-02 11:39:41 -07:00

source_info_util.py

Fix forward for PR #8392 (made source_info for new_jaxpr_eqn argument optional again)

2021-10-29 15:50:14 -07:00

test_util.py

Split jax.interpreters.xla up into three pieces:

2021-11-22 08:22:43 -08:00

traceback_util.py

Move contents of jax.lib to jax._src.lib.

2021-09-23 06:33:55 -07:00

tree_util.py

Fix jacfwd and jacrev for heterogeneous pytrees

2021-10-12 19:41:47 -04:00

util.py

Fix a number of bugs in MLIR translation rules.

2021-11-17 07:20:56 -08:00