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

Change to simpler import for jax.config

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

interpreters

[shape_poly] DynamicRngBitGenerator: add shape poly support with native serialization.

2023-06-19 00:39:37 -07:00

lax

[shape_poly] DynamicRngBitGenerator: add shape poly support with native serialization.

2023-06-19 00:39:37 -07:00

lib

Remove redundant stablehlo import.

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

Deduplicate references closed over across branches of a lax.cond.

2023-06-15 23:58:14 -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

Add live-analysis memory optimization to more jaxpr interpreters.

2023-06-16 06:08:51 -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

Errors: avoid printing tracer repr for concretization errors

2023-06-20 00:33:51 -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

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

Errors: avoid printing tracer repr for concretization errors

2023-06-20 00:33:51 -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

Fix the docs build

2023-06-16 13:14:38 -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

Generate topology from external TPU names and topology parameters.

2023-06-17 12:35:11 -07:00