From cc965d5e02f3741e51ed33572adaf68197c676cb Mon Sep 17 00:00:00 2001
From: 0urobor0s <0urobor0s@users.noreply.github.com>
Date: Tue, 10 Oct 2023 15:56:59 +0100
Subject: [PATCH 1/6] Add precision_recall_fscore_support function

---
 lib/scholar/metrics/classification.ex        | 229 +++++++++++++++++--
 lib/scholar/options.ex                       |   8 +
 test/scholar/metrics/classification_test.exs |   4 +-
 3 files changed, 216 insertions(+), 25 deletions(-)

diff --git a/lib/scholar/metrics/classification.ex b/lib/scholar/metrics/classification.ex
index d9bf2f7b..09462ba2 100644
--- a/lib/scholar/metrics/classification.ex
+++ b/lib/scholar/metrics/classification.ex
@@ -46,7 +46,66 @@ defmodule Scholar.Metrics.Classification do
         ]
       ]
 
-  fbeta_score_schema = f1_score_schema
+  fbeta_score_schema =
+    general_schema ++
+      [
+        average: [
+          type: {:in, [:micro, :macro, :weighted, :none]},
+          default: :none,
+          doc: """
+          This determines the type of averaging performed on the data.
+
+          * `:macro` - Calculate metrics for each label, and find their unweighted mean.
+          This does not take label imbalance into account.
+
+          * `:weighted` - Calculate metrics for each label, and find their average weighted by
+          support (the number of true instances for each label).
+
+          * `:micro` - Calculate metrics globally by counting the total true positives,
+          false negatives and false positives.
+
+          * `:none` - The F-score values for each class are returned.
+          """
+        ],
+        beta: [
+          type: {:custom, Scholar.Options, :beta, []},
+          doc: """
+          Determines the weight of recall in the combined score.
+          For values of `beta` > 1 it gives more weight to recall, while `beta` < 1 favors precision.
+          """
+        ]
+      ]
+
+  precision_recall_fscore_support_schema =
+    general_schema ++
+      [
+        average: [
+          type: {:in, [:micro, :macro, :weighted, :none]},
+          default: :none,
+          doc: """
+          This determines the type of averaging performed on the data.
+
+          * `:macro` - Calculate metrics for each label, and find their unweighted mean.
+          This does not take label imbalance into account.
+
+          * `:weighted` - Calculate metrics for each label, and find their average weighted by
+          support (the number of true instances for each label).
+
+          * `:micro` - Calculate metrics globally by counting the total true positives,
+          false negatives and false positives.
+
+          * `:none` - The F-score values for each class are returned.
+          """
+        ],
+        beta: [
+          type: {:custom, Scholar.Options, :beta, []},
+          default: 1,
+          doc: """
+          Determines the weight of recall in the combined score.
+          For values of `beta` > 1 it gives more weight to recall, while `beta` < 1 favors precision.
+          """
+        ]
+      ]
 
   confusion_matrix_schema =
     general_schema ++
@@ -167,6 +226,9 @@ defmodule Scholar.Metrics.Classification do
   @cohen_kappa_schema NimbleOptions.new!(cohen_kappa_schema)
   @fbeta_score_schema NimbleOptions.new!(fbeta_score_schema)
   @f1_score_schema NimbleOptions.new!(f1_score_schema)
+  @precision_recall_fscore_support_schema NimbleOptions.new!(
+                                            precision_recall_fscore_support_schema
+                                          )
   @brier_score_loss_schema NimbleOptions.new!(brier_score_loss_schema)
   @accuracy_schema NimbleOptions.new!(accuracy_schema)
   @top_k_accuracy_score_schema NimbleOptions.new!(top_k_accuracy_score_schema)
@@ -603,58 +665,56 @@ defmodule Scholar.Metrics.Classification do
 
       iex> y_true = Nx.tensor([0, 1, 1, 1, 1, 0, 2, 1, 0, 1], type: :u32)
       iex> y_pred = Nx.tensor([0, 2, 1, 1, 2, 2, 2, 0, 0, 1], type: :u32)
-      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, Nx.u32(1), num_classes: 3)
+      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, beta: Nx.u32(1), num_classes: 3)
       #Nx.Tensor<
         f32[3]
         [0.6666666865348816, 0.6666666865348816, 0.4000000059604645]
       >
