[typing] Avoid non-contiguous overloads in splash_attention_kernel.py #41856

Merged
stroxler opened 11:38am on October 21, 2026 wanted to merge 1 commit into google/jax main from
fix-noncontiguous-overload

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?

Fraction of total time Business days Phase
 
0.0 days Authoring 1 commit before pull request opened for review
 
0.1 days Awaiting first review
 
0.1 days Revising work with 0 commits in response to 1 review that left 1 comment
Set up deploy tracking to begin tracking full PR start-to-deploy times

Total time for pull request 0.2 business days from first commit authored to merge
Deploy tracking has not been set up

Author avatar

[typing] Avoid non-contiguous overloads in splash_attention_kernel.py

The typing spec says that overload function definitions and the implementation must be contiguous.

While there are pros and cons to this (I can see how defining a helper like _div right above the implementation is handy) it does seem like a reasonable rule for readability. A PR to Pyrefly adding this missing check flagged this kernel in mypy_primer, so I figured I'd put up a fix since the next release will flag it.

Comments Threads Pending Resolution

jakevdp reviewed and approved this work on October 21, 2026
jakevdp left a comment

Thanks!

Resolved Comment Threads

No resolved comments have been left on this PR.