Gleb Pobudzey 43b6be0e81 [Mosaic GPU] Add lowering for log, and a fast path using log2. ...

PiperOrigin-RevId: 733411276

2025-03-04 11:50:50 -08:00

..

fuser

[Pallas/Fuser] Add support for custom_call_jvp/pjit to push_block_spec

2025-03-03 17:43:13 -08:00

mosaic

Chnages for kernel export

2025-03-01 00:32:39 -08:00

mosaic_gpu

[Mosaic GPU] Add lowering for log, and a fast path using log2.

2025-03-04 11:50:50 -08:00

triton

[pallas:triton] Emit a better error message for matmul with non-2D operands

2025-03-04 05:46:29 -08:00

__init__.py

[Pallas] Upstream pallas to JAX

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

BUILD

[JAX] Allow pallas to accept scalar shape semaphores.

2025-02-14 23:20:47 -08:00

core.py

[Pallas] Add name parameter to core_map

2025-02-26 18:59:01 -08:00

cost_estimate.py

[better_errors] Continue adding debug info to Jaxprs (step 7)

2025-02-09 18:14:33 +02:00

helpers.py

[Pallas TPU] Add helpers for writing collectives

2025-02-04 15:39:10 -08:00

hlo_interpreter.py

[better_errors] Port the Pallas debug info mechanisms to the new JAX DebugInfo.

2025-02-25 14:43:17 +01:00

pallas_call.py

[Pallas] Add option for async DMAs in the new TPU interpret mode

2025-02-25 18:19:20 -08:00

primitives.py

Refactor Jax FFI lowering to prepare for implementing CPU/GPU callbacks using XLA's FFI.

2025-02-21 09:45:59 -08:00

utils.py

[Pallas TPU] Add helpers for writing collectives

2025-02-04 15:39:10 -08:00