[AxelNoun]

Summary

Fixes #23797

interpolate_pos_emb passed a permuted view of the bf16 positional table (strides (1536, 1, 49152, 1536)) to F.grid_sample. For Half and BFloat16 the CPU kernel computes gather offsets in int16, and only float32 is guarded against offset overflow, so the H stride of 49152 makes the offsets wrap. On the 32x32 grid about 98% of the outputs are wrong, some NaN, and the reads that land before the storage are consistent with the intermittent grid_sample segfault in the macOS unittest job. The issue has the details and references.

The underlying kernel bug is reported upstream in pytorch/pytorch#200949.

This change makes the view contiguous before grid_sample. The largest gather offset becomes 1023, and the interpolation stays in the table's dtype, as the docstring requires. The tests only checked shapes and dtypes: test_identity_grid_matches_table now requires the 32x32 interpolation to return the table exactly, and the new test_downscaled_matches_float32 compares the 8x8 bf16 result with a float32 reference, within 3 * eps times the weighted absolute sum of the texels. That bound covers the three bf16 roundings of the accumulation with a factor of two to spare.

Computing in float32 and casting once would also avoid the kernel issue. However, it departs from the documented computation in the table's dtype and changes the outputs: about a third of the 8x8 values, by up to 0.0156, and about 97% of the 74x74 values, by up to about 0.75 (the separate question in the issue). I kept the smaller change.

cc @mergennachin @iseeyuan @lucylq @helunwencser @tarun292 @kimishpatel @jackzhxng @nil-is-all @digantdesai, who added this example in #21711.

Test plan

python -m pytest examples/models/muse-glimmer/tests/test_vision_precompute.py

I ran this test file locally on Windows (x86 AVX2), three times under torch 2.14.0+cpu (Python 3.12) and three times under torch 2.11.0+cu128 (Python 3.14). Neither environment has an editable install, so a small wrapper made executorch.examples.models.muse_glimmer resolve to this branch's sources before calling pytest with the repository's pytest.ini. With this change, every test passes, except ResizeReferenceParityTest::test_matches_pinned_golden under torch 2.14, which fails in that environment only because PIL is not installed there. With precompute.py from main and the new tests, test_identity_grid_matches_table fails (97.5% to 97.8% of the values mismatched), and so does test_downscaled_matches_float32 (more than 96000 of 98304 values out of the bound).

Over 50 seeds, the identity case is exact and the 8x8 error stays below 1.34 times eps times the weighted absolute sum, against 3 in the test. ufmt check and flake8, at the versions pinned in requirements-lintrunner.txt, report no findings; the torchfix plugin was not installed. I could not run the macOS arm64 job.

Disclosure: this PR was authored with AI assistance (Claude Code).

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/23798

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 87b7e52 with merge base a90db24 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[AxelNoun]

@pytorchbot label "release notes: examples"

[AxelNoun]

The test-muse-glimmer-cuda-e2e-kquant-17gb-dflash-image failure happens in install_executorch.sh, before the test starts: pip rejected nvidia_nvjitlink-13.4.92 with a hash mismatch while the download was at 16.8 of 42.5 MB. The three other muse-glimmer jobs of the same attempt installed the same wheel without error, and this job passed on the last eight completed main runs that included it, including at the merge base 060075b. Could you re-run it once the workflow run has completed? I don't have the permissions to do so.