jax authors 414449e142 Merge pull request #26078 from gnecula:debug_info_jaxpr ...

PiperOrigin-RevId: 723151082

2025-02-04 10:54:26 -08:00

..

mosaic

[Pallas] Support promise_in_bounds mode in jnp.take_along_axis.

2025-02-03 22:06:19 -08:00

mosaic_gpu

[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

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] 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

[better_errors] Add debug info to more Jaxprs and WrappedFun (step 1)

2025-02-04 10:02:35 +02: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