Peter Hawkins 64eae324ee Migrate JAX MLIR Python dialect extensions to nanobind. ...

Now that https://github.com/llvm/llvm-project/pull/117922 has landed upstream, we can work towards removing our last uses of pybind11.

PiperOrigin-RevId: 705872751

2024-12-13 07:08:28 -08:00

..

mosaic

Cleanup: remove uses of no-op raise_to_shaped

2024-12-12 09:49:06 -08:00

mosaic_gpu

jax.core: deprecate a number of APIs

2024-12-10 11:11:32 -08:00

triton

Migrate JAX MLIR Python dialect extensions to nanobind.

2024-12-13 07:08:28 -08:00

__init__.py

[Pallas] Upstream pallas to JAX

2023-08-01 16:43:13 -07:00

BUILD

[Pallas] Add a cost estimator for Pallas/JAX functions.

2024-11-11 11:13:58 -08:00

core.py

[sharding_in_types] Add slice_p and squeeze_p sharding rule to make flash attention work in backward pass

2024-11-20 14:31:07 -08:00

cost_estimate.py

[Pallas] Add support for run_state to cost estimator.

2024-12-06 10:36:02 -08:00

pallas_call.py

Cleanup: remove uses of no-op raise_to_shaped

2024-12-12 09:49:06 -08:00

primitives.py

Stop using generators for linear_util transformations.

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

utils.py

[Pallas TPU] Add lowering for lax.nextafter

2024-10-31 17:34:38 -07:00