-      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, Nx.u32(2), num_classes: 3)
+      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, beta: Nx.u32(2), num_classes: 3)
       #Nx.Tensor<
         f32[3]
         [0.6666666865348816, 0.5555555820465088, 0.625]
       >
-      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, Nx.f32(0.5), num_classes: 3)
+      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, beta: Nx.f32(0.5), num_classes: 3)
       #Nx.Tensor<
         f32[3]
         [0.6666666865348816, 0.8333333134651184, 0.29411765933036804]
       >
-      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, Nx.u32(2), num_classes: 3, average: :macro)
+      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, beta: Nx.u32(2), num_classes: 3, average: :macro)
       #Nx.Tensor<
         f32
         0.6157407760620117
       >
-      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, Nx.u32(2), num_classes: 3, average: :weighted)
+      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, beta: Nx.u32(2), num_classes: 3, average: :weighted)
       #Nx.Tensor<
         f32
         0.5958333611488342
       >
-      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, Nx.f32(0.5), num_classes: 3, average: :micro)
+      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, beta: Nx.f32(0.5), num_classes: 3, average: :micro)
       #Nx.Tensor<
         f32
         0.6000000238418579
       >
-      iex> Scholar.Metrics.Classification.fbeta_score(Nx.tensor([1, 0, 1, 0]), Nx.tensor([0, 1, 0, 1]), Nx.tensor(0.5), num_classes: 2, average: :none)
+      iex> Scholar.Metrics.Classification.fbeta_score(Nx.tensor([1, 0, 1, 0]), Nx.tensor([0, 1, 0, 1]), beta: Nx.tensor(0.5), num_classes: 2, average: :none)
       #Nx.Tensor<
         f32[2]
         [0.0, 0.0]
       >
