[alejandroschuler]

Problem

On master, NGBClassifier(Dist=BetaBernoulli).fit(X, y) with binary y raises

ValueError: NGBClassifier Dist expects 3 classes, got 2.

PR #413 added NGBClassifier._fit_label_encoder, which counts the classes as Dist.n_params + 1. The __setstate__ fallback uses the same count. That holds for Bernoulli and k_categorical(K), but BetaBernoulli has 2 parameters (alpha and beta) for 2 classes. The example in docs/sympy_distributions.md fails for this reason. #413 is not in a release yet, so no user has seen this.

Fix

• ClassificationDistn now has an n_classes class attribute. A subclass can set it. If it does not, __init_subclass__ sets it to n_params + 1, so custom classification distributions keep the count they had before.
• k_categorical(K) sets n_classes = K.
• The SymPy factory sets n_classes = len(class_prob_exprs) on classification distributions.
• NGBClassifier._fit_label_encoder and NGBClassifier.__setstate__ use Dist.n_classes.

I checked the other places that use n_params + 1. NGBoost.__getstate__ (ngboost/ngboost.py:152) and Categorical.params (ngboost/distns/categorical.py:91) apply only to k_categorical distributions, where the count is correct, so I did not change them.

Tests

• test_classification_distns_set_n_classes checks n_classes for Bernoulli, k_categorical(4), BetaBernoulli and a custom subclass that uses the default.
• test_classifier_beta_bernoulli_binary_labels fits NGBClassifier(Dist=BetaBernoulli) on string labels and checks classes_, predict and the (n, 2) shape of predict_proba.

There is no pickle round trip test for this model yet. Factory distributions cannot be pickled until #445 merges. After that, a round trip of a fitted NGBClassifier(Dist=BetaBernoulli) is a good test to add.

black, isort, flake8 and pylint pass on the changed files, and pytest tests --slow passes.

🤖 Generated with Claude Code