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.

Author: RicardoSantos-99

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

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__{

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__{