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.
0 total changed files
Thatβs the last hunk
Approve, request changes, or stare into the void (reviewer's choice) π
google/jax