From d9dda6b3c1c13eb29945001299b7b2f22ff8be12 Mon Sep 17 00:00:00 2001 From: Ricardo Carvalho Santos Date: Fri, 28 Aug 2026 09:22:58 -0300 Subject: [PATCH] Fix confusion_matrix raising on tensor sample weights validate_weights reads :type from opts, which is optional. The nil and list branches work without it, but the tensor branch calls Nx.as_type with a nil type, which raises. confusion_matrix is the only caller that passes no type, so it is the only one that hits this. Convert only when a type is given, so the tensor branch behaves like the list branch instead. --- lib/scholar/shared.ex | 4 ++- test/scholar/metrics/classification_test.exs | 34 ++++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/lib/scholar/shared.ex b/lib/scholar/shared.ex index d7a237e0..f36a5e53 100644 --- a/lib/scholar/shared.ex +++ b/lib/scholar/shared.ex @@ -59,7 +59,9 @@ defmodule Scholar.Shared do Nx.tensor(weights, type: type) Nx.is_tensor(weights) and Nx.shape(weights) in [{}, {num_samples}] -> - weights |> Nx.broadcast({num_samples}) |> Nx.as_type(type) + weights = Nx.broadcast(weights, {num_samples}) + # :type is optional, and the branches above already work without it + if type, do: Nx.as_type(weights, type), else: weights true -> raise ArgumentError, diff --git a/test/scholar/metrics/classification_test.exs b/test/scholar/metrics/classification_test.exs index 975abed0..5ba7946d 100644 --- a/test/scholar/metrics/classification_test.exs +++ b/test/scholar/metrics/classification_test.exs @@ -4,6 +4,40 @@ defmodule Scholar.Metrics.ClassificationTest do alias Scholar.Metrics.Classification doctest Classification + test "confusion_matrix - sample weights as a tensor match the same weights as a list" do + y_true = Nx.tensor([0, 0, 1, 1], type: :u32) + y_pred = Nx.tensor([0, 1, 1, 1], type: :u32) + + from_list = + Classification.confusion_matrix(y_true, y_pred, + num_classes: 2, + sample_weights: [1, 2, 1, 1] + ) + + from_tensor = + Classification.confusion_matrix(y_true, y_pred, + num_classes: 2, + sample_weights: Nx.tensor([1, 2, 1, 1]) + ) + + assert from_tensor == from_list + end + + test "confusion_matrix - a scalar sample weight applies to every sample" do + y_true = Nx.tensor([0, 0, 1, 1], type: :u32) + y_pred = Nx.tensor([0, 1, 1, 1], type: :u32) + + weighted = + Classification.confusion_matrix(y_true, y_pred, + num_classes: 2, + sample_weights: Nx.tensor(2) + ) + + unweighted = Classification.confusion_matrix(y_true, y_pred, num_classes: 2) + + assert Nx.to_flat_list(weighted) == Nx.to_flat_list(Nx.multiply(unweighted, 2)) + end + test "roc_curve - y_score with repeated elements" do y_score = Nx.tensor([0.1, 0.1, 0.2, 0.2, 0.3, 0.3]) y_true = Nx.tensor([0, 0, 1, 1, 1, 1])