rocm_jax

Use pytype_strict_library() in Bazel build rules.

2023-03-27 10:16:08 -07:00

set cli debugger use_rawinput to True if optional stdin is not provided. This enables the use of arrow keys in the jdb prompt.

2023-02-27 01:57:38 -05:00

image

Typo fix in ResizeMethod docstring, scale.py

2023-03-20 23:10:51 +01:00

internal_test_util

lazy_loader_module: Move to new internal_test_util directory.

2023-03-10 10:29:33 -08:00

interpreters

Remove global_arg_shapes from pmap since it was only used for sharded_jit and sharded_jit was removed from JAX a long time ago

2023-03-29 09:23:22 -07:00

lax

Remove global_arg_shapes from pmap since it was only used for sharded_jit and sharded_jit was removed from JAX a long time ago

2023-03-29 09:23:22 -07:00

lib

Use pytype_strict_library() in Bazel build rules.

2023-03-27 10:16:08 -07:00

nn

update relu6 grad at 0 and 6 to match pytorch convention

2023-03-07 17:30:17 -08:00

numpy

Silence some pytype errors.

2023-03-28 15:24:48 -07:00

ops

Remove leading underscores in jax._src.numpy.util

2023-03-13 12:18:36 -07:00

scipy

jax.scipy.linalg.expm: support batched inputs

2023-03-27 16:39:48 -07:00

state

Move jax.interpreters.partial_eval to jax._src.interpreters.partial_eval.

2023-03-27 13:30:47 -07:00

third_party

Remove leading underscores in jax._src.numpy.util

2023-03-13 12:18:36 -07: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

Move jax.interpreters.partial_eval to jax._src.interpreters.partial_eval.

2023-03-27 13:30:47 -07:00

ad_util.py

custom_jvp symbolic zeros support

2023-02-24 07:33:49 -08:00

api_util.py

make mlir arg and result names work with pmap

2023-03-20 13:51:36 -07:00

api.py

Remove global_arg_shapes from pmap since it was only used for sharded_jit and sharded_jit was removed from JAX a long time ago

2023-03-29 09:23:22 -07:00

array.py

Bump minimum jaxlib version from 0.4.6 to 0.4.7.

2023-03-28 13:43:01 -07:00

basearray.py

Split basearray into separate Bazel module.

2023-03-13 11:14:41 -07:00

basearray.pyi

Relax type of the values argument to .at[...].set(...) and friends.

2023-03-17 10:46:36 -07:00

callback.py

Refactor effects system to use effect types, not objects

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

checkify.py

Move jax.interpreters.partial_eval to jax._src.interpreters.partial_eval.

2023-03-27 13:30:47 -07:00

cloud_tpu_init.py

Turn on PJRT C API by default.

2023-03-28 15:28:13 -07:00

config.py

Bump minimum jaxlib version from 0.4.6 to 0.4.7.

2023-03-28 13:43:01 -07:00

core.py

Split core.py and several files in an SCC with it into a separate Bazel build target.

2023-03-28 18:31:13 -07: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

Move jax.interpreters.partial_eval to jax._src.interpreters.partial_eval.

2023-03-27 13:30:47 -07:00

custom_derivatives.py

Move jax.interpreters.partial_eval to jax._src.interpreters.partial_eval.

2023-03-27 13:30:47 -07:00

custom_transpose.py

Move jax.interpreters.partial_eval to jax._src.interpreters.partial_eval.

2023-03-27 13:30:47 -07:00

debugging.py

Move jax.interpreters.partial_eval to jax._src.interpreters.partial_eval.

2023-03-27 13:30:47 -07: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

Split core.py and several files in an SCC with it into a separate Bazel build target.

2023-03-28 18:31:13 -07:00

distributed.py

Set coordinator address to allow it to later be used to initialize OCDBT coordinator server. Allow user to pass ts.Context when serializing or deserializing.

2023-03-28 10:25:07 -07:00

dlpack.py

Bump minimum jaxlib version from 0.4.6 to 0.4.7.

2023-03-28 13:43:01 -07:00

dtypes.py

[XLA:Python] Change JAX and the XLA Python extension to get NumPy bfloat16/float8 types from ml_dtypes.

2023-03-23 05:13:39 -07:00

effects.py

Add JaxprInputEffect and refactor StateEffects to use it

2023-02-21 16:30:06 -08:00

environment_info.py

Split _src files custom_api_util, deprecations, effects and environment_info into separate Bazel targets.

2023-03-10 12:26:05 -08:00

errors.py

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

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

flatten_util.py

use Partial to make ravel_pytree unflatteners jit-friendly

2023-03-13 11:06:56 -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

Add padding option "SAME_LOWER" for ticket https://github.com/google/jax/pull/14990

2023-03-23 15:50:16 -07:00

lazy_loader.py

Split _src modules cloud_tpu_init, lazy_loader, path, monitoring into their own pytype_library Bazel targets.

2023-03-09 13:11:04 -08:00

linear_util.py

Split core.py and several files in an SCC with it into a separate Bazel build target.

2023-03-28 18:31:13 -07:00

maps.py

Move jax.interpreters.partial_eval to jax._src.interpreters.partial_eval.

2023-03-27 13:30:47 -07:00

mesh.py

Split Mesh and ResourceEnv into a new module jax._src.mesh.

2023-03-10 10:08:21 -08:00

monitoring.py

Split _src modules cloud_tpu_init, lazy_loader, path, monitoring into their own pytype_library Bazel targets.

2023-03-09 13:11:04 -08:00

path.py

Make the etils[epath] dependency of JAX optional.

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

pjit.py

Deprecate FROM_GDA and remove its support from pjit's code since jax.Array inside pjit has sharding inference capabilities by default.

2023-03-28 10:33:00 -07:00

pretty_printer.py

Split _src/pretty_printer.py into its own Bazel target.

2023-03-09 08:51:30 -08:00

prng.py

internal: refactor array methods into separate private submodule

2023-03-23 10:57:53 -07:00

profiler.py

Split _src/profiler into a separate BUILD target.

2023-03-10 11:38:46 -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

Bump minimum jaxlib version from 0.4.6 to 0.4.7.

2023-03-28 13:43:01 -07:00

sharding_impls.py

Remove MeshPspecSharding since it has been more than 3 months since it was deprecated (Nov 2, 2022). The replacement name is NamedSharding.

2023-03-28 10:47:42 -07:00

sharding.py

Improve pytype inference for Sharding type.

2023-03-14 14:20:17 -07:00

source_info_util.py

Split source_info_util into its own Bazel target.

2023-03-10 08:41:06 -08:00

stages.py

Remove references to jax.config.jax_array, which is always True at head.

2023-03-15 17:09:11 -07:00

test_util.py

[Jax cleanup]

2023-03-16 15:47:28 -07:00

traceback_util.py

Split _src/traceback_util.py into its own Bazel target.

2023-03-09 10:33:47 -08:00

tree_util.py

improve scan error messages

2023-03-23 14:53:05 -07:00

typing.py

Split basearray into separate Bazel module.

2023-03-13 11:14:41 -07:00

util.py

make mlir arg and result names work with static_argnums/argnames

2023-03-20 11:50:30 -07:00

xla_bridge.py

Bump minimum jaxlib version from 0.4.6 to 0.4.7.

2023-03-28 13:43:01 -07:00