Sharad Vikram c6b164dc09 [Pallas/Fuser] Add custom evaluate to allow/disallow transposes ...

PiperOrigin-RevId: 735931978

2025-03-11 16:35:49 -07:00

..

fuser

[Pallas/Fuser] Add custom evaluate to allow/disallow transposes

2025-03-11 16:35:49 -07:00

mosaic

[Pallas TPU] Fix preferred_element_type propagation in dot_general with const

2025-03-11 15:06:07 -07:00

mosaic_gpu

[pallas:mosaic_gpu] Added support for math functions to the WG lowering

2025-03-10 05:08:19 -07: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

[Mosaic TPU][Pallas] Add pl.reciprocal

2025-03-07 18:29:30 -08:00

utils.py

[Pallas TPU] Add helpers for writing collectives

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