jax authors 11c4e2c820 [JAX] Add an option subset_by_index that allows computing a contiguous subset of eigenvalues from eigh. ...

PiperOrigin-RevId: 577222219

2023-10-27 09:29:33 -07:00

..

array_serialization

Add mitigation techniques in the error message when a barrier timeout occurs.

2023-10-27 08:57:39 -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

Cleanup the handling of single- and multi-platform lowering in ModuleContext

2023-10-25 10:40:41 -07:00

jax2tf

[JAX] Add an option subset_by_index that allows computing a contiguous subset of eigenvalues from eigh.

2023-10-27 09:29:33 -07:00

mosaic

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

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

pallas

[Pallas] Add Slice to public interface

2023-10-26 05:50:59 -07: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

Cleanup the handling of single- and multi-platform lowering in ModuleContext

2023-10-25 10:40:41 -07:00

host_callback.py

Cleanup the handling of single- and multi-platform lowering in ModuleContext

2023-10-25 10:40:41 -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

factor out subs_list and subs_list2

2023-10-19 09:32:44 -07:00

topologies.py

Apply pyupgrade --py39-plus.

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

x64_context.py

MAINT Do not import the config object in JAX internals

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