-      iex> Scholar.Metrics.Classification.fbeta_score(Nx.tensor([1, 0, 1, 0]), Nx.tensor([0, 1, 0, 1]), 0.5, num_classes: 2, average: :none)
+      iex> Scholar.Metrics.Classification.fbeta_score(Nx.tensor([1, 0, 1, 0]), Nx.tensor([0, 1, 0, 1]), beta: 0.5, num_classes: 2, average: :none)
       #Nx.Tensor<
         f32[2]
         [0.0, 0.0]
       >
   """
-  deftransform fbeta_score(y_true, y_pred, beta, opts \\ []) do
-    fbeta_score_n(y_true, y_pred, beta, NimbleOptions.validate!(opts, @fbeta_score_schema))
+  deftransform fbeta_score(y_true, y_pred, opts \\ []) do
+    fbeta_score_n(y_true, y_pred, NimbleOptions.validate!(opts, @fbeta_score_schema))
   end
 
-  defnp fbeta_score_n(y_true, y_pred, beta, opts) do
-    check_shape(y_pred, y_true)
-    num_classes = check_num_classes(opts[:num_classes])
-    average = opts[:average]
+  defnp fbeta_score_n(y_true, y_pred, opts) do
+    check_beta(opts[:beta])
 
-    {_precision, _recall, per_class_fscore} =
-      precision_recall_fscore_n(y_true, y_pred, beta, num_classes, average)
+    {_precision, _recall, per_class_fscore, _support} =
+      precision_recall_fscore_support_n(y_true, y_pred, opts)
 
     per_class_fscore
   end
@@ -677,7 +737,119 @@ defmodule Scholar.Metrics.Classification do
     end
   end
 
-  defnp precision_recall_fscore_n(y_true, y_pred, beta, num_classes, average) do
+  @doc """
+  Calculates precision, recall, F-score and support for each
+  class. It also supports a `beta` argument which weights
+  recall more than precision by it's value.
+
+  ## Options
+
+  #{NimbleOptions.docs(@precision_recall_fscore_support_schema)}
+
+  ## Examples
+
+      iex> y_true = Nx.tensor([0, 1, 1, 1, 1, 0, 2, 1, 0, 1], type: :u32)
+      iex> y_pred = Nx.tensor([0, 2, 1, 1, 2, 2, 2, 0, 0, 1], type: :u32)
+      iex> Scholar.Metrics.Classification.precision_recall_fscore_support(y_true, y_pred, num_classes: 3)
+      {#Nx.Tensor<
+       f32[3]
+       [0.6666666865348816, 1.0, 0.25]
+      >,
+      #Nx.Tensor<
+       f32[3]
+       [0.6666666865348816, 0.5, 1.0]
+      >,
+      #Nx.Tensor<
+       f32[3]
+       [0.6666666865348816, 0.6666666865348816, 0.4000000059604645]
+      >,
+      #Nx.Tensor<
+       u64[3]
+       [3, 6, 1]
+      >}
+      iex> Scholar.Metrics.Classification.precision_recall_fscore_support(y_true, y_pred, num_classes: 3, average: :macro)
+      {#Nx.Tensor<
+       f32[3]
+       [0.6666666865348816, 1.0, 0.25]
+      >,
+      #Nx.Tensor<
+       f32[3]
+       [0.6666666865348816, 0.5, 1.0]
+      >,
+      #Nx.Tensor<
+       f32
+       0.5777778029441833
+      >,
+      #Nx.Tensor<
+       f32
+       NaN
+      >}
+      iex> Scholar.Metrics.Classification.precision_recall_fscore_support(y_true, y_pred, num_classes: 3, average: :weighted)
+      {#Nx.Tensor<
+       f32[3]
+       [0.6666666865348816, 1.0, 0.25]
+      >,
+      #Nx.Tensor<
+       f32[3]
+       [0.6666666865348816, 0.5, 1.0]
+      >,
+      #Nx.Tensor<
+       f32
+       0.6399999856948853
+      >,
+      #Nx.Tensor<
+       f32
+       NaN
+      >}
+      iex> Scholar.Metrics.Classification.precision_recall_fscore_support(y_true, y_pred, num_classes: 3, average: :micro)
+      {#Nx.Tensor<
+       f32
+       0.6000000238418579
+      >,
+      #Nx.Tensor<
+       f32
+       0.6000000238418579
+      >,
+      #Nx.Tensor<
+       f32
+       0.6000000238418579
+      >,
+      #Nx.Tensor<
+       f32
+       NaN
+      >}
+      iex> Scholar.Metrics.Classification.precision_recall_fscore_support(Nx.tensor([1, 0, 1, 0]), Nx.tensor([0, 1, 0, 1]), beta: 2, num_classes: 2, average: :none)
+      {#Nx.Tensor<
+       f32[2]
+       [0.0, 0.0]
+      >,
+      #Nx.Tensor<
+       f32[2]
+       [0.0, 0.0]
+      >,
+      #Nx.Tensor<
+       f32[2]
+       [0.0, 0.0]
+      >,
+      #Nx.Tensor<
+       u64[2]
+       [2, 2]
+      >}
+  """
+  deftransform precision_recall_fscore_support(y_true, y_pred, opts) do
+    precision_recall_fscore_support_n(
+      y_true,
+      y_pred,
+      NimbleOptions.validate!(opts, @precision_recall_fscore_support_schema)
+    )
+  end
+
+  defnp precision_recall_fscore_support_n(y_true, y_pred, opts) do
+    check_shape(y_pred, y_true)
+    num_classes = check_num_classes(opts[:num_classes])
+    beta = opts[:beta]
+    average = opts[:average]
+
     confusion_matrix = confusion_matrix(y_true, y_pred, num_classes: num_classes)
     {true_positive, false_positive, false_negative} = fbeta_score_v(confusion_matrix, average)
 
@@ -700,13 +872,15 @@ defmodule Scholar.Metrics.Classification do
 
     case average do
       :none ->
-        {precision, recall, per_class_fscore}
+        support = (y_true == Nx.iota({num_classes, 1})) |> Nx.sum(axes: [1])
+
+        {precision, recall, per_class_fscore, support}
 
       :micro ->
-        {precision, recall, per_class_fscore}
+        {precision, recall, per_class_fscore, Nx.Constants.nan()}
 
       :macro ->
-        {precision, recall, Nx.mean(per_class_fscore)}
+        {precision, recall, Nx.mean(per_class_fscore), Nx.Constants.nan()}
 
       :weighted ->
         support = (y_true == Nx.iota({num_classes, 1})) |> Nx.sum(axes: [1])
@@ -716,7 +890,7 @@ defmodule Scholar.Metrics.Classification do
           |> safe_division(Nx.sum(support))
           |> Nx.sum()
 
-        {precision, recall, per_class_fscore}
+        {precision, recall, per_class_fscore, Nx.Constants.nan()}
     end
   end
 
@@ -762,7 +936,12 @@ defmodule Scholar.Metrics.Classification do
       >
   """
   deftransform f1_score(y_true, y_pred, opts \\ []) do
