Christos Perivolaropoulos b23f8f414b [pallas/pallas_mgpu] Discharging run_scoped should not be discharging the intermediates ...

When we do run_scoped[jaxpr, R1,R2], it can't be assumed that references
corresponding to R1 and R2 can be safely discharged. Sometimes they can (eg
Accumulator) but sometimes they can't (eg SMEM scratch). It should be up to the
lowering rule to do such discharging.

This further means that during lowering there is no guarantee that the
references will not be used/returned by nested scoped blocks so we also remove
that check.

PiperOrigin-RevId: 722137352

2025-02-01 09:37:03 -08:00

..

mosaic

[pallas] Support DMA start partial discharge and run_scoped() does its own partial discharge.

2025-02-01 08:23:23 -08:00

mosaic_gpu

[pallas/pallas_mgpu] Discharging run_scoped should not be discharging the intermediates

2025-02-01 09:37:03 -08:00

triton

Rename out_type -> out_sharding parameter on einsum

2025-01-16 18:16:52 -08:00

__init__.py

[Pallas] Upstream pallas to JAX

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

BUILD

[Pallas] Refactor Pallas HLO interpret mode to a standalone file.

2025-01-27 17:52:27 -08:00

core.py

[Pallas] Refactor Pallas HLO interpret mode to a standalone file.

2025-01-27 17:52:27 -08:00

cost_estimate.py

[Pallas] Add support for run_state to cost estimator.

2024-12-06 10:36:02 -08:00

helpers.py

[Pallas] Add empty/empty_like helper functions

2025-01-08 10:49:11 -08:00

hlo_interpreter.py

[Pallas] Refactor Pallas HLO interpret mode to a standalone file.

2025-01-27 17:52:27 -08:00

pallas_call.py

[Pallas] Refactor Pallas HLO interpret mode to a standalone file.

2025-01-27 17:52:27 -08:00

primitives.py

[pallas/pallas_mgpu] Discharging run_scoped should not be discharging the intermediates

2025-02-01 09:37:03 -08:00

utils.py

[Pallas TPU] Add lowering for lax.nextafter

2024-10-31 17:34:38 -07:00