George Necula 2228115cf4 [host_callback] Flip the JAX_HOST_CALLBACK_LEGACY flag to False ...

`jax.experimental.host_callback` has been deprecated since March 2024
 (JAX version 0.4.26). Now we set the default value of the `--jax_host_callback_legacy` configuration value to `True`, which means that if your code uses `jax.experimental.host_callback` APIs, those API calls will be implemented in terms of the new `jax.experimental.io_callback` API.

If this breaks your code, for a very limited time, you can set the `--jax_host_callback_legacy` to `True`. Soon we will remove that configuration option, so you should instead transition to using the new JAX callback APIs.

See https://github.com/google/jax/issues/20385 for a discussion.

PiperOrigin-RevId: 681004255

2024-10-01 07:07:29 -07:00

..

array_api

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

array_serialization

Rename jtu.create_global_mesh to jtu.create_mesh and use jax.make_mesh inside jtu.create_mesh to get maximum test coverage of the new API.

2024-09-03 16:23:07 -07:00

compilation_cache

Cleanup deprecated compilation cache APIs.

2024-01-12 22:44:48 -08:00

export

Re-land https://github.com/google/jax/pull/21847

2024-06-13 19:53:53 -07:00

jax2tf

Reverts eff00cc4499cfe3f3f24bafda6c1ecf908232ff3

2024-09-25 10:33:53 -07:00

key_reuse

Run pyupgrade --py310-plus.

2024-06-26 16:10:18 -04:00

mosaic

Replace disable_backends with enable_backends on jax_multiplatform_test.

2024-09-27 06:15:31 -07:00

pallas

[Pallas] Minor cleanup of memory spaces. Also add ANY as a general memory space

2024-09-19 19:08:18 -07:00

slab

Run pyupgrade --py310-plus.

2024-06-26 16:10:18 -04:00

sparse

Add support for setting a dot product "algorithm" for lax.dot_general.

2024-09-25 06:17:09 -07:00

__init__.py

add jax.experimental.primal_tangent_dtype helper

2024-09-21 20:35:20 +00:00

attrs.py

Clean up and fix primal type to tangent type mapping

2024-09-18 13:43:54 -07:00

checkify.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

compute_on.py

Initial commit for jax.experimental.compute_on API.

2024-05-17 15:59:21 -07:00

custom_partitioning.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

host_callback.py

[host_callback] Flip the JAX_HOST_CALLBACK_LEGACY flag to False

2024-10-01 07:07:29 -07:00

jet.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

layout.py

Rename layout.AUTO to DeviceLocalLayout.AUTO

2024-04-03 17:23:35 -07:00

mesh_utils.py

Create jax.make_mesh(axis_shapes: Sequence[int], axis_names: Sequence[str], devices: Sequence[jax.Device] | None = None) API to make it easier to create a mesh and reduce a ton of boilerplate.

2024-09-03 14:32:03 -07:00

multihost_utils.py

[Take 2] Generalize global jit cpp cache keys so we can add more keys than the current donate_argnums.

2024-09-17 16:11:28 -07:00

ode.py

Update jnp.clip to Array API 2023 standard

2024-04-04 22:55:10 +00:00

pjit.py

Remove a number of deprecated APIs

2023-12-18 10:08:47 -08:00

profiler.py

Add get_profiled_instructions_proto API in jax.experimental.profiler to get profiled instructions proto.

2023-06-29 21:02:05 -07:00

rnn.py

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

serialize_executable.py

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00:00

shard_alike.py

Expose shard_alike via jax.experimental. The API is x, y = shard_like(x, y).

2023-12-19 16:31:33 -08:00

shard_map.py

#sdy add JAX Shardy support for shard_map.

2024-09-26 08:45:40 -07:00

topologies.py

[PJRT C API] Supports plugins to register a method to create topology.

2024-02-22 16:55:34 -08:00

x64_context.py

MAINT Do not import the config object in JAX internals

2023-10-18 10:55:13 +01:00

xla_metadata.py

Add frontend attributes to Jax. This allows Jax users to annotate Jax code with frontend_attributes which can be traced down to the HLO level, to be used for numerical debugging purposes.

2024-09-06 16:44:56 -07:00