-    fbeta_score_n(y_true, y_pred, 1, NimbleOptions.validate!(opts, @f1_score_schema))
+    opts =
+      opts
+      |> NimbleOptions.validate!(@f1_score_schema)
+      |> Keyword.put(:beta, 1)
+
+    fbeta_score_n(y_true, y_pred, opts)
   end
 
   @doc """
@@ -1235,6 +1414,10 @@ defmodule Scholar.Metrics.Classification do
     num_classes || raise ArgumentError, "missing option :num_classes"
   end
 
+  deftransformp check_beta(beta) do
+    beta || raise ArgumentError, "missing option :beta"
+  end
+
   defnp safe_division(nominator, denominator) do
     is_zero? = denominator == 0
     nominator = Nx.select(is_zero?, 0, nominator)
diff --git a/lib/scholar/options.ex b/lib/scholar/options.ex
index a779d572..b4505eb3 100644
--- a/lib/scholar/options.ex
+++ b/lib/scholar/options.ex
@@ -100,4 +100,12 @@ defmodule Scholar.Options do
     {:error,
      "expected metric to be a :cosine or tuple {:minkowski, p} where p is a positive number or :infinity, got: #{inspect(metric)}"}
   end
+
+  def beta(beta) do
+    if (is_number(beta) and beta >= 0) or (Nx.is_tensor(beta) and Nx.size(beta) == 1) do
+      {:ok, beta}
+    else
+      {:error, "expect 'beta' to be in the range [0, inf]"}
+    end
+  end
 end
diff --git a/test/scholar/metrics/classification_test.exs b/test/scholar/metrics/classification_test.exs
index 7667e2df..a1c827ed 100644
--- a/test/scholar/metrics/classification_test.exs
+++ b/test/scholar/metrics/classification_test.exs
@@ -20,7 +20,7 @@ defmodule Scholar.Metrics.ClassificationTest do
       beta = Nx.tensor(:infinity)
       y_true = Nx.tensor([0, 0, 0, 0, 0, 1, 1, 1, 1, 1], type: :u32)
       y_pred = Nx.tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1], type: :u32)
-      fbeta_scores = Classification.fbeta_score(y_true, y_pred, beta, num_classes: 2)
+      fbeta_scores = Classification.fbeta_score(y_true, y_pred, beta: beta, num_classes: 2)
 
       assert_all_close(fbeta_scores, Classification.recall(y_true, y_pred, num_classes: 2))
     end
@@ -29,7 +29,7 @@ defmodule Scholar.Metrics.ClassificationTest do
       beta = 0
       y_true = Nx.tensor([0, 0, 0, 0, 0, 1, 1, 1, 1, 1], type: :u32)
       y_pred = Nx.tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1], type: :u32)
-      fbeta_scores = Classification.fbeta_score(y_true, y_pred, beta, num_classes: 2)
+      fbeta_scores = Classification.fbeta_score(y_true, y_pred, beta: beta, num_classes: 2)
 
       assert_all_close(fbeta_scores, Classification.precision(y_true, y_pred, num_classes: 2))
     end

From 16fd84255480971c18c0bb3202aa7ad6bc93d755 Mon Sep 17 00:00:00 2001
From: 0urobor0s <0urobor0s@users.noreply.github.com>
Date: Thu, 12 Oct 2023 19:09:40 +0100
Subject: [PATCH 2/6] Remove beta from opts in defn context

---
 lib/scholar/metrics/classification.ex | 26 +++++++++++++-------------
 1 file changed, 13 insertions(+), 13 deletions(-)

diff --git a/lib/scholar/metrics/classification.ex b/lib/scholar/metrics/classification.ex
index 09462ba2..025895ec 100644
--- a/lib/scholar/metrics/classification.ex
+++ b/lib/scholar/metrics/classification.ex
@@ -707,14 +707,16 @@ defmodule Scholar.Metrics.Classification do
       >
   """
   deftransform fbeta_score(y_true, y_pred, opts \\ []) do
-    fbeta_score_n(y_true, y_pred, NimbleOptions.validate!(opts, @fbeta_score_schema))
+    opts = NimbleOptions.validate!(opts, @fbeta_score_schema)
+    {beta, opts} = Keyword.pop(opts, :beta)
+    fbeta_score_n(y_true, y_pred, beta, opts)
   end
 
