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

PiperOrigin-RevId: 738218113

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

..

__init__.py

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

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

block_spec.py

[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

BUILD

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

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

custom_evaluate.py

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

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

fusable_dtype.py

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

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

fusable.py

[Pallas] Add experimental (private for now) API for manual fusion into Pallas kernels

2025-03-03 17:05:51 -08:00

fuser_utils.py

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

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

fusion.py

Add fuser to jax.experimental.pallas

2025-03-03 17:26:44 -08:00

jaxpr_fusion.py

[Pallas/Fuser] Add debug option to fuser.fuse that prints out jaxpr

2025-03-10 14:07:26 -07:00