[pallas:fuser] Do not scalar prefetch inside nested calls in `push_block_spec` #41852

Merged
copybara-service[bot] opened 11:00am on October 21, 2026 wanted to merge 99 ฮ” into google/jax main from
test_1003941194
copybara-service[bot] authored
of work during October 22
Diff Delta:
99
About 42 Diff Delta/hour
Classified as:  General

copybara-service[bot]'s Description of Work

[pallas:fuser] Do not scalar prefetch inside nested calls in push_block_spec

Scalar prefetch values must come from fusion consts, but a nested call receives
all of its inputs as arguments. pull_block_spec already accounts for this,
while push_block_spec failed an assertion on nested ops with dynamic indices,
e.g. dynamic_update_slice, even when the push rule did not need them.

2 total changed files
Loading changes...
PR terminus reached:
Thus concludes todayโ€™s episode of โ€œJust One Small Changeโ€