George Necula 904b74860c [better_errors] Continue adding debug info to Jaxprs (step 3) ...

This follows after #26078, and #26313, adding `debug_info` to
more calls to `lu.wrap_init`.

As part of this I have changed the primitives `custom_vjp_call_jaxpr`
and `custom_lin` to take the `bwd` parameter as a `lu.WrappedFun`,
which carries debug info. Previously, this was a `Callable`, but in
almost all cases if was really ` lu.WrappedFun.call_wrapped`.

2025-02-06 16:26:49 +02:00

..

mosaic

[Pallas TPU] Add helpers for writing collectives

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

mosaic_gpu

[pallas:mgpu] Cast all indices to i32 during lowering.

2025-02-05 07:37:04 -08:00

triton

[sharding_in_types] Add a canonicalize_value step before dispatching bind so that we can insert mesh_casts under the following conditions:

2025-02-03 18:00:19 -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

Merge pull request #26078 from gnecula:debug_info_jaxpr

2025-02-04 10:54:26 -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 TPU] Add helpers for writing collectives

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

hlo_interpreter.py

[better_errors] Continue adding debug info to Jaxprs (step 3)

2025-02-06 16:26:49 +02:00

pallas_call.py

[Pallas TPU] Add helpers for writing collectives

2025-02-04 15:39:10 -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 helpers for writing collectives

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