Fix KNNImputer wrong donors for multi-NaN rows and f64 crash - #364
Merged
josevalim merged 5 commits intoAug 27, 2026
Merged
Conversation
nan_euclidean only zeroed nan_col (the single value currently being filled) before comparing a row to a candidate donor. Any OTHER missing value in that same row was left as NaN, so it poisoned the squared euclidean sum with NaN for every donor, regardless of how close that donor actually was. The result: every candidate distance came out as infinity, and top_k on a constant array just returns the first num_neighbors indices - so any row with two or more missing values was imputed from the first rows in the tensor instead of its real nearest neighbors. Rewrote the usable-columns mask to exclude any column where either the target row or the donor has NaN, not just nan_col, matching the present-features weighting sklearn's nan_euclidean_distances uses. The donor-missing-nan_col guard is unchanged and still correctly rules out the row itself as its own neighbor, since a row's own nan_col value is NaN by construction. Verified against sklearn.impute.KNNImputter directly, matches exactly on an adversarial case where the true nearest donors are far apart in row index from the ones the old code always picked.
Updated the two existing tests: their expected values encoded the bug (a row with 3 missing values was imputed from the first num_neighbors rows in the tensor). Recomputed against sklearn.impute.KNNImputter on the exact same matrix. Confirmed all 3 fail against the code they replace and pass here. Added a dedicated adversarial case where the true nearest donors sit at high row indices and the ones the old code always picked sit at low indices with wildly different values, so the two can't coincide by accident.
row_distances and the NaN placeholder used to hardcode type: {:f, 32}
regardless of the input's actual type. An f64 tensor crashed inside
the while loop as soon as a value needed imputing, since the distance
returned by nan_euclidean carried the input's real type while the loop
accumulator was fixed at f32. Both now take their type from the input
tensor.
Confirmed this fails with a CompileError against the code it replaces and passes here, with the result checked bit-exact against a plain f64 average (not just close within tolerance).
RicardoSantos-99
changed the title
Contributor
|
💚 💙 💜 💛 ❤️ |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
KNNImputercomputes distance to a candidate donor by zeroing out only the single column currently being filled. Any other missing value in that same row is left as NaN, so it poisons the distance sum for every donor. Every candidate ends up at infinity, and picking "nearest" from a constant array just returns the firstnum_neighborsrows in the tensor. So any row with two or more missing values gets imputed from the first rows in the data instead of its real nearest neighbors, silently.Fixed by excluding any column where either row has NaN, not just the one being filled, matching the present-features weighting
sklearn.impute.KNNImputer'snan_euclidean_distancesuses. Verified against it directly: on a case where the true nearest donors sit far from the ones the old code always picked, the old code returns[1.05, 2.05, 5.0], this returns[55.0, 85.0, 5.0], matching sklearn exactly.Also fixed a separate bug found while testing: the distance accumulator and the NaN placeholder both hardcoded f32, so an f64 tensor crashed as soon as a value needed imputing. Both now take their type from the input.