rocm_jax

clusters

Don't set TPU topology env vars in cloud_tpu_init.py

2023-01-13 22:02:16 +00:00

debugger

Remove typing_extensions dependency

2022-12-05 15:42:26 -08:00

image

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

interpreters

Merge pull request #14418 from mattjj:vmap-spmd-axis-name-tuples

2023-02-10 16:08:32 -08:00

lax

migrate internal dependencies from jax.interpreters.batching to jax._src.interpreters.batching

2023-02-09 15:11:57 -08:00

lib

Add back loading TPU plugin for older jaxlib versions.

2023-02-10 16:16:20 -08:00

nn

Added shape error checking for compute_fans

2023-01-18 20:59:11 -08:00

numpy

jax.numpy reductions: better validation of initial value

2023-02-13 08:43:25 -08:00

ops

Remove typing_extensions dependency

2022-12-05 15:42:26 -08:00

scipy

added scipy special spence

2023-02-10 20:33:47 +00:00

state

simpler pretty-print for pjit, tweak custom pp rule signature

2023-02-09 12:45:51 -08:00

third_party

Address drastic slowdown in mypy runtime

2022-12-16 14:48:26 -08:00

__init__.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

abstract_arrays.py

Expose fp8 in jax dtypes and mlir builder.

2023-01-13 18:12:12 -08:00

ad_checkpoint.py

migrate internal dependencies from jax.interpreters.batching to jax._src.interpreters.batching

2023-02-09 15:11:57 -08:00

ad_util.py

Merge pull request #13693 from jakevdp:typing-ad-util

2022-12-20 16:50:11 -08:00

api_util.py

Make error_test a jax_test so that we can test other configs and fix it with jit/pjit merge.

2023-01-17 18:43:05 -08:00

api.py

Add in_shardings and out_shardings argument to pjit and jit to start deprecating in_axis_resources and out_axis_resources.

2023-02-11 15:30:14 -08:00

array.py

[Rollback] Convert _arrays to return PyArray instead of PyBuffer.

2023-02-10 21:36:56 -08:00

basearray.py

migrate internal dependencies from jax.core to jax._src.core

2022-12-16 21:00:14 -08:00

basearray.pyi

[typing] define additional methods & properties on jax.Array

2023-02-10 09:42:32 -08:00

callback.py

migrate internal dependencies from jax.interpreters.batching to jax._src.interpreters.batching

2023-02-09 15:11:57 -08:00

checkify.py

migrate internal dependencies from jax.interpreters.batching to jax._src.interpreters.batching

2023-02-09 15:11:57 -08:00

cloud_tpu_init.py

Add environment variable check in additional to libtpu check for cloud tpu vm

2023-01-25 09:51:37 -08:00

config.py

Remove jax_experimental_subjaxpr_lowering_cache since it was only for jit and was False by default. Now that jit/pjit are merged, this cache is not needed since pjit does the caching and we get it for free.

2023-02-08 14:55:56 -08:00

core.py

[typing] define additional methods & properties on jax.Array

2023-02-10 09:42:32 -08:00

custom_api_util.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

custom_batching.py

migrate internal dependencies from jax.interpreters.batching to jax._src.interpreters.batching

2023-02-09 15:11:57 -08:00

custom_derivatives.py

migrate internal dependencies from jax.interpreters.batching to jax._src.interpreters.batching

2023-02-09 15:11:57 -08:00

custom_transpose.py

migrate internal dependencies from jax.interpreters.ad to jax._src.interpreters.ad

2023-02-06 22:52:36 -08:00

debugging.py

migrate internal dependencies from jax.interpreters.batching to jax._src.interpreters.batching

2023-02-09 15:11:57 -08:00

deprecations.py

Deprecate jax.interpreters.xla.Device and jax.interpreters.xla.DeviceArray.

2023-02-09 16:11:48 -08:00

device_array.py

migrate internal dependencies from jax.core to jax._src.core

2022-12-16 21:00:14 -08:00

dispatch.py

Minimize the set of names exported from jax.experimental.pjit.

2023-02-11 07:37:32 -08:00

distributed.py

Add Open MPI automatic distributed initialization

2023-01-11 17:08:09 -06:00

dlpack.py

migrate internal dependencies from jax.core to jax._src.core

2022-12-16 21:00:14 -08:00

dtypes.py

Expose fp8 in jax dtypes and mlir builder.

2023-01-13 18:12:12 -08:00

environment_info.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

errors.py

migrate internal dependencies from jax.core to jax._src.core

2022-12-16 21:00:14 -08:00

flatten_util.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

iree.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

lax_reference.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

lazy_loader.py

Avoid loading scipy eagerly.

2022-10-12 19:51:09 +00:00

linear_util.py

Move jax.linear_util to jax._src.linear_util

2022-12-20 14:49:27 -08:00

monitoring.py

Move contents of jax.monitoring to jax._src.monitoring.

2023-02-08 15:03:22 -08:00

path.py

Make the etils[epath] dependency of JAX optional.

2022-10-12 21:16:48 +00:00

pjit.py

Change the axis_resources argument of with_sharding_constraint to shardings to match pjit and jit.

2023-02-13 10:53:57 -08:00

pretty_printer.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

prng.py

migrate internal dependencies from jax.interpreters.batching to jax._src.interpreters.batching

2023-02-09 15:11:57 -08:00

profiler.py

Add a number of missing function cross-references in the docs.

2022-11-07 12:00:26 -05:00

public_test_util.py

Roll-back https://github.com/google/jax/pull/14144 due to downstream test failures

2023-01-25 12:15:36 -08:00

random.py

migrate internal dependencies from jax.interpreters.batching to jax._src.interpreters.batching

2023-02-09 15:11:57 -08:00

sharding.py

[Rollback] Convert _arrays to return PyArray instead of PyBuffer.

2023-02-10 21:36:56 -08:00

source_info_util.py

Let is_user_frame ignore frames from stdlib.

2022-11-08 10:50:08 -08:00

stages.py

Move jax.interpreters.xla to jax._src.interpreters.xla.

2023-02-07 15:01:32 -08:00

test_util.py

Move jax.interpreters.mlir to jax._src.interpreters.mlir.

2023-02-08 14:39:01 -08:00

traceback_util.py

add blank line, mainly to trigger/test source sync

2023-02-01 13:56:29 -08:00

tree_util.py

Use jnp.arange to break the pjit cache (when jit and pjit are merged) because pytest runs tests non-hermetically.

2023-02-08 10:17:37 -08:00

typing.py

Remove typing_extensions dependency

2022-12-05 15:42:26 -08:00

util.py

[Rollback] Convert _arrays to return PyArray instead of PyBuffer.

2023-02-10 21:36:56 -08:00