-  defnp fbeta_score_n(y_true, y_pred, opts) do
-    check_beta(opts[:beta])
+  defnp fbeta_score_n(y_true, y_pred, beta, opts) do
+    check_beta(beta)
 
     {_precision, _recall, per_class_fscore, _support} =
-      precision_recall_fscore_support_n(y_true, y_pred, opts)
+      precision_recall_fscore_support_n(y_true, y_pred, beta, opts)
 
     per_class_fscore
   end
@@ -837,17 +839,20 @@ defmodule Scholar.Metrics.Classification do
       >}
   """
   deftransform precision_recall_fscore_support(y_true, y_pred, opts) do
+    opts = NimbleOptions.validate!(opts, @precision_recall_fscore_support_schema)
+    {beta, opts} = Keyword.pop(opts, :beta)
+
     precision_recall_fscore_support_n(
       y_true,
       y_pred,
-      NimbleOptions.validate!(opts, @precision_recall_fscore_support_schema)
+      beta,
+      opts
     )
   end
 
-  defnp precision_recall_fscore_support_n(y_true, y_pred, opts) do
+  defnp precision_recall_fscore_support_n(y_true, y_pred, beta, opts) do
     check_shape(y_pred, y_true)
     num_classes = check_num_classes(opts[:num_classes])
-    beta = opts[:beta]
     average = opts[:average]
 
     confusion_matrix = confusion_matrix(y_true, y_pred, num_classes: num_classes)
@@ -936,12 +941,7 @@ defmodule Scholar.Metrics.Classification do
       >
   """
   deftransform f1_score(y_true, y_pred, opts \\ []) do
