Sergei Lebedev 0cfed4efad [pallas:mosaic_gpu] Shrink max_concurrent_iteration based on the total number of steps ...

PiperOrigin-RevId: 680990842

2024-10-01 06:19:43 -07:00

..

mosaic

[Pallas TPU] Add lowering for lax.cos_p

2024-09-30 16:12:11 -07:00

mosaic_gpu

[pallas:mosaic_gpu] Shrink max_concurrent_iteration based on the total number of steps

2024-10-01 06:19:43 -07:00

triton

[pallas:triton] Do not DCE the jaxpr in the lowering pass

2024-09-30 14:04:37 -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/MGPU] Undo transforms before giving refs back to users

2024-10-01 03:32:40 -07:00

pallas_call.py

[Pallas/MGPU] Undo transforms before giving refs back to users

2024-10-01 03:32:40 -07:00

primitives.py

[Pallas GPU] Enable Pallas OpsExtraTest in 64-bit mode

2024-09-20 16:18:31 -07:00

utils.py

[Pallas TPU] Add lowering for 64 bit

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