jax authors 01a110c4c9 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. ...

PiperOrigin-RevId: 738171437

2025-03-18 15:51:15 -07:00

..

fuser

Add a variant of safe_map() that has no return value, named foreach().

2025-03-14 07:42:48 -07:00

mosaic

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

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