rocm_jax

clusters

[JAX] Move jax._src.lib.xla_bridge to jax._src.xla_bridge.

2023-02-28 07:01:57 -08:00

debugger

Replace uses of jnp.ndarray with jax.Array inside JAX.

2023-02-15 14:53:00 -08:00

image

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

2023-02-14 23:01:11 -08:00

internal_test_util

lax_test: Create a separate module for lax-specific test utils in a new package.

2023-02-16 15:29:41 -08:00

interpreters

Internal change

2023-03-04 13:24:11 +00:00

lax

[shape_poly] Add support for jnp.cum{sum,prod,max,min} with shape polymorphism

2023-03-03 10:22:45 +01:00

lib

[JAX] Move jax._src.lib.xla_bridge to jax._src.xla_bridge.

2023-02-28 07:01:57 -08:00

nn

Replace jax._src.util.prod with math.prod.

2023-02-28 12:41:00 -08:00

numpy

Reenable pytype for numpy ufuncs.

2023-03-03 05:01:03 -08:00

ops

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

2023-02-14 23:01:11 -08:00

scipy

Reenable pytype for numpy ufuncs.

2023-03-03 05:01:03 -08:00

state

Replace jax._src.util.prod with math.prod.

2023-02-28 12:41:00 -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

Internal change

2023-03-04 13:24:11 +00:00

ad_util.py

custom_jvp symbolic zeros support

2023-02-24 07:33:49 -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

Internal change

2023-03-04 13:24:11 +00:00

array.py

Add sharding to the signature of shard_args and update

2023-03-02 13:29:03 -08:00

basearray.py

TMP

2023-02-17 09:49:11 -08:00

basearray.pyi

Deprecate three jax.Array methods:

2023-02-23 16:15:09 -08:00

callback.py

Refactor effects system to use effect types, not objects

2023-02-17 17:40:08 -08:00

checkify.py

attach debug info to jaxpr, pass to mlir/mhlo

2023-03-02 17:23:58 -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 pytype suppression for jax/_src/config.py

2023-02-27 10:39:55 -08:00

core.py

attach debug info to jaxpr, pass to mlir/mhlo

2023-03-02 17:23:58 -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

[custom_vjp] bwd function should not be WrappedFun, may run multiple times

2023-03-01 16:23:30 -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

Replace usage of {in|out}_axis_resources with {in|out}_shardings

2023-02-28 14:29:09 -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

Fix stale reference to util.prod.

2023-03-02 08:24:30 -08:00

dispatch.py

Merge pull request #14711 from gnecula:tf_cross_platform2

2023-03-03 01:02:28 -08:00

distributed.py

Add Open MPI automatic distributed initialization

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

dlpack.py

[JAX] Move jax._src.lib.xla_bridge to jax._src.xla_bridge.

2023-02-28 07:01:57 -08:00

dtypes.py

DOC: add docs for jax.dtypes module

2023-02-14 11:18:59 -08:00

effects.py

Add JaxprInputEffect and refactor StateEffects to use it

2023-02-21 16:30:06 -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

global_device_array.py

Internal change

2023-03-04 13:24:11 +00: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

maps.py

Always default keep_unused to True if going via lower_mesh_computation

2023-03-02 22:13:24 -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

Internal change

2023-03-04 13:24:11 +00: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

Add sharding to the signature of shard_args and update

2023-03-02 13:29:03 -08:00

profiler.py

[JAX] Move jax._src.lib.xla_bridge to jax._src.xla_bridge.

2023-02-28 07:01:57 -08:00

public_test_util.py

[JAX] Move jax._src.lib.xla_bridge to jax._src.xla_bridge.

2023-02-28 07:01:57 -08:00

random.py

Merge pull request #14658 from JiaYaobo:chisq_and_f_dist

2023-03-01 06:35:34 -08:00

sharding.py

Add a helpful error message when device_putting with a Sharding that is incompatible with the shape of the input

2023-02-23 15:37:50 -08:00

source_info_util.py

Remove circular dependency between source_info_util and util.

2023-02-27 11:41:46 -08:00

stages.py

Remove use_stablehlo as minimum mlir_api_version >= 43

2023-02-24 15:20:09 -08:00

test_util.py

Shorten alias chains for names exported in jax. namespace.

2023-03-01 09:19:44 -08:00

traceback_util.py

Change definition of util.wraps so pytype can understand it.

2023-03-02 18:41:52 -08:00

tree_util.py

Internal change

2023-03-04 13:24:11 +00:00

typing.py

DOC: improve usage recommendation in jax.typing

2023-02-21 04:58:21 -08:00

util.py

Change definition of util.wraps so pytype can understand it.

2023-03-02 18:41:52 -08:00

xla_bridge.py

Shorten alias chains for names exported in jax. namespace.

2023-03-01 09:19:44 -08:00