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.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Author: RicardoSantos-99

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])

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)