Yash Katariya 34d5a6259f Default jax_spmd_mode to allow_jit which will allow explicit jax.jit to not raise the multihost error (since jit and pjit have been merged). ...

Implicit jit and apply_primitive will still raise an error though (which is recognized via inline parameter). Majority of jnp operations in JAX should be inlined.

PiperOrigin-RevId: 527398394

2023-04-26 15:56:46 -07:00

..

clusters

Use pytype_strict_library() in Bazel build rules.

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

debugger

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

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

interpreters

Default jax_spmd_mode to allow_jit which will allow explicit jax.jit to not raise the multihost error (since jit and pjit have been merged).

2023-04-26 15:56:46 -07:00

lax

KeyArray: implement __eq__ and __ne__

2023-04-26 13:12:24 -07:00

lib

[JAX] Improve handling of metadata in compilation cache.

2023-04-19 13:27:04 -07:00

nn

add custom_jvp for jax.nn.softmax

2023-04-22 11:28:03 -07:00

numpy

PRNGKeyArray: allow opaque dtype passed as dtype arg

2023-04-25 12:27:12 -07:00

ops

Remove some dead code and unused imports

2023-04-12 12:15:15 -07:00

scipy

grad(entr)(0.0): return inf instead of NaN

2023-04-25 08:32:37 -07:00

state

Avoid imports from the public jax.* namespace in more places internally.

2023-04-04 11:41:40 -07:00

third_party

Fix typo in betaln.py

2023-03-28 14:44:43 +09: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

Avoid imports from the public jax.* namespace in more places internally.

2023-04-04 11:41:40 -07:00

ad_util.py

Split more targets out the main JAX Bazel target.

2023-03-30 06:12:45 -07:00

api_util.py

Add arg_names to aval mismatch error raised during AOT compilation to raise better error messages

2023-04-19 15:08:53 -07:00

api.py

PRNGKeyArray: fix compatibility with eval_shape

2023-04-25 10:46:39 -07:00

array.py

Don't explode into individual shards if an Array is fully replicated and addressable_data is called.

2023-04-17 10:05:37 -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

DOC: add docstrings for callback functions

2023-04-12 07:33:09 -07:00

checkify.py

Merge pull request #15639 from jakevdp:checkify-dynamic

2023-04-18 08:28:33 -07:00

cloud_tpu_init.py

Turn on PJRT C API by default.

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

compilation_cache_interface.py

Move implementation of compilation cache out of jax/experimental and into jax/_src.

2023-04-17 08:35:53 -07:00

compilation_cache.py

Include the device_kind in the compilation cache key.

2023-04-20 06:16:45 -07:00

config.py

Default jax_spmd_mode to allow_jit which will allow explicit jax.jit to not raise the multihost error (since jit and pjit have been merged).

2023-04-26 15:56:46 -07:00

core.py

[djax] handle simple reshapes and size-0 checks

2023-04-21 19:20:48 -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

Fix type failures under an upcoming pytype change.

2023-04-07 07:10:09 -07:00

custom_derivatives.py

Copybara import of the project:

2023-04-12 15:11:55 -07:00

custom_transpose.py

Avoid imports from the public jax.* namespace in more places internally.

2023-04-04 11:41:40 -07:00

debugging.py

Convert inspect_sharding to register the handler directly in c++ so that it can

2023-04-26 11:22:28 -07:00

deprecations.py

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

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

dispatch.py

Default jax_spmd_mode to allow_jit which will allow explicit jax.jit to not raise the multihost error (since jit and pjit have been merged).

2023-04-26 15:56:46 -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

[JAX] Delete _DeviceArray and DeviceArray.

2023-03-29 15:07:14 -07:00

dtypes.py

Merge pull request #15746 from jakevdp:dtype-opaque

2023-04-25 14:40:32 -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

Add test framework for module attribute

2023-04-21 13:20:16 -07:00

flatten_util.py

Avoid imports from the public jax.* namespace in more places internally.

2023-04-04 11:41:40 -07:00

gfile_cache.py

Move implementation of compilation cache out of jax/experimental and into jax/_src.

2023-04-17 08:35:53 -07:00

iree.py

Remove PyBuffer type and its bindings.

2023-04-04 14:24:23 -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

Merge pull request #15221 from froystig:custom-vjp-symbolic-zeros2

2023-04-08 09:49:45 -07:00

maps.py

Default jax_spmd_mode to allow_jit which will allow explicit jax.jit to not raise the multihost error (since jit and pjit have been merged).

2023-04-26 15:56:46 -07:00

mesh.py

Change np.prod->math.prod

2023-04-13 11:48:11 -07: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

op_shardings.py

Small optimization: Add a C++ helper for testing if an OpSharding is fully replicated.

2023-04-13 06:13:35 -07:00

partition_spec.py

Move PartitionSpec into its own file (jax/_src/partition_spec.py).

2023-04-06 11:43:25 -07:00

path.py

Make the etils[epath] dependency of JAX optional.

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

pjit.py

Default jax_spmd_mode to allow_jit which will allow explicit jax.jit to not raise the multihost error (since jit and pjit have been merged).

2023-04-26 15:56:46 -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

KeyArray: implement __eq__ and __ne__

2023-04-26 13:12:24 -07:00

profiler.py

Split _src/profiler into a separate BUILD target.

2023-03-10 11:38:46 -08:00

public_test_util.py

Avoid imports from the public jax.* namespace in more places internally.

2023-04-04 11:41:40 -07:00

random.py

fix pareto docstring x domain

2023-04-22 11:36:16 +08:00

sharding_impls.py

Fix a bug in converting GSPMDSharding to PositionalSharding. Also assert that we are creating correct OpShardings (like the check in hlo_sharding.cc).

2023-04-14 15:51:58 -07:00

sharding_specs.py

Fix a bug in converting GSPMDSharding to PositionalSharding. Also assert that we are creating correct OpShardings (like the check in hlo_sharding.cc).

2023-04-14 15:51:58 -07:00

sharding.py

Add is_fully_replicated method to Shardings. This allows to scrub the usage of is_op_sharding_replicated from JAX because we can just query it on Shardings and save an expensive round trip to OpSharding creation.

2023-04-14 13:56:33 -07:00

source_info_util.py

Micro-optimization: avoid _replace on sourceinfo namedtuples.

2023-04-10 19:28:45 -07:00

stages.py

Small cleanups to pxla.py.

2023-04-21 14:38:27 -07:00

test_util.py

KeyArray: use assertArraysEqual in place of assertKeysEqual

2023-04-26 13:15:03 -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

[jax2tfO] Add support for calling an exported JAX function

2023-04-25 14:14:55 +02:00

typing.py

Split basearray into separate Bazel module.

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

util.py

jax.errors: set obj.__module__ to public module

2023-04-21 09:31:02 -07:00

xla_bridge.py

Add an experimental API that allows compiling AOT for TPUs.

2023-04-19 13:33:59 -07:00