Skip to content

Numba sparse dot: use final precision in intermediate computations#2159

Open
ricardoV94 wants to merge 1 commit into
pymc-devs:mainfrom
ricardoV94:flaky_test
Open

Numba sparse dot: use final precision in intermediate computations#2159
ricardoV94 wants to merge 1 commit into
pymc-devs:mainfrom
ricardoV94:flaky_test

Conversation

@ricardoV94
Copy link
Copy Markdown
Member

@ricardoV94 ricardoV94 commented May 20, 2026

Fixes a flaky test (see https://github.com/pymc-devs/pytensor/actions/runs/26150694899/job/76918627812#step:6:5422)

It doesn't cause meaningful slowdown in some benchmarks I did and it matches scipy in precision for mixed dtypes

Numba does scalar int32|64 * array float32 in float32 precision, whereas numpy does it in float64 precision. Similar for discrete * complex64

Numba underpromotes relative to numpy/scipy mixed scalar * array dtypes
col_idx = x_ind[idx]
value = x_data[idx]
value = out_type(x_data[idx])
z[row_idx] += value * y[col_idx]
Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

these are the branches where we do scalar * array, that numba may underpromote relative to numpy/scipy

@ricardoV94 ricardoV94 changed the title Numba sparse dot: Use final precision in intermediate computations Numba sparse dot: use final precision in intermediate computations May 20, 2026
@ricardoV94 ricardoV94 requested a review from tomicapretto May 20, 2026 12:12
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant