Fix NaiveBayes prediction paths and prior dtype - #359
Merged
josevalim merged 5 commits intoAug 23, 2026
Conversation
fit stores the binarize threshold and applies it before counting features, but predict/predict_log_probability/predict_probability/ predict_joint_log_probability passed x straight to the likelihood computation. Any input that wasn't already 0/1 (the common case, since binarize defaults to 0.0 specifically so raw counts can be passed in) produced wrong log-likelihoods, since the formula assumes binary x. Stores binarize on the struct as a plain (non-tensor) field and applies it once, centrally, in joint_log_likelihood, so every predict path gets it. Verified against sklearn.naive_bayes.BernoulliNB on the same data: predict_log_probability/predict_probability now match sklearn exactly (previously off by orders of magnitude), and predict_joint_log_probability matches sklearn's internal _joint_log_likelihood computed on the properly-binarized input.
…tion
check_dim compared x's feature count against
Nx.axis_size(model.feature_count, 1), but feature_count here is
{num_features, num_classes, num_categories} (unlike the 2D
{num_classes, num_features} in the other NaiveBayes modules this was
copied from), so axis 1 is num_classes. Every predict/*_probability
call rejected correctly-shaped input whenever num_features didn't
happen to equal num_classes, and would silently accept wrongly-shaped
input in the reverse case. Fixed to read axis 0.
Separately, fit's min_categories option was dead: Keyword.pop removed
:min_categories from opts and rebound it, so the very next read of
opts[:min_categories] (used to size the feature_count/
feature_log_probability tensors) always saw nil and silently fell back
to inferring the category count from the training data alone. Kept a
flag captured before the pop and branched on it instead, matching
sklearn.naive_bayes.CategoricalNB's min_categories semantics (a literal
per-feature category count, not a max index) - verified shapes match
sklearn exactly for both list and tensor min_categories inputs.
joint_log_likelihood sized its accumulator from x, as
{num_samples, num_features}, but each loop iteration contributes a
{num_samples, num_classes} term (feature_log_probability is
{num_features, num_classes, num_categories}, so the class axis is 1).
The two only agree when num_features happens to equal num_classes,
which is the case in every existing doctest and test - x is
Nx.iota({4, 3}) with num_classes: 3 - so the whole predict path was
only ever exercised in the one shape where the bug is invisible. Any
other dataset raised "cannot broadcast tensor of dimensions {n, c} to
{n, f}".
Same num_features/num_classes confusion as the check_dim bug. Also
pins the squeeze to axis 0 so a model with a single class or a single
category does not lose an extra axis.
Verified against sklearn.naive_bayes.CategoricalNB on a 3-feature,
2-class dataset across default/min_categories/class_prior/fit_prior:
predictions identical and probabilities within f32 precision.
class_log_priors was computed in f32 regardless of the input type on
two of its three branches, while class_count and
feature_log_probability followed the input, so an f64 model silently
carried an f32 prior into every joint log-likelihood:
* explicit priors: class_priors was built with Nx.tensor/1 without
type:, unlike alpha and sample_weights next to it, which both pass
type: to_float_type(x).
* fit_priors: false: Nx.log/1 on the bare num_classes integer
defaults to f32.
Multiplying a typed literal by num_classes keeps the log in the target
type rather than widening an f32 result afterwards, which would keep
f32 precision. Verified that -log(3) and log(0.3) now come back
bit-exact against the f64 reference (previously off by ~2e-8), and
that class_log_priors matches feature_log_probability's type across
all three prior branches, both backends.
Applies to all four modules since the block is identical in each.
Each test was checked against the pre-fix code: seven of them fail there and pass here, so they detect the bugs rather than merely documenting current behaviour. Expected values come from sklearn.naive_bayes, not from Scholar's own output.
RicardoSantos-99
force-pushed
the
fix-naive-bayes-predict-bugs
branch
from
August 23, 2026 20:55
2a73fb0 to
54889c1
Compare
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.
Five bugs in the NaiveBayes modules, four on the prediction path.
:binarizeat predict time. The threshold is not kept on the model, so everypredict*passes rawxinto a likelihood that assumes binary features.:binarizedefaults to0.0, so this hits the default configuration.check_dimreads axis 1 offeature_count, which is the class count here (the tensor is{num_features, num_classes, num_categories}, not 2-D as in the other modules). It rejects valid input and accepts invalid input whenevernum_features != num_classes.joint_log_likelihoodsizes its accumulator fromxrather than from the class axis, so predict raises for any dataset where the two differ. Every doctest in the module usesnum_features == num_classes, the one shape where this is invisible.:min_categoriesis dead code.Keyword.poprebindsoptsbefore the value is read, so the category count always falls back to being inferred from the data.class_log_priorsis computed in f32 regardless of input type, whileclass_countandfeature_log_probabilityfollow the input. The block is identical in all four modules, so all four are fixed.:alphaalso gets a non-negative check through the existingScholar.Options.non_negative_number. A negative value previously produced NaN with no error.Behaviour changes:
%Bernoulli{}gains a:binarizefield; BernoulliNB predictions change for input that is not already binary, and its doctests are updated;fit/3raises on a negative:alpha; CategoricalNB models built with:min_categorieschange shape.