rocm_jax

clusters

Automate arguments for jax.distributed.initialize for cloud TPU environments.

2023-11-30 22:25:00 -08:00

debugger

jax.debug.breakpoint now supports being DCE'd. Drive-by: fix #16186

2023-07-31 15:36:43 -07:00

extend

[jex] replace extend.random.PRNGImpl with extend.random.define_prng_impl

2023-10-31 17:21:54 -07:00

image

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04:00

internal_test_util

Use trace_to_jaxpr_dynamic for the apply_primitive path. trace_to_jaxpr_final is only for final style primitives. Also do some cleanup.

2023-11-28 14:35:44 -08:00

interpreters

Deprecate the device() method of JAX arrays

2023-11-30 11:43:02 -08:00

lax

Always lower reduce_scatter_p as an HLO ReduceScatter.

2023-11-29 05:37:58 -08:00

lib

Increase minimum jaxlib version to 0.4.19.

2023-11-17 09:38:31 -08:00

nn

Add nn.squareplus.

2023-11-14 23:52:41 -05:00

numpy

Merge pull request #18708 from jakevdp:array-equal-dep

2023-11-29 08:58:29 -08:00

ops

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

pallas

[Pallas/Mosaic] Add support for barrier semaphores

2023-11-29 04:04:11 -08:00

scipy

Fix sign of jax.scipy.special.gamma for negative inputs

2023-11-27 14:08:02 -08:00

state

Migrate a subset of internal modules to use state objects

2023-10-09 07:29:53 -07:00

third_party

Add a type stub for jax.numpy.

2023-08-22 11:50:49 -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

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04:00

ad_checkpoint.py

Migrate a subset of internal modules to use state objects

2023-10-09 07:29:53 -07:00

ad_util.py

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04:00

api_util.py

jax.vmap: improve docs & error for structured in_axes

2023-11-15 11:56:53 -08:00

api.py

Merge pull request #18736 from mattjj:device-put-fixes

2023-11-29 16:51:15 -08:00

array.py

Deprecate the device() method of JAX arrays

2023-11-30 11:43:02 -08:00

basearray.py

Ensure sharding-related array properties are documented

2023-11-03 09:56:33 -07:00

basearray.pyi

Ensure sharding-related array properties are documented

2023-11-03 09:56:33 -07:00

cache_key.py

Increase minimum jaxlib version to 0.4.19.

2023-11-17 09:38:31 -08:00

callback.py

DOC: mention ShapeDtypeStruct in callback docstrings

2023-10-11 19:55:39 -07:00

checkify.py

MAINT Do not use str() and repr() in f-string replacement fields

2023-10-23 15:12:04 +01:00

cloud_tpu_init.py

Remove calling configure_library_path during jax import and get libtpu path from libtpu_module.get_library_path().

2023-10-10 10:59:37 -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

Downgrade a bunch of logging to DEBUG

2023-11-29 12:10:53 +00:00

compiler.py

Downgrade a bunch of logging to DEBUG

2023-11-29 12:10:53 +00:00

config.py

Default-enable the Jax persistent compilation cache.

2023-11-27 14:53:20 -08:00

core.py

Replace apply_primitive internals with jax.jit.

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

Fix type failures under an upcoming pytype change.

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

custom_derivatives.py

Migrate a subset of internal modules to use state objects

2023-10-09 07:29:53 -07:00

custom_transpose.py

Use lower-case PEP 585 names for types.

2023-06-23 15:12:14 -07:00

debugging.py

small tweaks to jax.debug.print docstring

2023-11-22 15:03:41 -08:00

deprecations.py

Make deprecation warnings warn at appropriate stacklevel

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

dispatch.py

Replace apply_primitive internals with jax.jit.

2023-11-29 18:07:13 -08:00

distributed.py

Increase minimum jaxlib version to 0.4.19.

2023-11-17 09:38:31 -08:00

dlpack.py

Increase minimum jaxlib version to 0.4.19.

2023-11-17 09:38:31 -08:00

dtypes.py

Finish deprecation of allow_opaque_dtype

2023-11-02 09:51:06 -07:00

effects.py

[callbacks] Add support for shardable ordered effects.

2023-09-18 02:50:25 -07: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

Fix typos across the package

2023-09-22 14:54:31 -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

jaxpr_util.py

Switch from flake8 to Ruff

2023-11-15 22:35:52 -05:00

lax_reference.py

add an exp2 primitive and lax.exp2

2023-07-28 12:33:49 -07:00

layout.py

Make the SpecifiedLayout class opaque.

2023-11-18 15:17:16 -08:00

lazy_loader.py

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04:00

linear_util.py

Migrate a subset of internal modules to use state objects

2023-10-09 07:29:53 -07:00

logging_config.py

Add jax_debug_log_modules config option.

2023-07-28 18:11:12 +00:00

maps.py

Replace apply_primitive internals with jax.jit.

2023-11-29 18:07:13 -08:00

mesh.py

Allow threads to race setting attributes on Mesh.

2023-11-22 05:47:56 -08:00

monitoring.py

Allow duration event listeners to take extra keyword arguments.

2023-10-04 13:15:49 -07:00

op_shardings.py

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04:00

partition_spec.py

Improve the sharding documentation.

2023-08-03 11:12:02 -04: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

Replace apply_primitive internals with jax.jit.

2023-11-29 18:07:13 -08:00

pretty_printer.py

Update flags to use the ABSL typed flag API.

2023-07-27 12:15:58 -07:00

prng.py

Use MLIR generated convenience functions athing(...) instead of writing AThingOp(...).result.

2023-11-17 11:47:14 -08:00

profiler.py

Increase minimum jaxlib version to 0.4.19.

2023-11-17 09:38:31 -08:00

public_test_util.py

Properly pack and unpack int4 arrays on CPU in PJRT.

2023-11-01 17:39:24 -07:00

random.py

Merge pull request #18563 from jakevdp:cleanup-fold-in

2023-11-21 13:37:51 -08:00

sharding_impls.py

Increase minimum jaxlib version to 0.4.19.

2023-11-17 09:38:31 -08:00

sharding_specs.py

Replace references to DeviceArray with Array.

2023-08-18 17:46:00 -04:00

sharding.py

Make is_fully_addressable an abstract method and implement it on each concrete Sharding.

2023-09-05 17:28:22 -07:00

source_info_util.py

MAINT Do not use str() and repr() in f-string replacement fields

2023-10-23 15:12:04 +01:00

stages.py

Fix indexing bug when querying _input_layouts

2023-11-27 21:12:07 -08:00

test_util.py

tests: improve warnings-related tests

2023-11-30 10:35:24 -08:00

tpu_custom_call.py

[Mosaic] Use C++ apply-vector-layout as default setting

2023-11-27 08:39:42 -08:00

traceback_util.py

Migrate a subset of internal modules to use state objects

2023-10-09 07:29:53 -07:00

tree_util.py

Switch from flake8 to Ruff

2023-11-15 22:35:52 -05:00

typing.py

[shape_poly] Simplify the API for processing polymorphic_shape specifications

2023-11-28 12:45:59 +02:00

util.py

Fix hashcode for HashablePartial to get equal hashes for equal objects

2023-11-06 17:27:19 +08:00

xla_bridge.py

Use process_count(backend) in local_devices().

2023-11-29 14:24:06 -08:00