jax authors a9460f2461 Merge pull request #21130 from Micky774:reshape ...

PiperOrigin-RevId: 632277406

2024-05-09 14:51:56 -07:00

..

array_api

Deprecate newshape argument of jnp.reshape

2024-05-09 21:02:07 +00:00

array_serialization

Move dtype settings out of metadata field into the root of Tensorstore spec

2024-05-01 14:48:55 -07:00

compilation_cache

Cleanup deprecated compilation cache APIs.

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

export

[export] Add backwards compatibility test for Pallas call on GPUs.

2024-05-02 05:38:33 -07:00

jax2tf

Fix mlir.merge_mlir_modules to properly remember the inlined symbols

2024-04-29 19:56:41 +03:00

key_reuse

[key reuse] call key reuse logic directly in dispatch

2024-04-11 17:08:32 -07:00

mosaic

[jax:mosaic-gpu] FragmentedArray can do tiled load.

2024-05-07 18:13:55 -07:00

pallas

Removed get_compute_capability from jax.experimental.pallas.gpu

2024-05-08 21:10:43 +01:00

sparse

Merge pull request #21050 from rajasekharporeddy:test_branch3

2024-05-03 03:24:22 -07:00

__init__.py

start adding EArray, a jax.Array analog that can contain extended dtypes

2024-04-06 13:09:25 -07:00

attrs.py

[attrs] fix tracer lifetime bug, fixes #20082

2024-03-05 12:08:44 -08:00

checkify.py

Document why 'import name as name' is used

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

custom_partitioning.py

Cleanup second registration of custom_partitioning callbacks now that

2024-05-08 10:45:39 -07:00

host_callback.py

Add instructions to report issues to GitHub for unsupported features in PJRT C API.

2024-04-30 17:15:40 -07:00

jet.py

Add Layout support to jax.jit.

2024-04-05 20:09:34 -07:00

layout.py

Rename layout.AUTO to DeviceLocalLayout.AUTO

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

maps.py

Register maps module deprecation outside of module

2024-03-19 09:21:48 -07:00

mesh_utils.py

Update the arg description.

2024-03-22 20:10:18 -07:00

multihost_utils.py

Deprecate jax.tree_map in favor of jax.tree.map

2024-02-22 11:35:39 -08: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

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

Allow nested shard_map.

2024-05-09 14: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