fix: linear fit overflow

[ianna]

Follow-up to #4232.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 83.98%. Comparing base (1b161d3) to head (447eefa).
✅ All tests successful. No failed tests found.

Additional details and impacted files

Files with missing lines	Coverage Δ
src/awkward/operations/ak_linear_fit.py	100.00% <100.00%> (ø)

[github-actions]

The documentation preview is ready to be viewed at http://preview.awkward-array.org.s3-website.us-east-1.amazonaws.com/PR4257

[TaiSakuma]

Thank you for the PR.

The test test_linear_fit_int64_overflow doesn't seem to have overflowing samples and appears to have passed without the fix, according to the review by Claude below.

If so, please include values that make the test fail without the fix and pass with it.

I copied the Claude review below.

---

The int32 fix is verified: on main, the inputs of test_linear_fit_int32_overflow give slope 2.0523 instead of 2.0; replaying the branch's promoted computation on the same inputs gives exactly slope 2.0, intercept 1.0. Scope first, then one requested change, one optional suggestion, and two minor notes.

Scope. No public API changes: the ak.linear_fit signature, defaults, and the returned LinearFit record (fields and dtypes) are unchanged; the new import comes from the private awkward._layout module. Behavior changes only for integer inputs — previously wrapped results become correct, and int64 sums above 2⁵³ shift from exact integer accumulation to float64 rounding, the same tradeoff #4232 accepted for ak.covar/ak.corr. No existing tests are modified: the diff adds two new test files and touches no existing assertion or expected value; the three existing files that exercise linear_fit (test_0115_generic_reducer_operation.py, test_1320_mask_identity_defaults.py, test_2754_highlevel_behavior_missing.py) are unchanged and pass in this PR's CI.

Requested change: give test_linear_fit_int64_overflow values large enough to wrap. 5_000_000**2 is 2.5×10¹³, far below the int64 maximum of about 9.2×10¹⁸; on main these inputs already give exactly slope=3.0, intercept=-7.0, so the test passes without the fix and cannot detect the overflow it is named for. Values around 3–5×10⁹ do wrap (main gives slope −0.0745) and would make the test detect a regression, but with them y = 3 * x - 7 no longer passes on the fixed code: the promoted fit's intercept is −7.000048, outside pytest.approx's default 10⁻⁶ relative tolerance, so the intercept needs an absolute tolerance or a larger magnitude:

x = np.array([3_000_000_000, 4_000_000_000, 5_000_000_000], dtype=np.int64)
y = 3 * x - 7_000_000  # promoted fit: slope 3.0, intercept -7000000.000018

Optional suggestion: promote per array, or drop the guard, so mixed integer/complex inputs do not wrap. The guard is_complex = "complex" in str(x.type) or "complex" in str(y.type) skips promotion of both arrays when either one is complex, so with int32 x and complex128 y, xp * xp still wraps:

import numpy as np
import awkward as ak

xi = np.array([100000, 200000, 300000, 400000], dtype=np.int32)
x = ak.Array(xi)
y = ak.Array((2 * xi + 1).astype(np.complex128))
ak.linear_fit(x, y)["slope"]  # (-0.4059-0j); expected (2+0j)

Running this snippet on main gives slope (-0.4059-0j); this branch takes the same path for these inputs (is_complex is true, so neither array is promoted), so the result is unchanged. The guard is also not needed for pure-complex inputs: promote_integral_to_float64 casts only when a leaf has dtype.kind in "biu", so complex (and float) data already passes through unchanged, and in their one-pass fallback paths ak.covar and ak.corr call the helper unconditionally — so the code comment "Matches how ak.covar/ak.corr promote in their one-pass paths" does not hold as written. Two options, either of which fixes the case above and removes the string comparison (which treats any type string containing "complex" as complex): drop the guard entirely, or guard per array with the structural _has_complex_leaf from ak_var.py (it could move to _layout.py next to _has_integral_leaf). The options differ only for a union layout mixing integral and complex leaves: dropping the guard promotes such a union and casts its complex leaf to float64, discarding imaginary parts, while the per-array guard keeps ak.var's protection there.

Minor notes:

• test_linear_fit_uint_and_negative contains no negative values, so the negative half of the name is untested.
• test_linear_fit_complex_not_regressed asserts only is not None; comparing against a hand-computed fit would also check the values.

🤖 Generated with Claude Code