Peter Hawkins afb750cf76 Remove unnecessary use of xla_client.OpMetadata class. ...

We create this object and immediately turn it into a different object. We can cut out a step here!

PiperOrigin-RevId: 718023353

2025-01-21 12:24:51 -08:00

..

array_serialization

[JAX] Add a new jax_num_cpu_devices flag that allows the user to specify the number of CPU directly.

2025-01-08 06:37:44 -08:00

colocated_python

Don't monkey-patch functions in test_utils to count events for tests.

2024-12-12 09:58:14 -05:00

compilation_cache

Cleanup deprecated compilation cache APIs.

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

jax2tf

Remove unnecessary use of xla_client.OpMetadata class.

2025-01-21 12:24:51 -08:00

key_reuse

[key reuse] fix signature for device_put

2025-01-17 09:47:50 -08:00

mosaic

[mosaic_gpu] a function for bitwidth as well as bytewidht.

2025-01-21 06:52:04 -08:00

pallas

Fix the backwards pass and support more block sizes.

2025-01-21 02:48:37 +00:00

roofline

Cleanup: toward merging core.concrete_aval & xla.abstractify

2024-12-17 09:27:00 -08:00

slab

Remove ConcreteArray from JAX. It's easy to do trace-time concretization without it.

2024-10-31 14:06:54 -07:00

sparse

Rename out_type -> out_sharding parameter on einsum

2025-01-16 18:16:52 -08:00

__init__.py

add jax.experimental.primal_tangent_dtype helper

2024-09-21 20:35:20 +00:00

attrs.py

Even more linearize fixes

2025-01-13 15:30:49 -05:00

checkify.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

compute_on.py

Initial commit for jax.experimental.compute_on API.

2024-05-17 15:59:21 -07:00

custom_partitioning.py

Implement the extension to the custom_partitioning API.

2025-01-08 13:34:47 -08:00

host_callback.py

Remove remaining implementations of jax.experimental.host_callback.call.

2024-10-08 04:22:20 -07:00

jet.py

Remove internal uses of api_util.shaped_abstractify

2024-12-19 07:06:36 -08:00

layout.py

Rename layout.AUTO to DeviceLocalLayout.AUTO

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

mesh_utils.py

Create jax.make_mesh(axis_shapes: Sequence[int], axis_names: Sequence[str], devices: Sequence[jax.Device] | None = None) API to make it easier to create a mesh and reduce a ton of boilerplate.

2024-09-03 14:32:03 -07:00

multihost_utils.py

Match the behavior on single host wrt multi-host if tiled=False. Fixes https://github.com/jax-ml/jax/issues/25783

2025-01-08 13:31:35 -08:00

ode.py

Stop using generators for linear_util transformations.

2024-11-13 13:47:07 -08:00

pjit.py

Remove _pjit_lower_cached cache. We can simplify the caching of jit as we have downstream caches and a cpp cache too.

2024-11-25 11:40:50 -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

[ROCm] Implement RNN support

2025-01-14 19:04:49 -06:00

serialize_executable.py

Cleanup: fix unused imports & mark exported names

2024-10-16 17:42:41 -07: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

More linearization fixes

2025-01-15 10:27:21 -05:00

shard.py

Add reshard API in experimental. Currently for sharding_in_types we have 2 APIs: mesh_cast and reshard. Both work in sharding_in_types mode and affect the sharding of the aval. Following are the semantics of both:

2025-01-20 11:39:25 -08:00

topologies.py

Cleanup: fix unused imports & mark exported names

2024-10-16 17:42:41 -07:00

x64_context.py

MAINT Do not import the config object in JAX internals

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

xla_metadata.py

Add frontend attributes to Jax. This allows Jax users to annotate Jax code with frontend_attributes which can be traced down to the HLO level, to be used for numerical debugging purposes.

2024-09-06 16:44:56 -07:00