Fall back to caller location for bare-NameLoc MLIR callsites in source_mapper. #41865

Merged
copybara-service[bot] opened 3:41pm on October 21, 2026 wanted to merge 1 commit into google/jax main from
test_1003520552

Pull Request Overview

  • Opened on October 21, 2026
  • Status Merged
  • Commit count 1 with first commit October 21, 2026

Total Delta

0 Total Diff Delta

Open Days

Open 17 weekdays

Test Delta

0 Diff Delta in Test Files
Breakdown by Phase

How long has this pull request spent in each phase of its lifecycle?

Data pending calculation for pull request

Author avatar

Fall back to caller location for bare-NameLoc MLIR callsites in source_mapper.

Fall back to caller location for bare-NameLoc MLIR callsites in source_mapper.

In modern Pallas and Mosaic MLIR debug dumps (such as pltpu.emit_pipeline), callsite locations frequently have a bare NameLoc or scoped NameLoc callee with no FileLineColLoc, while the caller location carries the originating Python source line. Previously, parse_mlir_locations treated callsites as unconditional redirects to the callee ID, raising KeyError whenever a bare NameLoc was absent from locations.

Record both callee and caller on Callsite entries and resolve locations recursively with cycle detection, preferring callee when it resolves to a Location and falling back to caller otherwise. Unresolved bare NameLoc entries without file locations are skipped cleanly.

PR was closed without comments.