feat: use vectorization for overlap_and_add (#11)

* feat: use vectorization for overlap_and_add

* fix: mix.exs version

* fix: ci elixir version

* refactor: use Nx.revectorize

* fix: support type opt

* docs

Author: polvalente

.github/workflows/ci.yml

@@ -11,7 +11,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        elixir: ["1.13.0"]
+        elixir: ["1.14.0"]
         otp: ["24.0"]
     env:
       MIX_ENV: test
@@ -42,7 +42,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        elixir: ["1.13.0"]
+        elixir: ["1.14.0"]
         otp: ["24.0"]
     env:
       MIX_ENV: test

lib/nx_signal.ex

@@ -718,7 +718,7 @@ defmodule NxSignal do
 
   @doc """
   Performs the overlap-and-add algorithm over
-  an M by N tensor, where M is the number of
+  an {..., M, N}-shaped tensor, where M is the number of
   windows and N is the window size.
 
   The tensor is zero-padded on the right so
@@ -736,60 +736,80 @@ defmodule NxSignal do
         s64[12]
         [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
       >
+
       iex> NxSignal.overlap_and_add(Nx.iota({3, 4}), overlap_length: 3)
       #Nx.Tensor<
         s64[6]
         [0, 5, 15, 18, 17, 11]
       >
+
+      iex> t = Nx.tensor([[[[0, 1, 2, 3], [4, 5, 6, 7]]], [[[10, 11, 12, 13], [14, 15, 16, 17]]]]) |> Nx.vectorize(x: 2, y: 1)
+      iex> NxSignal.overlap_and_add(t, overlap_length: 3)
+      #Nx.Tensor<
+        vectorized[x: 2][y: 1]
+        s64[5]
+        [
+          [
+            [0, 5, 7, 9, 7]
+          ],
+          [
+            [10, 25, 27, 29, 17]
+          ]
+        ]
+      >
   """
   @doc type: :windowing
   defn overlap_and_add(tensor, opts \\ []) do
-    opts = keyword!(opts, [:overlap_length])
-
-    {num_windows, window_length} = Nx.shape(tensor)
+    opts = keyword!(opts, [:overlap_length, type: Nx.type(tensor)])
     overlap_length = opts[:overlap_length]
 
+    %{vectorized_axes: vectorized_axes, shape: input_shape} = tensor
+    num_windows = Nx.axis_size(tensor, -2)
+    window_length = Nx.axis_size(tensor, -1)
+
     if overlap_length >= window_length do
       raise ArgumentError,
             "overlap_length must be a number less than the window size #{window_length}, got: #{inspect(window_length)}"
     end
 
+    tensor =
+      Nx.revectorize(tensor, [condensed_vectors: :auto, windows: num_windows],
+        target_shape: {window_length}
+      )
+
     stride = window_length - overlap_length
     output_holder_shape = {num_windows * stride + overlap_length}
 
-    {output, _, _, _, _, _} =
-      while {
-              out =
-                Nx.broadcast(
-                  Nx.tensor(0, type: tensor.type),
-                  output_holder_shape
-                ),
-              tensor,
-              i = 0,
-              idx_template = Nx.iota({window_length, 1}),
-              stride,
-              num_windows
-            },
-            i < num_windows do
-        current_window = tensor[i]
-        idx = idx_template + i * stride
-
-        {
-          Nx.indexed_add(out, idx, current_window),
-          tensor,
-          i + 1,
-          idx_template,
-          stride,
-          num_windows
-        }
-      end
+    out =
+      Nx.broadcast(
+        Nx.tensor(0, type: tensor.type),
+        output_holder_shape
+      )
 
-    case opts[:type] do
-      nil ->
-        output
+    idx_template = Nx.iota({window_length, 1}, vectorized_axes: [windows: 1])
+    i = Nx.iota({num_windows}) |> Nx.vectorize(:windows)
+    idx = idx_template + i * stride
 
-      t ->
-        Nx.as_type(output, t)
-    end
+    [%{vectorized_axes: [condensed_vectors: n, windows: _]} = tensor, idx] =
+      Nx.broadcast_vectors([tensor, idx])
+
+    tensor = Nx.revectorize(tensor, [condensed_vectors: n], target_shape: {:auto})
+    idx = Nx.revectorize(idx, [condensed_vectors: n], target_shape: {:auto, 1})
+
+    out_shape = overlap_and_add_output_shape(out.shape, input_shape)
+
+    out
+    |> Nx.indexed_add(idx, tensor)
+    |> Nx.as_type(opts[:type])
+    |> Nx.revectorize(vectorized_axes, target_shape: out_shape)
+  end
+
+  deftransformp overlap_and_add_output_shape({out_len}, in_shape) do
+    idx = tuple_size(in_shape) - 2
+
+    in_shape
+    |> Tuple.delete_at(idx)
+    |> Tuple.delete_at(idx)
+    |> Tuple.append(out_len)
   end
 end

mix.exs

@@ -8,7 +8,7 @@ defmodule NxSignal.MixProject do
     [
       app: :nx_signal,
       version: @version,
-      elixir: "~> 1.13",
+      elixir: "~> 1.14",
       start_permanent: Mix.env() == :prod,
       elixirc_paths: elixirc_paths(Mix.env()),
       deps: deps(),
@@ -56,7 +56,7 @@ defmodule NxSignal.MixProject do
   # Run "mix help deps" to learn about dependencies.
   defp deps do
     [
-      {:nx, "~> 0.5"},
+      {:nx, github: "elixir-nx/nx", sparse: "nx"},
       {:ex_doc, "~> 0.29", only: :docs}
     ]
   end

mix.lock

@@ -7,7 +7,7 @@
   "makeup_elixir": {:hex, :makeup_elixir, "0.16.0", "f8c570a0d33f8039513fbccaf7108c5d750f47d8defd44088371191b76492b0b", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}, {:nimble_parsec, "~> 1.2.3", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "28b2cbdc13960a46ae9a8858c4bebdec3c9a6d7b4b9e7f4ed1502f8159f338e7"},
   "makeup_erlang": {:hex, :makeup_erlang, "0.1.1", "3fcb7f09eb9d98dc4d208f49cc955a34218fc41ff6b84df7c75b3e6e533cc65f", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}], "hexpm", "174d0809e98a4ef0b3309256cbf97101c6ec01c4ab0b23e926a9e17df2077cbb"},
   "nimble_parsec": {:hex, :nimble_parsec, "1.2.3", "244836e6e3f1200c7f30cb56733fd808744eca61fd182f731eac4af635cc6d0b", [:mix], [], "hexpm", "c8d789e39b9131acf7b99291e93dae60ab48ef14a7ee9d58c6964f59efb570b0"},
