add utility functions to NX

nx/lib/nx.ex

@@ -9167,6 +9167,70 @@ defmodule Nx do
     end
   end
 
+  @doc """
+  Finds the variance of a tensor.
+
+  The variance is the average of the squared deviations from the mean.
+  The mean is typically calculated as sum(x) / N, where N = total of elements.
+  If, however, ddof is specified, the divisor N - ddof is used instead.
+
+  ## Examples
+
+      iex> Nx.variance(Nx.tensor([[1, 2], [3, 4]]))
+      #Nx.Tensor<
+        f32
+        1.25
+      >
+
+      iex> Nx.variance(Nx.tensor([[1, 2], [3, 4]]), ddof: 1)
+      #Nx.Tensor<
+        f32
+        1.6666666269302368
+      >
+  """
+  @doc type: :aggregation
+  @spec variance(tensor :: Nx.Tensor.t(), opts :: Keyword.t()) :: Nx.Tensor.t()
+  def variance(tensor, opts \\ []) do
+    %T{shape: shape} = tensor = to_tensor(tensor)
+
+    total = Tuple.product(shape)
+    ddof = Keyword.get(opts, :ddof, 0)
+    mean = mean(tensor)
+
+    tensor
+    |> subtract(mean)
+    |> power(2)
+    |> sum()
+    |> divide(total - ddof)
+  end
+
+  @doc """
+  Finds the standard deviation of a tensor.
+
+  if ddof is specified, the divisor N - ddof is used to calculate the variance.
+
+  ## Examples
+
+      iex> Nx.standard_deviation(Nx.tensor([[1, 2], [3, 4]]))
+      #Nx.Tensor<
+        f32
+        1.1180340051651
+      >
+
+      iex> Nx.standard_deviation(Nx.tensor([[1, 2], [3, 4]]), ddof: 1)
+      #Nx.Tensor<
+        f32
+        1.29099440574646
+      >
+  """
+  @doc type: :aggregation
+  @spec standard_deviation(tensor :: Nx.Tensor.t(), ddof :: Keyword.t()) :: Nx.Tensor.t()
+  def standard_deviation(tensor, opts \\ []) do
+    tensor = to_tensor(tensor)
+
+    sqrt(variance(tensor, opts))
+  end
+
   ## Sigils
 
   @doc """

nx/test/nx_test.exs

@@ -1920,4 +1920,28 @@ defmodule NxTest do
       |> elem(0)
     end
   end
+
+  describe "variance/1" do
+    test "should calculate the variance of a tensor" do
+      t = Nx.tensor([[4, 5], [2, 3], [1, 0]])
+      assert Nx.variance(t) == Nx.tensor(2.9166667461395264)
+    end
+
+    test "should use the optional ddof" do
+      t = Nx.tensor([[4, 5], [2, 3], [1, 0]])
+      assert Nx.variance(t, ddof: 1) == Nx.tensor(3.5)
+    end
+  end
+
+  describe "standard_deviation/1" do
+    test "should calculate the standard deviation of a tensor" do
+      t = Nx.tensor([[4, 5], [2, 3], [1, 0]])
+      assert Nx.standard_deviation(t) == Nx.tensor(1.707825127659933)
+    end
+
+    test "should use the optional ddof" do
+      t = Nx.tensor([[4, 5], [2, 3], [1, 0]])
+      assert Nx.standard_deviation(t, ddof: 1) == Nx.tensor(1.8708287477493286)
+    end
+  end
 end