Peter Hawkins 74e4db47da Change the default IR dialect returned by .compiler_ir() to MHLO. ...

PiperOrigin-RevId: 423091674

2022-01-20 09:50:17 -08:00

..

image

Added workaround for bug in XLA

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

lax

NotImplementedError for new_checkpoint-of-xmap

2022-01-19 20:40:38 -08:00

lib

[JAX] Add support for generating "equation profiles" in JAX.

2022-01-12 14:27:57 -08:00

nn

Add where= arg to jax.nn.{softmax, log_softmax, normalize}.

2021-12-29 15:49:30 +01:00

numpy

Merge pull request #9184 from jakevdp:unique-nan

2022-01-16 23:57:40 -08:00

ops

[x64] ensure scatter functionality preserves weak_type

2021-11-30 15:43:06 -08:00

scipy

[x64] minor weak_type changes to linalg.py

2021-12-07 16:27:29 -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

Move dtype canonicalization out of core.AbstractValue subclasses.

2021-12-07 06:13:07 -08:00

ad_checkpoint.py

Clean up backwards-compatibility support for remat lowering

2022-01-04 09:53:24 +02:00

ad_util.py

Move jax.ad_util to jax._src.ad_util.

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

api_util.py

Use imports relative to the jax package consistently, rather than .-relative imports.

2021-11-24 07:48:29 -08:00

api.py

Change the default IR dialect returned by .compiler_ir() to MHLO.

2022-01-20 09:50:17 -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

[jax2tf] Add jax2tf_associative_scan_reductions flag

2022-01-13 15:52:18 +02:00

custom_api_util.py

enable custom transformation "stacking"

2022-01-11 17:55:08 -08:00

custom_batching.py

custom vmap: abstract eval and translation rules

2022-01-18 15:48:29 -08:00

custom_derivatives.py

Merge pull request #9175 from froystig:custom-xform-wrappers-forward-attrs

2022-01-12 19:11:22 -08:00

custom_transpose.py

enable custom transformation "stacking"

2022-01-11 17:55:08 -08:00

device_array.py

Move dtype canonicalization out of core.AbstractValue subclasses.

2021-12-07 06:13:07 -08:00

dispatch.py

[JAX] Change signature of .mhlo() method on compiler IR objects to return an ir.Module object instead of its string representation.

2022-01-19 11:04:48 -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

Only use config.x64_enabled as the memo cache key for canonicalize_dtype, not any other fields.

2022-01-18 13:24:15 -08:00

errors.py

Use imports relative to the jax package consistently, rather than .-relative imports.

2021-11-24 07:48:29 -08:00

flatten_util.py

make ravel_pytree unraveler dtype-polymorphic

2021-12-14 14:47:35 -08: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

custom_prng: generalize indexing of PRNGKeyArray

2021-12-20 10:16:32 -08:00

profiler.py

Add profiler annotations

2021-12-06 15:13:35 -08:00

random.py

add random.default_prng_impl to retrieve the default PRNG implementation

2022-01-12 19:13:14 -08:00

source_info_util.py

[JAX] Add support for generating "equation profiles" in JAX.

2022-01-12 14:27:57 -08:00

test_util.py

Benchmarks for GDA. Also move create_global_mesh to test_utils since it was replicated in a lot of places.

2022-01-11 15:43:05 -08:00

traceback_util.py

[jax] completely truncate trivial filtered tracebacks

2022-01-19 19:42:02 -08:00

tree_util.py

Add is_leaf to tree_{leaves,structure}.

2021-12-22 02:56:56 -08:00

util.py

de-duplicate util.memoize and util.cache

2022-01-10 14:28:28 -08:00