Sergei Lebedev e5102957b0 [pallas:mosaic_gpu] Do not store the grid mapping in ModuleContext ...

We really only ever use the grid names.

PiperOrigin-RevId: 703108864

2024-12-05 07:32:43 -08:00

..

mosaic

[Mosaic TPU] Add support for modeling loads/stores and fix minor issues in model extraction

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

mosaic_gpu

[pallas:mosaic_gpu] Do not store the grid mapping in ModuleContext

2024-12-05 07:32:43 -08:00

triton

[Pallas] Fix type annotation on TritonCompilerParams

2024-12-04 12:47:41 -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] Use Pallas cost estimator for flash attention.

2024-11-20 17:12:37 -08:00

pallas_call.py

[pallas] Minor simplifications to Pallas interpreter.

2024-11-18 17:10:10 -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