-    opts =
-      opts
-      |> NimbleOptions.validate!(@f1_score_schema)
-      |> Keyword.put(:beta, 1)
-
-    fbeta_score_n(y_true, y_pred, opts)
+    fbeta_score_n(y_true, y_pred, 1, opts)
   end
 
   @doc """

From 9271b3bc7b2363426acf52ce8f224584ac83d36f Mon Sep 17 00:00:00 2001
From: 0urobor0s <0urobor0s@users.noreply.github.com>
Date: Thu, 12 Oct 2023 22:11:20 +0100
Subject: [PATCH 3/6] Fix opts and improve docs

---
 lib/scholar/metrics/classification.ex | 108 +++++++-------------------
 1 file changed, 26 insertions(+), 82 deletions(-)

diff --git a/lib/scholar/metrics/classification.ex b/lib/scholar/metrics/classification.ex
index 025895ec..4f7ce7cf 100644
--- a/lib/scholar/metrics/classification.ex
+++ b/lib/scholar/metrics/classification.ex
@@ -753,90 +753,34 @@ defmodule Scholar.Metrics.Classification do
       iex> y_true = Nx.tensor([0, 1, 1, 1, 1, 0, 2, 1, 0, 1], type: :u32)
       iex> y_pred = Nx.tensor([0, 2, 1, 1, 2, 2, 2, 0, 0, 1], type: :u32)
       iex> Scholar.Metrics.Classification.precision_recall_fscore_support(y_true, y_pred, num_classes: 3)
-      {#Nx.Tensor<
-       f32[3]
-       [0.6666666865348816, 1.0, 0.25]
-      >,
-      #Nx.Tensor<
-       f32[3]
-       [0.6666666865348816, 0.5, 1.0]
-      >,
-      #Nx.Tensor<
-       f32[3]
-       [0.6666666865348816, 0.6666666865348816, 0.4000000059604645]
-      >,
-      #Nx.Tensor<
-       u64[3]
-       [3, 6, 1]
-      >}
+      {Nx.f32([0.6666666865348816, 1.0, 0.25]),
+       Nx.f32([0.6666666865348816, 0.5, 1.0]),
+       Nx.f32([0.6666666865348816, 0.6666666865348816, 0.4000000059604645]),
+       Nx.u64([3, 6, 1])}
       iex> Scholar.Metrics.Classification.precision_recall_fscore_support(y_true, y_pred, num_classes: 3, average: :macro)
-      {#Nx.Tensor<
-       f32[3]
-       [0.6666666865348816, 1.0, 0.25]
-      >,
-      #Nx.Tensor<
-       f32[3]
-       [0.6666666865348816, 0.5, 1.0]
-      >,
-      #Nx.Tensor<
-       f32
-       0.5777778029441833
-      >,
-      #Nx.Tensor<
-       f32
-       NaN
-      >}
+      {Nx.f32([0.6666666865348816, 1.0, 0.25]),
+       Nx.f32([0.6666666865348816, 0.5, 1.0]),
+       Nx.f32(0.5777778029441833),
+       Nx.Constants.nan()}
       iex> Scholar.Metrics.Classification.precision_recall_fscore_support(y_true, y_pred, num_classes: 3, average: :weighted)
-      {#Nx.Tensor<
-       f32[3]
-       [0.6666666865348816, 1.0, 0.25]
-      >,
-      #Nx.Tensor<
-       f32[3]
-       [0.6666666865348816, 0.5, 1.0]
-      >,
-      #Nx.Tensor<
-       f32
-       0.6399999856948853
-      >,
-      #Nx.Tensor<
-       f32
-       NaN
-      >}
+      {Nx.f32([0.6666666865348816, 1.0, 0.25]),
+       Nx.f32([0.6666666865348816, 0.5, 1.0]),
+       Nx.f32(0.6399999856948853),
+       Nx.Constants.nan()}
       iex> Scholar.Metrics.Classification.precision_recall_fscore_support(y_true, y_pred, num_classes: 3, average: :micro)
-      {#Nx.Tensor<
-       f32
-       0.6000000238418579
-      >,
-      #Nx.Tensor<
-       f32
-       0.6000000238418579
-      >,
-      #Nx.Tensor<
-       f32
-       0.6000000238418579
-      >,
-      #Nx.Tensor<
-       f32
-       NaN
-      >}
-      iex> Scholar.Metrics.Classification.precision_recall_fscore_support(Nx.tensor([1, 0, 1, 0]), Nx.tensor([0, 1, 0, 1]), beta: 2, num_classes: 2, average: :none)
-      {#Nx.Tensor<
-       f32[2]
-       [0.0, 0.0]
-      >,
-      #Nx.Tensor<
-       f32[2]
-       [0.0, 0.0]
-      >,
-      #Nx.Tensor<
-       f32[2]
-       [0.0, 0.0]
-      >,
-      #Nx.Tensor<
-       u64[2]
-       [2, 2]
-      >}
+      {Nx.f32(0.6000000238418579),
+       Nx.f32(0.6000000238418579),
+       Nx.f32(0.6000000238418579),
+       Nx.Constants.nan()}
+
+      iex> y_true = Nx.tensor([1, 0, 1, 0], type: :u32)
+      iex> y_pred = Nx.tensor([0, 1, 0, 1], type: :u32)
+      iex> opts = [beta: 2, num_classes: 2, average: :none]
+      iex> Scholar.Metrics.Classification.precision_recall_fscore_support(y_true, y_pred, opts)
+      {Nx.f32([0.0, 0.0]),
+       Nx.f32([0.0, 0.0]),
+       Nx.f32([0.0, 0.0]),
+       Nx.u64([2, 2])}
   """
   deftransform precision_recall_fscore_support(y_true, y_pred, opts) do
     opts = NimbleOptions.validate!(opts, @precision_recall_fscore_support_schema)
@@ -941,7 +885,7 @@ defmodule Scholar.Metrics.Classification do
       >
   """
   deftransform f1_score(y_true, y_pred, opts \\ []) do
-    fbeta_score_n(y_true, y_pred, 1, opts)
+    fbeta_score_n(y_true, y_pred, 1, NimbleOptions.validate!(opts, @f1_score_schema))
   end
 
   @doc """

From f40789a3df8716f3c87f360fe9ed8e3af86c8a84 Mon Sep 17 00:00:00 2001
From: 0urobor0s <0urobor0s@users.noreply.github.com>
Date: Thu, 12 Oct 2023 22:41:51 +0100
Subject: [PATCH 4/6] Require beta for fbeta_score

---
 lib/scholar/metrics/classification.ex | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/lib/scholar/metrics/classification.ex b/lib/scholar/metrics/classification.ex
