Parker Schuh 7ba8622719 For custom_partitioning, directly emit call when inside of a shard_map. ...

PiperOrigin-RevId: 592011427

2023-12-18 14:32:38 -08:00

..

array_api

jnp.linalg: add matrix_norm, matrix_transpose, vector_norm, vector_transpose

2023-12-15 14:17:36 -08:00

array_serialization

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00: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

[export] Fix handling of float0 when exporting

2023-12-18 14:51:27 +02:00

jax2tf

Move export backwards compatibility tests out of jax2tf. Step 2.

2023-12-18 10:16:54 -08:00

key_reuse

Upgrade remaining sources to Python 3.9

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

mosaic

Increase minimum jaxlib version to 0.4.19.

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

pallas

Upgrade remaining sources to Python 3.9

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

sparse

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00: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

For custom_partitioning, directly emit call when inside of a shard_map.

2023-12-18 14:32:38 -08:00

host_callback.py

Upgrade remaining sources to Python 3.9

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

jet.py

Use lower-case PEP 585 names for types.

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

layout.py

Remove DefaultLayout and make None same as DefaultLayout

2023-11-16 18:01:27 -08:00

maps.py

Copybara import of the project:

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

mesh_utils.py

Upgrade remaining sources to Python 3.9

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

multihost_utils.py

Upgrade remaining sources to Python 3.9

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

ode.py

Remove leading underscores in jax._src.numpy.util

2023-03-13 12:18:36 -07: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

Increase minimum jaxlib version to 0.4.19.

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

serialize_executable.py

Upgrade remaining sources to Python 3.9

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

shard_map.py

Merge pull request #18961 from mattjj:issue18955

2023-12-13 11:04:24 -08:00

topologies.py

Upgrade remaining sources to Python 3.9

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

x64_context.py

MAINT Do not import the config object in JAX internals

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