Yash Katariya 6007698f4e Allow None to be passed to in_shardings and out_shardings. The default is still UNSPECIFIED to handle edge cases around the old semantics where None is treated as fully replicated. ...

The semantics are as follow:

* if the mesh context manager is not provided, None will be treated as UNSPECIFIED for both in_shardings and out_shardings

* If the mesh context manager is provided, None will be treated as fully replicated as per the old semantics.

This will make sure that we don't break existing code depending on None meaning replicated but also start making the transition to None meaning UNSPECIFIED for jit and pjit.

PiperOrigin-RevId: 540705660

2023-06-15 15:22:22 -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

Raise a good error message when a ShapeDtypeStruct is closed over as a const which is not a valid arg during execution.

2023-06-14 09:40:37 -07:00

lax

Merge pull request #16392 from cloudhan:fix-multi-line

2023-06-15 14:38:25 -07:00

lib

cleanup old jaxlib version check

2023-06-06 21:51:58 -07:00

nn

jax.nn.softmax: fix fill value when where is specified

2023-06-01 10:18:05 -07:00

numpy

Merge pull request #16265 from jakevdp:zeros-like

2023-06-06 13:01:54 -07:00

ops

Remove some dead code and unused imports

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

scipy

Fix test failures in JAX under NumPy 1.25.0rc1.

2023-06-09 14:10:35 -07:00

state

Verify that slices are trivial before discarding them in state primitives

2023-05-04 05:59:47 -07:00

third_party

Change uses of np.product to np.prod

2023-06-02 03:57:30 -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

Add some preliminary support for int4/uint4 types to JAX.

2023-05-18 14:27:33 -07:00

ad_checkpoint.py

In jax.remat/jax.checkpoint, don't cache on Tracers in static args

2023-05-04 13:42:00 -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

Allow None to be passed to in_shardings and out_shardings. The default is still UNSPECIFIED to handle edge cases around the old semantics where None is treated as fully replicated.

2023-06-15 15:22:22 -07:00

array.py

Resolve unsoundness caught by pytype --strict-none-binding

2023-06-06 06:52:21 -07:00

basearray.py

Split basearray into separate Bazel module.

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

basearray.pyi

Add jnp.matrix_transpose() and jax.Array.mT

2023-05-25 09:02:05 -07:00

callback.py

Add maximal sharding for pure_callback not inside of a shard_map

2023-05-11 13:28:37 -07:00

checkify.py

Checkify: release memory by deleting references to variables after last use.

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

Bump minimum jaxlib version to 0.4.11. xla_extension_version is 158 and mlir_api_version is 49. It will subsume https://github.com/google/jax/pull/16161#issuecomment-1564977332

2023-06-01 09:42:55 -07:00

config.py

disable custom_jvp for softmax by default

2023-05-23 11:56:50 -07:00

core.py

Update todo comments based on offline discussion.

2023-06-13 10:44:52 -04: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

functools.partial should not unpack curried kwargs in custom_jvp.

2023-06-06 13:23:13 -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

Make inspect_sharding_lowering_rule work only with HloSharding in it's callback. Also remove get_replicated_op_sharding since it is not needed anymore.

2023-06-07 14:39:52 -07:00

deprecations.py

Make deprecation warnings warn at appropriate stacklevel

2023-05-23 14:43:38 -07:00

dispatch.py

Add missing %s parameter to error message in _cache_write.

2023-06-09 15:12:38 -07:00

distributed.py

[PJRT PLUGIN] Add num_processes to distributed.global_state.

2023-05-26 13:14:40 -07:00

dlpack.py

[JAX] Delete _DeviceArray and DeviceArray.

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

dtypes.py

Add support for dtype float8_e4m3b11fnuz

2023-06-06 00:50:57 -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

jax errors: improve leading info in error docs

2023-06-14 02:07:00 -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

Fix FileExistsError for tests/gfile_cache_test.py::FileSystemCacheTest::test_threads

2023-06-14 08:07:22 +08: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

Fix the lu.clear_all_cache function by adding the memoized_fun to the global weakref set rather than the function local fun_caches weakrefDict.

2023-05-24 13:58:51 -07:00

maps.py

Use _to_xla_hlo_sharding everywhere in JAX. Remove _to_xla_op_sharding in favor of _to_xla_hlo_sharding since constructing a C++ class is faster than protos and will help with further changes coming to HloSharding.

2023-06-05 13:41:31 -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

Remove xc._xla.is_op_sharding_fully_replicated since JAX should now produce HloSharding's everywhere.

2023-06-06 07:44:14 -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

pickle_util.py

Use cloudpickle for serializing host callbacks

2023-06-09 13:34:43 -07:00

pjit.py

Allow None to be passed to in_shardings and out_shardings. The default is still UNSPECIFIED to handle edge cases around the old semantics where None is treated as fully replicated.

2023-06-15 15:22:22 -07:00

pretty_printer.py

Resolve unsoundness caught by pytype --strict-none-binding

2023-06-06 06:52:21 -07:00

prng.py

Simplify sharding types input to physical_hlo_sharding and lower_jaxpr_to_fun.

2023-06-06 18:03:12 -07:00

profiler.py

Split _src/profiler into a separate BUILD target.

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

public_test_util.py

Fix test_util for new float8 type

2023-06-08 00:30:45 -07:00

random.py

outline jitted jax.random functions

2023-05-25 15:01:04 -07:00

sharding_impls.py

Allow None to be passed to in_shardings and out_shardings. The default is still UNSPECIFIED to handle edge cases around the old semantics where None is treated as fully replicated.

2023-06-15 15:22:22 -07:00

sharding_specs.py

Convert OpSharding to HloSharding in the constructor of GSPMDSharding. Also make op_sharding_to_indices to work with only HloSharding.

2023-06-06 06:35:29 -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

Add option jax_include_full_tracebacks_in_locations.

2023-05-22 07:41:29 -07:00

stages.py

Small cleanups to pxla.py.

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

test_util.py

[shape_poly] Improve compile-time shape checking.

2023-06-13 20:47:30 +03: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

Add jax.tree_util.tree_leaves_with_path(tree).

2023-06-12 04:13:37 -07:00

typing.py

Allow duck-typed inputs to zeros_like, ones_like, etc.

2023-06-06 00:59:51 -07:00

util.py

Fix the lu.clear_all_cache function by adding the memoized_fun to the global weakref set rather than the function local fun_caches weakrefDict.

2023-05-24 13:58:51 -07:00

xla_bridge.py

Remove older plugin device integration.

2023-06-14 15:26:58 -04:00