index 4f7ce7cf..e6eb9366 100644
--- a/lib/scholar/metrics/classification.ex
+++ b/lib/scholar/metrics/classification.ex
@@ -69,6 +69,7 @@ defmodule Scholar.Metrics.Classification do
         ],
         beta: [
           type: {:custom, Scholar.Options, :beta, []},
+          required: true,
           doc: """
           Determines the weight of recall in the combined score.
           For values of `beta` > 1 it gives more weight to recall, while `beta` < 1 favors precision.
@@ -713,8 +714,6 @@ defmodule Scholar.Metrics.Classification do
   end
 
   defnp fbeta_score_n(y_true, y_pred, beta, opts) do
-    check_beta(beta)
-
     {_precision, _recall, per_class_fscore, _support} =
       precision_recall_fscore_support_n(y_true, y_pred, beta, opts)
 
@@ -1358,10 +1357,6 @@ defmodule Scholar.Metrics.Classification do
     num_classes || raise ArgumentError, "missing option :num_classes"
   end
 
-  deftransformp check_beta(beta) do
-    beta || raise ArgumentError, "missing option :beta"
-  end
-
   defnp safe_division(nominator, denominator) do
     is_zero? = denominator == 0
     nominator = Nx.select(is_zero?, 0, nominator)

From 92710d66c2db5e2b1f925a87e56dd9f9b3699e7c Mon Sep 17 00:00:00 2001
From: 0urobor0s <0urobor0s@users.noreply.github.com>
Date: Mon, 16 Oct 2023 17:06:24 +0100
Subject: [PATCH 5/6] Have beta as function argument

---
 lib/scholar/metrics/classification.ex        | 53 ++++----------------
 test/scholar/metrics/classification_test.exs |  4 +-
 2 files changed, 13 insertions(+), 44 deletions(-)

diff --git a/lib/scholar/metrics/classification.ex b/lib/scholar/metrics/classification.ex
index e6eb9366..373c1b45 100644
--- a/lib/scholar/metrics/classification.ex
+++ b/lib/scholar/metrics/classification.ex
@@ -46,36 +46,7 @@ defmodule Scholar.Metrics.Classification do
         ]
       ]
 
-  fbeta_score_schema =
-    general_schema ++
-      [
-        average: [
-          type: {:in, [:micro, :macro, :weighted, :none]},
-          default: :none,
-          doc: """
-          This determines the type of averaging performed on the data.
-
-          * `:macro` - Calculate metrics for each label, and find their unweighted mean.
-          This does not take label imbalance into account.
-
-          * `:weighted` - Calculate metrics for each label, and find their average weighted by
-          support (the number of true instances for each label).
-
-          * `:micro` - Calculate metrics globally by counting the total true positives,
-          false negatives and false positives.
-
-          * `:none` - The F-score values for each class are returned.
-          """
-        ],
-        beta: [
-          type: {:custom, Scholar.Options, :beta, []},
-          required: true,
-          doc: """
-          Determines the weight of recall in the combined score.
-          For values of `beta` > 1 it gives more weight to recall, while `beta` < 1 favors precision.
-          """
-        ]
-      ]
+  fbeta_score_schema = f1_score_schema
 
   precision_recall_fscore_support_schema =
     general_schema ++
@@ -666,51 +637,49 @@ defmodule Scholar.Metrics.Classification do
 
       iex> y_true = Nx.tensor([0, 1, 1, 1, 1, 0, 2, 1, 0, 1], type: :u32)
       iex> y_pred = Nx.tensor([0, 2, 1, 1, 2, 2, 2, 0, 0, 1], type: :u32)
-      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, beta: Nx.u32(1), num_classes: 3)
+      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, Nx.u32(1), num_classes: 3)
       #Nx.Tensor<
         f32[3]
         [0.6666666865348816, 0.6666666865348816, 0.4000000059604645]
       >
-      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, beta: Nx.u32(2), num_classes: 3)
+      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, Nx.u32(2), num_classes: 3)
       #Nx.Tensor<
         f32[3]
         [0.6666666865348816, 0.5555555820465088, 0.625]
       >
-      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, beta: Nx.f32(0.5), num_classes: 3)
+      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, Nx.f32(0.5), num_classes: 3)
       #Nx.Tensor<
         f32[3]
         [0.6666666865348816, 0.8333333134651184, 0.29411765933036804]
       >
-      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, beta: Nx.u32(2), num_classes: 3, average: :macro)
+      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, Nx.u32(2), num_classes: 3, average: :macro)
       #Nx.Tensor<
         f32
         0.6157407760620117
       >
-      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, beta: Nx.u32(2), num_classes: 3, average: :weighted)
+      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, Nx.u32(2), num_classes: 3, average: :weighted)
       #Nx.Tensor<
         f32
         0.5958333611488342
       >
-      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, beta: Nx.f32(0.5), num_classes: 3, average: :micro)
+      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, Nx.f32(0.5), num_classes: 3, average: :micro)
       #Nx.Tensor<
         f32
         0.6000000238418579
       >
-      iex> Scholar.Metrics.Classification.fbeta_score(Nx.tensor([1, 0, 1, 0]), Nx.tensor([0, 1, 0, 1]), beta: Nx.tensor(0.5), num_classes: 2, average: :none)
+      iex> Scholar.Metrics.Classification.fbeta_score(Nx.tensor([1, 0, 1, 0]), Nx.tensor([0, 1, 0, 1]), Nx.tensor(0.5), num_classes: 2, average: :none)
       #Nx.Tensor<
         f32[2]
         [0.0, 0.0]
       >
-      iex> Scholar.Metrics.Classification.fbeta_score(Nx.tensor([1, 0, 1, 0]), Nx.tensor([0, 1, 0, 1]), beta: 0.5, num_classes: 2, average: :none)
+      iex> Scholar.Metrics.Classification.fbeta_score(Nx.tensor([1, 0, 1, 0]), Nx.tensor([0, 1, 0, 1]), 0.5, num_classes: 2, average: :none)
       #Nx.Tensor<
         f32[2]
         [0.0, 0.0]
       >
   """
-  deftransform fbeta_score(y_true, y_pred, opts \\ []) do
-    opts = NimbleOptions.validate!(opts, @fbeta_score_schema)
-    {beta, opts} = Keyword.pop(opts, :beta)
-    fbeta_score_n(y_true, y_pred, beta, opts)
+  deftransform fbeta_score(y_true, y_pred, beta, opts \\ []) do
+    fbeta_score_n(y_true, y_pred, beta, NimbleOptions.validate!(opts, @fbeta_score_schema))
   end
 
   defnp fbeta_score_n(y_true, y_pred, beta, opts) do
diff --git a/test/scholar/metrics/classification_test.exs b/test/scholar/metrics/classification_test.exs
index a1c827ed..7667e2df 100644
--- a/test/scholar/metrics/classification_test.exs
+++ b/test/scholar/metrics/classification_test.exs
@@ -20,7 +20,7 @@ defmodule Scholar.Metrics.ClassificationTest do
       beta = Nx.tensor(:infinity)
       y_true = Nx.tensor([0, 0, 0, 0, 0, 1, 1, 1, 1, 1], type: :u32)
       y_pred = Nx.tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1], type: :u32)