-  "nx": {:hex, :nx, "0.5.0", "c5e62e82606ff372d986e72cce505c98421bb4305ce9cc8e439fe6cc1966c6ad", [:mix], [{:complex, "~> 0.5", [hex: :complex, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.0 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "b29c246318181c3ebfcf0f230a0d33783ac4c92dfa34ca3aa5b9b38ae58c187e"},
+  "nx": {:git, "https://github.com/elixir-nx/nx.git", "16ecbc6dbbde5fc5e122f8013601bcc4af2ef4c1", [sparse: "nx"]},
   "telemetry": {:hex, :telemetry, "1.2.1", "68fdfe8d8f05a8428483a97d7aab2f268aaff24b49e0f599faa091f1d4e7f61c", [:rebar3], [], "hexpm", "dad9ce9d8effc621708f99eac538ef1cbe05d6a874dd741de2e689c47feafed5"},
   "xla": {:hex, :xla, "0.4.3", "cf6201aaa44d990298996156a83a16b9a87c5fbb257758dbf4c3e83c5e1c4b96", [:make, :mix], [{:elixir_make, "~> 0.4", [hex: :elixir_make, repo: "hexpm", optional: false]}], "hexpm", "caae164b56dcaec6fbcabcd7dea14303afde07623b0cfa4a3cd2576b923105f5"},
 }