Fix NaiveBayes prediction paths and prior dtype

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/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,
@@ -272,10 +277,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 ++
@@ -332,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__{
@@ -360,7 +367,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 +408,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 +460,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
 
@@ -476,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)

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

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