Skip to content

Fix inplace/layout bugs in Numba lapack routines #1304

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

Conversation

ricardoV94
Copy link
Member

@ricardoV94 ricardoV94 commented Mar 19, 2025

  1. Numba's _copy_to_fortran_order does not seem to do the right thing with vectors and doesn't always copy them.
  2. Most methods were forgetting inputs still need to be F-contiguous even if they don't need to be copied.

PR did some tricks to avoid copying C-contiguous arrays when not needed.

Proof of the _copy_to_fortran_order thing:

import numpy as np
import numba
from numba.np.linalg import _copy_to_fortran_order

@numba.njit
def foo(x):
    out = _copy_to_fortran_order(x)
    return out
    
# Actually copies if x is a matrix
x = np.zeros((10, 1))
y = foo(x)
y[0] = np.pi
assert (x == 0).all()

# Doesn't copy if it's a vector
x = np.zeros((10))
y = foo(x)
y[0] = np.pi
assert (x == 0).all()  # AssertionError

@ricardoV94 ricardoV94 added bug Something isn't working numba linalg Linear algebra labels Mar 19, 2025
@ricardoV94 ricardoV94 force-pushed the fix_contiguity_bugs_numba_lapack_routines branch 3 times, most recently from a27bbf4 to cbae16b Compare March 19, 2025 16:04
Copy link

codecov bot commented Mar 19, 2025

Codecov Report

Attention: Patch coverage is 11.76471% with 60 lines in your changes missing coverage. Please review.

Project coverage is 81.98%. Comparing base (a149f6c) to head (689edc2).
Report is 4 commits behind head on main.

Files with missing lines Patch % Lines
pytensor/link/numba/dispatch/slinalg.py 11.76% 55 Missing and 5 partials ⚠️

❌ Your patch status has failed because the patch coverage (11.76%) is below the target coverage (100.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##             main    #1304      +/-   ##
==========================================
- Coverage   82.00%   81.98%   -0.02%     
==========================================
  Files         188      188              
  Lines       48478    48496      +18     
  Branches     8665     8671       +6     
==========================================
+ Hits        39755    39761       +6     
- Misses       6575     6583       +8     
- Partials     2148     2152       +4     
Files with missing lines Coverage Δ
pytensor/link/numba/dispatch/slinalg.py 44.09% <11.76%> (-0.36%) ⬇️
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@ricardoV94 ricardoV94 force-pushed the fix_contiguity_bugs_numba_lapack_routines branch 5 times, most recently from 57e5aa4 to be0df57 Compare March 21, 2025 06:12
@ricardoV94 ricardoV94 marked this pull request as ready for review March 21, 2025 06:33
@ricardoV94 ricardoV94 changed the title Fix contiguity bugs in Numba lapack routines Fix inplace/layout bugs in Numba lapack routines Mar 21, 2025
Copy link
Member

@jessegrabowski jessegrabowski left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Small comments that don't matter, but also double-check the logic in posv to make sure we can't do better, since the input is symmetric

assert np.allclose(A_val_c_contig, A_val) == (not overwrite_a)
# b vectors are always f_contiguous if also c_contiguous
assert np.allclose(b_val_c_contig, b_val) == (
not (overwrite_b and b_val_c_contig.flags.f_contiguous)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
not (overwrite_b and b_val_c_contig.flags.f_contiguous)
not (overwrite_b and b_val_c_contig.ndims == 1)

Is more clear? Checking that the c and f flags are both true is not obvious

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the col row matrices also have this property

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

then .squeee().ndims. I get that checking the other flag is elegant, but it requires detailed knowledge of how all these flags work, which isn't going to obvious to someone reading through these tests for the first time in the future

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No, the point is even if you defined it as C-contiguous array, it may still be F-contiguous in which case it will override just like the previous block did. I think this is the most readable test

overwrite_b: bool,
):
if is_complex:
# TODO: Complex raises ValueError: To change to a dtype of a different size, the last axis must be contiguous,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder if these bug fixes also fix the error reported here. Not for this PR, just thinking out loud.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't know what code you were running to see the error

Copy link
Member

@jessegrabowski jessegrabowski left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Satisfied on the posv point, the rest of the comments can be addressed or not as you see fit

@ricardoV94 ricardoV94 force-pushed the fix_contiguity_bugs_numba_lapack_routines branch 3 times, most recently from bf7c72a to 33e00fb Compare March 21, 2025 10:21
@ricardoV94 ricardoV94 force-pushed the fix_contiguity_bugs_numba_lapack_routines branch from 33e00fb to 689edc2 Compare March 21, 2025 10:22
@ricardoV94 ricardoV94 merged commit 39704d1 into pymc-devs:main Mar 21, 2025
67 checks passed
@ricardoV94 ricardoV94 deleted the fix_contiguity_bugs_numba_lapack_routines branch April 21, 2025 13:36
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
bug Something isn't working linalg Linear algebra numba
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants