jax authors e9ce8fb92d Merge pull request #27227 from jburnim:jburnim_pallas_interpret_mode4 ...

PiperOrigin-RevId: 738235363

2025-03-18 20:22:27 -07:00

..

fuser

[Pallas/Fuser] DCE fusion jaxprs before pulling (to avoid unnecessary computations being staged out in block functions)

2025-03-18 19:00:41 -07:00

mosaic

Merge pull request #27227 from jburnim:jburnim_pallas_interpret_mode4

2025-03-18 20:22:27 -07:00

mosaic_gpu

[pallas] pallas_call_p is now parameterized by a mesh

2025-03-17 16:30:40 -07:00

triton

[pallas] pallas_call_p is now parameterized by a mesh

2025-03-17 16:30:40 -07: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

Better mosaic lowering for dynamic shapes, extend an interpreter into shape_poly dimexpr and lower them alongside the graph if we are in a dynamic export regime.

2025-03-18 15:51:15 -07: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

[pallas] pallas_call_p is now parameterized by a mesh

2025-03-17 16:30:40 -07:00

pallas_call.py

[pallas] pallas_call_p is now parameterized by a mesh

2025-03-17 16:30:40 -07:00

primitives.py

Add explicit support for float8_e4m3b11fnuz in pl.dot

2025-03-14 02:51:55 -07:00

utils.py

[Pallas TPU] Add helpers for writing collectives

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