Jake VanderPlas fa6d3f26ff jnp.unique: make return_inverse shape match NumPy 2.0

2024-01-16 11:47:45 -08:00

..

array_api

[array api] add jax.numpy.concat

2024-01-12 13:12:09 -08:00

array_serialization

Upgrade remaining sources to Python 3.9

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

compilation_cache

Cleanup deprecated compilation cache APIs.

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

export

[shape_poly] Improve core.max_dim and core.min_dim

2024-01-15 15:10:28 +02:00

jax2tf

Deprecate support for the mhlo dialect.

2024-01-15 02:13:40 -08:00

key_reuse

Upgrade remaining sources to Python 3.9

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

mosaic

[Mosaic] Remove Python implementation of apply_vector_layout and infer_memref_layout.

2024-01-10 13:00:21 -08:00

pallas

[splash_attention] Fix empty mask corner cases in mask_info construction

2024-01-15 14:41:32 -08:00

sparse

jnp.unique: make return_inverse shape match NumPy 2.0

2024-01-16 11:47:45 -08: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

Cache ClosedJaxpr creation to minimize cache_misses. ClosedJaxpr should always be created under a cache.

2023-12-21 22:15:52 -08:00

host_callback.py

revise logic for tangent types of extended dtypes

2023-12-20 14:24:52 -08:00

jet.py

partial rollback of #19096 due to internal breakage (relying on jax internals)

2023-12-22 15:54:26 -08: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

Updates doc for host_local_array_to_global_array to reflect a few use case patterns and adds a few extra tests.

2024-01-08 16:51:26 -08: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_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

[shard-map] add approx_top_k replication rule

2024-01-16 03:59:40 -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