Sergei Lebedev ec745f48c8 Use the current minimum jaxlib version for type checking on the CI

2024-10-10 12:46:15 +01:00

..

mosaic

Use the current minimum jaxlib version for type checking on the CI

2024-10-10 12:46:15 +01:00

mosaic_gpu

[Pallas/MGPU] Allow delaying the release of pipelined buffers

2024-10-08 08:17:58 -07:00

triton

Simplify and consolidate dot algorithm control in lax.

2024-10-07 13:21:34 -07:00

__init__.py

[Pallas] Upstream pallas to JAX

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

BUILD

Clean up BUILD files.

2024-08-26 09:11:17 -07:00

core.py

[pallas] Removed support for the deprecated pl.BlockSpec argument order

2024-10-03 14:39:58 -07:00

pallas_call.py

[Pallas] Add support for runtime checking of grid bounds using checkify.

2024-10-08 15:48:16 -07:00

primitives.py

[pallas] Added MemoryRef and run_scoped to the API docs

2024-10-07 15:35:09 -07:00

utils.py

[Pallas TPU] Add lowering for 64 bit

2024-09-19 16:42:45 +01:00