Fix Scholar.Cluster.DBSCAN :weights

lib/scholar/cluster/dbscan.ex

@@ -88,13 +88,15 @@ defmodule Scholar.Cluster.DBSCAN do
       }
   """
   deftransform fit(x, opts \\ []) do
-    fit_n(x, NimbleOptions.validate!(opts, @opts_schema))
+    opts = NimbleOptions.validate!(opts, @opts_schema)
+    {weights, opts} = Keyword.pop(opts, :weights, nil)
+    weights = validate_weights(weights, Nx.axis_size(x, 0), type: to_float_type(x))
+    fit_n(x, weights, opts)
   end
 
-  defnp fit_n(x, opts) do
+  defnp fit_n(x, weights, opts) do
     num_samples = Nx.axis_size(x, 0)
-    weights = validate_weights(opts[:weights], num_samples, type: to_float_type(x))
-    weights = if Nx.rank(weights) == 0, do: weights, else: Nx.new_axis(weights, -1)
+    weights = if Nx.rank(weights) == 0, do: weights, else: Nx.new_axis(weights, 0)
     y_dummy = Nx.broadcast(Nx.tensor(0), {num_samples})
 
     neighbor_model =

test/scholar/cluster/dbscan_test.exs

@@ -40,6 +40,50 @@ defmodule Scholar.Cluster.DBSCANTest do
       assert model.core_sample_indices == Nx.tensor([0, 1, 0, 1, 0, 1, 0, 1, 0, 0], type: :u8)
     end
 
+    test "fit with weights" do
+      # Expected values from scikit-learn 1.6.1 on this data.
+      x = Nx.tensor([[1, 2], [2, 2], [2, 3], [8, 7], [8, 8], [25, 80]])
+      weights = Nx.tensor([5.0, 0.1, 0.1, 1.0, 1.0, 1.0])
+      model = DBSCAN.fit(x, eps: 3, min_samples: 2, weights: weights)
+
+      assert model.labels == Nx.tensor([0, 0, 0, 1, 1, -1])
+      assert model.core_sample_indices == Nx.tensor([1, 1, 1, 1, 1, 0], type: :u8)
+    end
+
+    test "fit with weights joining two clusters through a light sample" do
+      # Expected values from scikit-learn 1.6.1 on this data.
+      x = Nx.tensor([[0, 0], [1, 0], [2, 0]])
+      weights = Nx.tensor([5.0, 0.1, 5.0])
+      model = DBSCAN.fit(x, eps: 1.5, min_samples: 3, weights: weights)
+
+      assert model.labels == Nx.tensor([0, 0, 0])
+      assert model.core_sample_indices == Nx.tensor([1, 1, 1], type: :u8)
+    end
+
+    test "fit with negative weights" do
+      # Expected values from scikit-learn 1.6.1 on this data.
+      x = Nx.tensor([[0, 0], [1, 0], [2, 0]])
+      weights = Nx.tensor([3.0, -1.0, 3.0])
+      model = DBSCAN.fit(x, eps: 1.5, min_samples: 3, weights: weights)
+
+      assert model.labels == Nx.tensor([0, 0, 0])
+      assert model.core_sample_indices == Nx.tensor([0, 1, 0], type: :u8)
+    end
+
+    test "works with jit_apply" do
+      opts = [
+        eps: 3.3,
+        min_samples: 2,
+        weights: Nx.tensor([5.0, 0.1, 0.1, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0])
+      ]
+
+      model = DBSCAN.fit(x(), opts)
+      jitted = Nx.Defn.jit_apply(&DBSCAN.fit(&1, opts), [x()])
+
+      assert jitted.labels == model.labels
+      assert jitted.core_sample_indices == model.core_sample_indices
+    end
+
     test "test with artificially created 3 blobs" do
       x =
         Nx.tensor([