-      fbeta_scores = Classification.fbeta_score(y_true, y_pred, beta: beta, num_classes: 2)
+      fbeta_scores = Classification.fbeta_score(y_true, y_pred, beta, num_classes: 2)
 
       assert_all_close(fbeta_scores, Classification.recall(y_true, y_pred, num_classes: 2))
     end
@@ -29,7 +29,7 @@ defmodule Scholar.Metrics.ClassificationTest do
       beta = 0
       y_true = Nx.tensor([0, 0, 0, 0, 0, 1, 1, 1, 1, 1], type: :u32)
       y_pred = Nx.tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1], type: :u32)
-      fbeta_scores = Classification.fbeta_score(y_true, y_pred, beta: beta, num_classes: 2)
+      fbeta_scores = Classification.fbeta_score(y_true, y_pred, beta, num_classes: 2)
 
       assert_all_close(fbeta_scores, Classification.precision(y_true, y_pred, num_classes: 2))
     end

From 88582f2325ea8b5c9a18746b80f73b3001e4364a Mon Sep 17 00:00:00 2001
From: Daniel Tinoco <0urobor0s@users.noreply.github.com>
Date: Mon, 16 Oct 2023 23:57:28 +0100
Subject: [PATCH 6/6] Use rank instead of size

Co-authored-by: Mateusz Sluszniak <56299341+msluszniak@users.noreply.github.com>
---
 lib/scholar/options.ex | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/scholar/options.ex b/lib/scholar/options.ex
index b4505eb3..e0173ad1 100644
--- a/lib/scholar/options.ex
+++ b/lib/scholar/options.ex
@@ -102,7 +102,7 @@ defmodule Scholar.Options do
   end
 
   def beta(beta) do
-    if (is_number(beta) and beta >= 0) or (Nx.is_tensor(beta) and Nx.size(beta) == 1) do
+    if (is_number(beta) and beta >= 0) or (Nx.is_tensor(beta) and Nx.rank(beta) == 0) do
       {:ok, beta}
     else
       {:error, "expect 'beta' to be in the range [0, inf]"}