Standard Scaler fit-transform interface

lib/scholar/scaler/standard_scaler.ex

@@ -0,0 +1,39 @@
+defmodule Scholar.Scaler.StandardScaler do
+  import Nx.Defn
+
+  defstruct [:deviation, :mean]
+
+  opts_schema = [
+    axes: [
+      type: {:custom, Scholar.Options, :axes, []},
+      doc: """
+      Axes to calculate the distance over. By default the distance
+      is calculated between the whole tensors.
+      """
+    ]
+  ]
+
+  @opts_schema NimbleOptions.new!(opts_schema)
+
+  deftransform fit(tensor, opts \\ []) do
+    NimbleOptions.validate!(opts, @opts_schema)
+
+    std = Nx.standard_deviation(tensor, axes: opts[:axes], keep_axes: true)
+    mean_reduced = Nx.mean(tensor, axes: opts[:axes], keep_axes: true)
+    mean_reduced = Nx.select(Nx.equal(std, 0), 0.0, mean_reduced)
+    %__MODULE__{deviation: std, mean: mean_reduced}
+  end
+
+  deftransform transform(tensor, %__MODULE__{deviation: std, mean: mean}) do
+    scale(tensor, std, mean)
+  end
+
+  deftransform fit_transform(tensor, opts \\ []) do
+    scaler = __MODULE__.fit(tensor, opts)
+    __MODULE__.transform(tensor, scaler)
+  end
+
+  defnp scale(tensor, std, mean) do
+    (tensor - mean) / Nx.select(std == 0, 1.0, std)
+  end
+end

test/scholar/scaler/standard_scaler_test.exs

@@ -0,0 +1,25 @@
+defmodule StandardScalerTest do
+  use Scholar.Case, async: true
+  alias Scholar.Scaler.StandardScaler
+
+  describe "fit_transform/2" do
+    test "applies standard scaling to data" do
+      data = Nx.tensor([[1, -1, 2], [2, 0, 0], [0, 1, -1]])
+
+      expected =
+        Nx.tensor([
+          [0.5212860703468323, -1.3553436994552612, 1.4596009254455566],
+          [1.4596009254455566, -0.4170288145542145, -0.4170288145542145],
+          [-0.4170288145542145, 0.5212860703468323, -1.3553436994552612]
+        ])
+
+      assert_all_close(StandardScaler.fit_transform(data), expected)
+    end
+
+    test "leaves data as it is when variance is zero" do
+      data = 42.0
+      expected = Nx.tensor(data)
+      assert StandardScaler.fit_transform(data) == expected
+    end
+  end
+end