Skip to content

Commit

Permalink
docs: improve Nx.conv docs on convolution vs correlation
Browse files Browse the repository at this point in the history
  • Loading branch information
polvalente committed Dec 9, 2024
1 parent cb7fed4 commit b2fdb9a
Showing 1 changed file with 26 additions and 9 deletions.
35 changes: 26 additions & 9 deletions nx/lib/nx.ex
Original file line number Diff line number Diff line change
Expand Up @@ -13007,24 +13007,41 @@ defmodule Nx do
> #### Convolution vs Correlation {: .tip}
>
> `conv/3` does not perform reversion nor conjugation of the kernel.
> `conv/3` does not perform reversion of the kernel.
> This means that if you come from a Signal Processing background,
> you might call this operation "correlation" instead of convolution.
> you might treat it as a cross-correlation operation instead of a convolution.
>
> If you need the proper Signal Processing convolution, you can use
> `reverse/2` and `conjugate/1`, like in the example:
> This function is not exactly a cross-correlation function, as it does not
> perform conjugation of the kernel, as is done in [scipy.signal.correlate](https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.correlate.html).
> This can be remedied via `Nx.conjugate/1`, as seen below:
>
> ```elixir
> axes = Nx.axes(kernel) |> Enum.drop(2)
>
> kernel =
> if Nx.Type.complex?(Nx.type(kernel)) do
> Nx.conjugate(Nx.reverse(kernel, axes: axes))
> Nx.conjugate(kernel)
> else
> Nx.reverse(kernel, axes: axes)
> kernel
> end
>
> Nx.conv(tensor, kernel)
> ```
>
> If you need the proper Signal Processing convolution, such as the one in
> [scipy.signal.convolve](https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.convolve.html),
> you can use `reverse/2`, like in the example:
>
> ```elixir
> reversal_axes =
> case Nx.rank(kernel) do
> 0 -> []
> 1 -> [1]
> 2 -> [0, 1]
> _ -> Enum.drop(Nx.axes(kernel), 2)
> end
>
> Nx.conv(img, kernel)
> kernel = Nx.reverse(kernel, axes: reversal_axes)
>
> Nx.conv(tensor, kernel)
> ```
## Examples
Expand Down

0 comments on commit b2fdb9a

Please sign in to comment.