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.

Author: RicardoSantos-99

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

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