From 05bcabd6140726acda4dd558f725c3ef99e624a5 Mon Sep 17 00:00:00 2001 From: Ricardo Carvalho Santos Date: Sun, 23 Aug 2026 17:12:18 -0300 Subject: [PATCH 1/5] Fix BernoulliNB never binarizing input at predict time 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. --- lib/scholar/naive_bayes/bernoulli.ex | 33 ++++++++++++++++++---------- 1 file changed, 22 insertions(+), 11 deletions(-) diff --git a/lib/scholar/naive_bayes/bernoulli.ex b/lib/scholar/naive_bayes/bernoulli.ex index 61602334..e814b417 100644 --- a/lib/scholar/naive_bayes/bernoulli.ex +++ b/lib/scholar/naive_bayes/bernoulli.ex @@ -10,13 +10,14 @@ defmodule Scholar.NaiveBayes.Bernoulli do import Scholar.Shared @derive {Nx.Container, + keep: [:binarize], containers: [ :feature_count, :class_count, :class_log_priors, :feature_log_probability ]} - defstruct [:feature_count, :class_count, :class_log_priors, :feature_log_probability] + defstruct [:feature_count, :class_count, :class_log_priors, :feature_log_probability, :binarize] opts_schema = [ num_classes: [ @@ -127,7 +128,8 @@ defmodule Scholar.NaiveBayes.Bernoulli do [-1.0986123085021973, -1.0986123085021973, -0.40546512603759766], [-0.28768205642700195, -0.28768205642700195, -0.28768205642700195] ] - ) + ), + binarize: 1.0 } iex> x = Nx.iota({4, 3}) @@ -153,7 +155,8 @@ defmodule Scholar.NaiveBayes.Bernoulli do [-23.025850296020508, 0.0, 0.0], [0.0, 0.0, 0.0] ] - ) + ), + binarize: 0.0 } """ @@ -288,7 +291,8 @@ defmodule Scholar.NaiveBayes.Bernoulli do class_count: class_count, class_log_priors: class_log_priors, feature_count: feature_count, - feature_log_probability: feature_log_probability + feature_log_probability: feature_log_probability, + binarize: opts[:binarize] } end @@ -343,8 +347,8 @@ defmodule Scholar.NaiveBayes.Bernoulli do #Nx.Tensor< f32[2][3] [ - [-4.7047806, -12.329399, -0.009097099], - [-8.750494, -19.147701, -1.5830994e-4] + [-1.4696369, -2.162784, -0.42314053], + [-1.4696369, -2.162784, -0.42314053] ] > """ @@ -373,8 +377,8 @@ defmodule Scholar.NaiveBayes.Bernoulli do #Nx.Tensor< f32[2][3] [ - [0.0090519, 4.419875e-6, 0.99094415], - [1.5838306e-4, 4.8334696e-9, 0.9998417] + [0.23000899, 0.11500449, 0.65498656], + [0.23000899, 0.11500449, 0.65498656] ] > """ @@ -395,8 +399,8 @@ defmodule Scholar.NaiveBayes.Bernoulli do #Nx.Tensor< f32[2][3] [ - [3.6356335, -3.988985, 8.331317], - [10.567104, 0.16989732, 19.31744] + [-2.6026897, -3.295837, -1.5561934], + [-2.6026897, -3.295837, -1.5561934] ] > """ @@ -421,10 +425,17 @@ defmodule Scholar.NaiveBayes.Bernoulli do defnp joint_log_likelihood( %__MODULE__{ feature_log_probability: feature_log_probability, - class_log_priors: class_log_priors + class_log_priors: class_log_priors, + binarize: binarize }, x ) do + x = + case binarize do + nil -> x + threshold -> Scholar.Preprocessing.Binarizer.fit_transform(x, threshold: threshold) + end + neg_prob = Nx.log(1 - Nx.exp(feature_log_probability)) jll = Nx.dot(x, [1], feature_log_probability - neg_prob, [1]) jll + class_log_priors + Nx.sum(neg_prob, axes: [1]) From 40efdd1fbce3e947738691954299776316a354ab Mon Sep 17 00:00:00 2001 From: Ricardo Carvalho Santos Date: Sun, 23 Aug 2026 17:12:18 -0300 Subject: [PATCH 2/5] Fix CategoricalNB predict's feature-count check and min_categories option 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. --- lib/scholar/naive_bayes/categorical.ex | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/lib/scholar/naive_bayes/categorical.ex b/lib/scholar/naive_bayes/categorical.ex index 3fdd9d66..d9196d4c 100644 --- a/lib/scholar/naive_bayes/categorical.ex +++ b/lib/scholar/naive_bayes/categorical.ex @@ -272,10 +272,11 @@ defmodule Scholar.NaiveBayes.Categorical do end num_categories = - (opts[:min_categories] || x) - |> Nx.reduce_max() - |> Nx.add(1) - |> Nx.to_number() + if min_categories_flag do + min_categories |> Nx.reduce_max() |> Nx.to_number() |> trunc() + else + x |> Nx.reduce_max() |> Nx.add(1) |> Nx.to_number() |> trunc() + end opts = opts ++ @@ -360,7 +361,7 @@ defmodule Scholar.NaiveBayes.Categorical do """ defn predict(%__MODULE__{} = model, x, classes) do - check_dim(x, Nx.axis_size(model.feature_count, 1)) + check_dim(x, Nx.axis_size(model.feature_count, 0)) if Nx.rank(classes) != 1 do raise ArgumentError, @@ -401,7 +402,7 @@ defmodule Scholar.NaiveBayes.Categorical do """ defn predict_log_probability(%__MODULE__{} = model, x) do - check_dim(x, Nx.axis_size(model.feature_count, 1)) + check_dim(x, Nx.axis_size(model.feature_count, 0)) jll = joint_log_likelihood(model, x) log_proba_x = @@ -453,7 +454,7 @@ defmodule Scholar.NaiveBayes.Categorical do """ defn predict_joint_log_probability(%__MODULE__{} = model, x) do - check_dim(x, Nx.axis_size(model.feature_count, 1)) + check_dim(x, Nx.axis_size(model.feature_count, 0)) joint_log_likelihood(model, x) end From cd5435a9237f77db744ffa2cd15e5adbdd2a4fab Mon Sep 17 00:00:00 2001 From: Ricardo Carvalho Santos Date: Sun, 23 Aug 2026 17:40:23 -0300 Subject: [PATCH 3/5] Fix CategoricalNB predict for datasets where num_features != num_classes 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. --- lib/scholar/naive_bayes/bernoulli.ex | 33 +++++++++----------------- lib/scholar/naive_bayes/categorical.ex | 24 +++++++++++++++---- 2 files changed, 30 insertions(+), 27 deletions(-) diff --git a/lib/scholar/naive_bayes/bernoulli.ex b/lib/scholar/naive_bayes/bernoulli.ex index e814b417..61602334 100644 --- a/lib/scholar/naive_bayes/bernoulli.ex +++ b/lib/scholar/naive_bayes/bernoulli.ex @@ -10,14 +10,13 @@ defmodule Scholar.NaiveBayes.Bernoulli do import Scholar.Shared @derive {Nx.Container, - keep: [:binarize], containers: [ :feature_count, :class_count, :class_log_priors, :feature_log_probability ]} - defstruct [:feature_count, :class_count, :class_log_priors, :feature_log_probability, :binarize] + defstruct [:feature_count, :class_count, :class_log_priors, :feature_log_probability] opts_schema = [ num_classes: [ @@ -128,8 +127,7 @@ defmodule Scholar.NaiveBayes.Bernoulli do [-1.0986123085021973, -1.0986123085021973, -0.40546512603759766], [-0.28768205642700195, -0.28768205642700195, -0.28768205642700195] ] - ), - binarize: 1.0 + ) } iex> x = Nx.iota({4, 3}) @@ -155,8 +153,7 @@ defmodule Scholar.NaiveBayes.Bernoulli do [-23.025850296020508, 0.0, 0.0], [0.0, 0.0, 0.0] ] - ), - binarize: 0.0 + ) } """ @@ -291,8 +288,7 @@ defmodule Scholar.NaiveBayes.Bernoulli do class_count: class_count, class_log_priors: class_log_priors, feature_count: feature_count, - feature_log_probability: feature_log_probability, - binarize: opts[:binarize] + feature_log_probability: feature_log_probability } end @@ -347,8 +343,8 @@ defmodule Scholar.NaiveBayes.Bernoulli do #Nx.Tensor< f32[2][3] [ - [-1.4696369, -2.162784, -0.42314053], - [-1.4696369, -2.162784, -0.42314053] + [-4.7047806, -12.329399, -0.009097099], + [-8.750494, -19.147701, -1.5830994e-4] ] > """ @@ -377,8 +373,8 @@ defmodule Scholar.NaiveBayes.Bernoulli do #Nx.Tensor< f32[2][3] [ - [0.23000899, 0.11500449, 0.65498656], - [0.23000899, 0.11500449, 0.65498656] + [0.0090519, 4.419875e-6, 0.99094415], + [1.5838306e-4, 4.8334696e-9, 0.9998417] ] > """ @@ -399,8 +395,8 @@ defmodule Scholar.NaiveBayes.Bernoulli do #Nx.Tensor< f32[2][3] [ - [-2.6026897, -3.295837, -1.5561934], - [-2.6026897, -3.295837, -1.5561934] + [3.6356335, -3.988985, 8.331317], + [10.567104, 0.16989732, 19.31744] ] > """ @@ -425,17 +421,10 @@ defmodule Scholar.NaiveBayes.Bernoulli do defnp joint_log_likelihood( %__MODULE__{ feature_log_probability: feature_log_probability, - class_log_priors: class_log_priors, - binarize: binarize + class_log_priors: class_log_priors }, x ) do - x = - case binarize do - nil -> x - threshold -> Scholar.Preprocessing.Binarizer.fit_transform(x, threshold: threshold) - end - neg_prob = Nx.log(1 - Nx.exp(feature_log_probability)) jll = Nx.dot(x, [1], feature_log_probability - neg_prob, [1]) jll + class_log_priors + Nx.sum(neg_prob, axes: [1]) diff --git a/lib/scholar/naive_bayes/categorical.ex b/lib/scholar/naive_bayes/categorical.ex index d9196d4c..85d152e2 100644 --- a/lib/scholar/naive_bayes/categorical.ex +++ b/lib/scholar/naive_bayes/categorical.ex @@ -29,7 +29,12 @@ defmodule Scholar.NaiveBayes.Categorical do """ ], alpha: [ - type: {:or, [:float, {:list, :float}]}, + type: + {:or, + [ + {:custom, Scholar.Options, :non_negative_number, []}, + {:list, {:custom, Scholar.Options, :non_negative_number, []}} + ]}, default: 1.0, doc: ~S""" Additive (Laplace/Lidstone) smoothing parameter @@ -235,7 +240,7 @@ defmodule Scholar.NaiveBayes.Categorical do priors_flag = opts[:class_priors] != nil {class_priors, opts} = Keyword.pop(opts, :class_priors, :nan) - class_priors = Nx.tensor(class_priors) + class_priors = Nx.tensor(class_priors, type: type) if priors_flag and Nx.size(class_priors) != num_classes do raise ArgumentError, @@ -333,7 +338,8 @@ defmodule Scholar.NaiveBayes.Categorical do Nx.log(class_count) - Nx.log(Nx.sum(class_count)) true -> - Nx.broadcast(-Nx.log(num_classes), {num_classes}) + num_classes_t = Nx.tensor(1.0, type: type) * num_classes + Nx.broadcast(-Nx.log(num_classes_t), {num_classes}) end %__MODULE__{ @@ -477,14 +483,22 @@ defmodule Scholar.NaiveBayes.Categorical do }, x ) do + # jll accumulates one {num_samples, num_classes} term per feature. Note + # that this is not the shape of x: feature_log_probability is + # {num_features, num_classes, num_categories}, so num_classes comes from + # its axis 1, not from x's second axis. + num_samples = Nx.axis_size(x, 0) + num_classes = Nx.axis_size(feature_log_probability, 1) + {_, jll} = - while {{i = 0, feature_log_probability, x}, jll = Nx.broadcast(0.0, Nx.shape(x))}, + while {{i = 0, feature_log_probability, x}, + jll = Nx.broadcast(0.0, {num_samples, num_classes})}, i < Nx.axis_size(x, 1) do indices = Nx.slice_along_axis(x, i, 1, axis: 1) |> Nx.squeeze(axes: [1]) jll = Nx.slice_along_axis(feature_log_probability, i, 1, axis: 0) - |> Nx.squeeze() + |> Nx.squeeze(axes: [0]) |> Nx.take(indices, axis: 1) |> Nx.transpose() |> Nx.add(jll) From 0ebe9d8d57e5a366089b48b6e2fc57eef5f5cb93 Mon Sep 17 00:00:00 2001 From: Ricardo Carvalho Santos Date: Sun, 23 Aug 2026 17:40:23 -0300 Subject: [PATCH 4/5] Keep NaiveBayes class_log_priors in the model's float type 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. --- lib/scholar/naive_bayes/bernoulli.ex | 45 ++++++++++++++++++-------- lib/scholar/naive_bayes/complement.ex | 5 +-- lib/scholar/naive_bayes/multinomial.ex | 5 +-- 3 files changed, 37 insertions(+), 18 deletions(-) diff --git a/lib/scholar/naive_bayes/bernoulli.ex b/lib/scholar/naive_bayes/bernoulli.ex index 61602334..943e6b8d 100644 --- a/lib/scholar/naive_bayes/bernoulli.ex +++ b/lib/scholar/naive_bayes/bernoulli.ex @@ -10,13 +10,14 @@ defmodule Scholar.NaiveBayes.Bernoulli do import Scholar.Shared @derive {Nx.Container, + keep: [:binarize], containers: [ :feature_count, :class_count, :class_log_priors, :feature_log_probability ]} - defstruct [:feature_count, :class_count, :class_log_priors, :feature_log_probability] + defstruct [:feature_count, :class_count, :class_log_priors, :feature_log_probability, :binarize] opts_schema = [ num_classes: [ @@ -27,7 +28,12 @@ defmodule Scholar.NaiveBayes.Bernoulli do """ ], alpha: [ - type: {:or, [:float, {:list, :float}]}, + type: + {:or, + [ + {:custom, Scholar.Options, :non_negative_number, []}, + {:list, {:custom, Scholar.Options, :non_negative_number, []}} + ]}, default: 1.0, doc: ~S""" Additive (Laplace/Lidstone) smoothing parameter @@ -127,7 +133,8 @@ defmodule Scholar.NaiveBayes.Bernoulli do [-1.0986123085021973, -1.0986123085021973, -0.40546512603759766], [-0.28768205642700195, -0.28768205642700195, -0.28768205642700195] ] - ) + ), + binarize: 1.0 } iex> x = Nx.iota({4, 3}) @@ -153,7 +160,8 @@ defmodule Scholar.NaiveBayes.Bernoulli do [-23.025850296020508, 0.0, 0.0], [0.0, 0.0, 0.0] ] - ) + ), + binarize: 0.0 } """ @@ -203,7 +211,7 @@ defmodule Scholar.NaiveBayes.Bernoulli do priors_flag = opts[:class_priors] != nil {class_priors, opts} = Keyword.pop(opts, :class_priors, :nan) - class_priors = Nx.tensor(class_priors) + class_priors = Nx.tensor(class_priors, type: type) if priors_flag and Nx.size(class_priors) != num_classes do raise ArgumentError, @@ -281,14 +289,16 @@ defmodule Scholar.NaiveBayes.Bernoulli do Nx.log(class_count) - Nx.log(Nx.sum(class_count)) true -> - Nx.broadcast(-Nx.log(num_classes), {num_classes}) + num_classes_t = Nx.tensor(1.0, type: type) * num_classes + Nx.broadcast(-Nx.log(num_classes_t), {num_classes}) end %__MODULE__{ class_count: class_count, class_log_priors: class_log_priors, feature_count: feature_count, - feature_log_probability: feature_log_probability + feature_log_probability: feature_log_probability, + binarize: opts[:binarize] } end @@ -343,8 +353,8 @@ defmodule Scholar.NaiveBayes.Bernoulli do #Nx.Tensor< f32[2][3] [ - [-4.7047806, -12.329399, -0.009097099], - [-8.750494, -19.147701, -1.5830994e-4] + [-1.4696369, -2.162784, -0.42314053], + [-1.4696369, -2.162784, -0.42314053] ] > """ @@ -373,8 +383,8 @@ defmodule Scholar.NaiveBayes.Bernoulli do #Nx.Tensor< f32[2][3] [ - [0.0090519, 4.419875e-6, 0.99094415], - [1.5838306e-4, 4.8334696e-9, 0.9998417] + [0.23000899, 0.11500449, 0.65498656], + [0.23000899, 0.11500449, 0.65498656] ] > """ @@ -395,8 +405,8 @@ defmodule Scholar.NaiveBayes.Bernoulli do #Nx.Tensor< f32[2][3] [ - [3.6356335, -3.988985, 8.331317], - [10.567104, 0.16989732, 19.31744] + [-2.6026897, -3.295837, -1.5561934], + [-2.6026897, -3.295837, -1.5561934] ] > """ @@ -421,10 +431,17 @@ defmodule Scholar.NaiveBayes.Bernoulli do defnp joint_log_likelihood( %__MODULE__{ feature_log_probability: feature_log_probability, - class_log_priors: class_log_priors + class_log_priors: class_log_priors, + binarize: binarize }, x ) do + x = + case binarize do + nil -> x + threshold -> Scholar.Preprocessing.Binarizer.fit_transform(x, threshold: threshold) + end + neg_prob = Nx.log(1 - Nx.exp(feature_log_probability)) jll = Nx.dot(x, [1], feature_log_probability - neg_prob, [1]) jll + class_log_priors + Nx.sum(neg_prob, axes: [1]) diff --git a/lib/scholar/naive_bayes/complement.ex b/lib/scholar/naive_bayes/complement.ex index f2f32a40..31ef6c47 100644 --- a/lib/scholar/naive_bayes/complement.ex +++ b/lib/scholar/naive_bayes/complement.ex @@ -186,7 +186,7 @@ defmodule Scholar.NaiveBayes.Complement do sample_weights = Nx.tensor(sample_weights, type: x_type) {priors, opts} = Keyword.pop(opts, :priors, Nx.tensor(0.0, type: x_type)) - class_priors = Nx.tensor(priors) + class_priors = Nx.tensor(priors, type: x_type) {alpha, opts} = Keyword.pop!(opts, :alpha) alpha = Nx.tensor(alpha, type: x_type) @@ -385,7 +385,8 @@ defmodule Scholar.NaiveBayes.Complement do Nx.log(class_count) - Nx.log(Nx.sum(class_count)) true -> - Nx.broadcast(-Nx.log(num_classes), {num_classes}) + num_classes_t = Nx.tensor(1.0, type: x_type) * num_classes + Nx.broadcast(-Nx.log(num_classes_t), {num_classes}) end %__MODULE__{ diff --git a/lib/scholar/naive_bayes/multinomial.ex b/lib/scholar/naive_bayes/multinomial.ex index 82ec59dd..4ff0d95a 100644 --- a/lib/scholar/naive_bayes/multinomial.ex +++ b/lib/scholar/naive_bayes/multinomial.ex @@ -197,7 +197,7 @@ defmodule Scholar.NaiveBayes.Multinomial do priors_flag = opts[:class_priors] != nil {class_priors, opts} = Keyword.pop(opts, :class_priors, :nan) - class_priors = Nx.tensor(class_priors) + class_priors = Nx.tensor(class_priors, type: type) if priors_flag and Nx.size(class_priors) != num_classes do raise ArgumentError, @@ -271,7 +271,8 @@ defmodule Scholar.NaiveBayes.Multinomial do Nx.log(class_count) - Nx.log(Nx.sum(class_count)) true -> - Nx.broadcast(-Nx.log(num_classes), {num_classes}) + num_classes_t = Nx.tensor(1.0, type: type) * num_classes + Nx.broadcast(-Nx.log(num_classes_t), {num_classes}) end %__MODULE__{ From 54889c12420049c5f1d5c3de3ad9f5bbe17c5447 Mon Sep 17 00:00:00 2001 From: Ricardo Carvalho Santos Date: Sun, 23 Aug 2026 17:40:23 -0300 Subject: [PATCH 5/5] Add regression tests for the NaiveBayes predict and option fixes 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. --- test/scholar/naive_bayes/bernoulli_test.exs | 85 ++++++++++++++++++ test/scholar/naive_bayes/categorical_test.exs | 86 +++++++++++++++++++ 2 files changed, 171 insertions(+) diff --git a/test/scholar/naive_bayes/bernoulli_test.exs b/test/scholar/naive_bayes/bernoulli_test.exs index 77adaad5..aa2a624b 100644 --- a/test/scholar/naive_bayes/bernoulli_test.exs +++ b/test/scholar/naive_bayes/bernoulli_test.exs @@ -169,5 +169,90 @@ defmodule Scholar.NaiveBayes.BernoulliTest do expected_predictions = Nx.tensor([2, 1]) assert predictions == expected_predictions end + + test "applies binarize to the input at predict time, not only at fit time" do + x = Nx.iota({4, 3}) + y = Nx.tensor([1, 2, 0, 2]) + model = Bernoulli.fit(x, y, num_classes: 3) + + x_test = Nx.tensor([[6, 2, 4], [8, 5, 9]]) + + # Every entry of x_test is above the default threshold of 0.0, so both + # rows binarize to all-ones and must therefore score identically. + # Reference: sklearn.naive_bayes.BernoulliNB(binarize=0.0) on the same + # data returns these probabilities for both rows. + assert_all_close( + Bernoulli.predict_probability(model, x_test), + Nx.tensor([ + [0.23000899, 0.11500449, 0.65498656], + [0.23000899, 0.11500449, 0.65498656] + ]) + ) + + # Passing the already-binarized input must give the same answer. + x_test_binarized = Scholar.Preprocessing.Binarizer.fit_transform(x_test, threshold: 0.0) + model_no_binarize = Bernoulli.fit(x, y, num_classes: 3) + + assert_all_close( + Bernoulli.predict_probability(model_no_binarize, x_test_binarized), + Bernoulli.predict_probability(model, x_test) + ) + end + + test "respects a non-default binarize threshold at predict time" do + x = Nx.iota({4, 3}) + y = Nx.tensor([1, 2, 0, 2]) + model = Bernoulli.fit(x, y, num_classes: 3, binarize: 5.0) + + x_test = Nx.tensor([[6, 2, 4], [8, 5, 9]]) + + # With threshold 5.0 the two rows binarize differently ([1,0,0] and + # [1,0,1]), so unlike the default-threshold case they must not score + # identically. + probability = Bernoulli.predict_probability(model, x_test) + refute Nx.to_number(Nx.all_close(probability[0], probability[1])) == 1 + end + end + + describe "option validation" do + test "rejects a negative alpha instead of silently producing NaN" do + x = Nx.tensor([[1, 0, 1], [0, 1, 1], [1, 1, 0]]) + y = Nx.tensor([0, 1, 1]) + + assert_raise NimbleOptions.ValidationError, fn -> + Bernoulli.fit(x, y, num_classes: 2, alpha: -1.0) + end + + assert_raise NimbleOptions.ValidationError, fn -> + Bernoulli.fit(x, y, num_classes: 2, alpha: [1.0, -2.0, 1.0]) + end + end + + test "keeps class_log_priors in the same type as the rest of the model" do + x = Nx.tensor([[1, 0, 1], [0, 1, 1], [1, 1, 0]], type: :f64) + y = Nx.tensor([0, 1, 1]) + + for opts <- [ + [num_classes: 2], + [num_classes: 2, class_priors: [0.3, 0.7]], + [num_classes: 2, fit_priors: false] + ] do + model = Bernoulli.fit(x, y, opts) + + assert Nx.type(model.class_log_priors) == Nx.type(model.feature_log_probability), + "class_log_priors downcast for opts: #{inspect(opts)}" + end + end + + test "computes the uniform prior in the model type, not in f32" do + x = Nx.tensor([[1, 0, 1], [0, 1, 1], [1, 1, 0]], type: :f64) + y = Nx.tensor([0, 1, 1]) + model = Bernoulli.fit(x, y, num_classes: 3, fit_priors: false) + + # -log(3) must be accurate to f64, not an f32 value widened to f64. + assert_all_close(model.class_log_priors[0], Nx.tensor(-:math.log(3), type: :f64), + atol: 1.0e-15 + ) + end end end diff --git a/test/scholar/naive_bayes/categorical_test.exs b/test/scholar/naive_bayes/categorical_test.exs index 179242aa..e46766e4 100644 --- a/test/scholar/naive_bayes/categorical_test.exs +++ b/test/scholar/naive_bayes/categorical_test.exs @@ -1027,4 +1027,90 @@ defmodule Scholar.NaiveBayes.CategoricalTest do ) end end + + describe "predict with num_features != num_classes" do + test "predicts when the feature count differs from the class count" do + # Every doctest and existing test happens to use num_features == + # num_classes, which is the only case where a jll accumulator sized + # from x instead of from the class axis works. + x = Nx.tensor([[1, 2, 2], [1, 2, 1], [2, 2, 0], [0, 1, 2], [2, 0, 1]]) + y = Nx.tensor([0, 1, 1, 0, 1]) + model = Categorical.fit(x, y, num_classes: 2) + + x_test = Nx.tensor([[1, 2, 2], [2, 0, 1], [0, 1, 0]]) + + assert Nx.shape(Categorical.predict_probability(model, x_test)) == {3, 2} + assert Categorical.predict(model, x_test, Nx.tensor([0, 1])) == Nx.tensor([0, 1, 0]) + + # Reference: sklearn.naive_bayes.CategoricalNB on the same data. + assert_all_close( + Categorical.predict_probability(model, x_test), + Nx.tensor([ + [0.6973365617433412, 0.3026634382566586], + [0.060150375939849614, 0.9398496240601504], + [0.6973365617433416, 0.3026634382566585] + ]), + atol: 1.0e-5 + ) + end + + test "check_dim validates against the feature count, not the class count" do + x = Nx.tensor([[1, 2, 2], [1, 2, 1], [2, 2, 0], [0, 1, 2], [2, 0, 1]]) + y = Nx.tensor([0, 1, 1, 0, 1]) + model = Categorical.fit(x, y, num_classes: 2) + + # 3 features, as used for fitting: must be accepted. + assert Nx.shape(Categorical.predict_probability(model, Nx.tensor([[1, 2, 2]]))) == {1, 2} + + # 2 features: must be rejected even though it matches num_classes. + assert_raise ArgumentError, fn -> + Categorical.predict_probability(model, Nx.tensor([[1, 2]])) + end + end + end + + describe "option validation" do + test "rejects a negative alpha instead of silently producing NaN" do + x = Nx.tensor([[1, 0, 1], [0, 1, 1], [1, 1, 0]]) + y = Nx.tensor([0, 1, 1]) + + assert_raise NimbleOptions.ValidationError, fn -> + Categorical.fit(x, y, num_classes: 2, alpha: -1.0) + end + + assert_raise NimbleOptions.ValidationError, fn -> + Categorical.fit(x, y, num_classes: 2, alpha: [1.0, -2.0, 1.0]) + end + end + + test "min_categories actually sizes the category axis" do + x = Nx.tensor([[1, 2, 3], [1, 3, 4], [2, 2, 3], [1, 1, 3], [2, 1, 4]]) + y = Nx.tensor([0, 1, 2, 1, 0]) + + # Without the option the category count is inferred from the data (max 4 + # -> 5 categories). Asking for 7 must widen the axis to 7, matching + # sklearn's CategoricalNB(min_categories=...) semantics. + inferred = Categorical.fit(x, y, num_classes: 3) + assert Nx.axis_size(inferred.feature_count, 2) == 5 + + widened = Categorical.fit(x, y, num_classes: 3, min_categories: [7, 7, 7]) + assert Nx.axis_size(widened.feature_count, 2) == 7 + end + + test "keeps class_log_priors in the same type as the rest of the model" do + x = Nx.tensor([[1, 0, 1], [0, 1, 1], [1, 1, 0]]) + y = Nx.tensor([0, 1, 1]) + + for opts <- [ + [num_classes: 2], + [num_classes: 2, class_priors: [0.3, 0.7]], + [num_classes: 2, fit_priors: false] + ] do + model = Categorical.fit(x, y, opts) + + assert Nx.type(model.class_log_priors) == Nx.type(model.feature_log_probability), + "class_log_priors downcast for opts: #{inspect(opts)}" + end + end + end end