George Necula b65c1b293b [jax2tf] First step to enable multi-platform native lowering ...

Enable experiments with jax2tf native serialization for
multiple platforms. This feature is not yet fully functional
but we need this change to enable further testing.

Cleanup some of the places that are specific to single-platform
serialization, e.g., `lowering_platform`, and generalize
them to multiple platforms (`lowering_platforms`).

2023-10-16 07:01:23 -07:00

..

array_serialization

Fix typos across the package

2023-09-22 14:54:31 -07:00

compilation_cache

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

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

export

[jax2tf] First step to enable multi-platform native lowering

2023-10-16 07:01:23 -07:00

jax2tf

[jax2tf] First step to enable multi-platform native lowering

2023-10-16 07:01:23 -07:00

mosaic

Add Mosaic to Jaxlib and expose bindings in jax.experimental.mosaic

2023-07-20 18:28:51 -07:00

pallas

Fix mypy errors in flash_attention

2023-10-10 13:05:48 +00:00

sparse

MAINT Migrate remaining internal/test modules to use state objects

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

__init__.py

Copybara import of the project:

2023-05-24 10:35:37 -07:00

checkify.py

Document why 'import name as name' is used

2022-12-14 15:07:04 -08:00

custom_partitioning.py

Pad generated sharding specs with None up to ndims to simplify comparing dims

2023-08-10 17:02:31 -07:00

host_callback.py

fix typo: device_fun(c) -> device_fun(x)

2023-10-02 23:38:48 -07:00

jet.py

Use lower-case PEP 585 names for types.

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

maps.py

Copybara import of the project:

2023-05-24 10:35:37 -07:00

mesh_utils.py

Clarify that mesh_utils.create_device_mesh's contiguous_submeshes arg isn't necessary with jax.Array

2023-10-04 11:24:23 -07:00

multihost_utils.py

Add docs for is_fully_addressable to jax.Array and remove GDA from process_allgather docs and clarify it a bit more.

2023-08-20 18:58:22 -07:00

ode.py

Remove leading underscores in jax._src.numpy.util

2023-03-13 12:18:36 -07:00

pjit.py

Deprecate jax.experimental.pjit.with_sharding_constraint. Replacement is jax.lax.with_sharding_constraint which has been available since 1 year.

2023-09-14 09:23:03 -07: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

Fix typos across the package

2023-09-22 14:54:31 -07:00

serialize_executable.py

Stop passing CompileOptions when deserializing.

2023-05-10 16:22:54 -07:00

shard_map.py

Add support for (ordered) io_callable within shard_map

2023-09-25 13:56:43 -07:00

topologies.py

Apply pyupgrade --py39-plus.

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

x64_context.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00