`Nx.conv` shortcomings (vs `numpy.conv`)

[laszlokorte]

I recently spent some time translating some of my numpy image processing experiments to Elixir Nx.

I noticed a few unexpected missing features, especially regarding convolution. Maybe it is partly due my misunderstanding, so I would love to be corrected. If the features below are indeed missing, I would like to help implementing them.

1. Nx.conv/2 works only for 4+ dimensional tensors. I understand that Nx uses the basic concept of {batch, channel, width, height} dimensions and from this point of view a convolution across batches might not make sense. But from a more general mathematical perspective of discrete convolution I would expect it to to also work on any combination of 1D, 2D and 3D data and kernels.
2. cyclic padding: Nx.conv/2 supports only zero padding and Nx.pad/2 supports only constant padding, and Nx.reflect is a separate function for reflection padding. From perspective of signal processing and discrete fourier transforms, cyclic padding ([1,2,3] => [2,3, 1,2,3 ,1,2]) would be the most natural since discrete convolution with cyclic padding corresponds to multiplication in the discrete fourier spectra. But Nx does not support cyclic padding at all. The simplest improvement would be to add Nx.cycle/2 as alternative to Nx.reflect/2 or to allow more options to Nx.pad/2 like numpy.pad
3. Nx.grad/2 does not work with Nx.conv(image, kernel) if the the image is constant. I know that there is Axon for building neural networks, but at least for small experiments and tutorials it would be nice to quickly calculate the gradient only in respect to the kernel.

image = Nx.tensor([[0, 0, 0], [0, 1, 0], [0, 0, 0]]) |> Nx.new_axis(0) |> Nx.new_axis(0)
kernel = Nx.tensor([[1, 1, 1], [1, 1, 1], [1, 1, 1]]) |> Nx.new_axis(0) |> Nx.new_axis(0)

# this does not work
Nx.Defn.grad({kernel}, fn {k} ->
  Nx.conv(
    image,
    k
  )
  |> Nx.sum()
end)

# but this does
Nx.Defn.grad({kernel, image}, fn {k,i} ->
  Nx.conv(
    i,
    k
  )
  |> Nx.sum()
end)

# and this does
Nx.Defn.grad({image}, fn {i} ->
  Nx.conv(
    image,
    k
  )
  |> Nx.sum()
end)

As far as I understand from the error message the issue with grad/2 is that conv/2 calls Nx.pad internally (even for :same with 0 size padding) and this call to pad can not be serialized as required by grad if image is not part of the computation graph. On the other hand applying pad manually to image inside grad does work, even if image is constant:

# this works:
Nx.Defn.grad({kernel}, fn {k} ->
  Nx.pad(
    image,
    0,
    [{1, 1, 1}, {1, 1, 1}, {0, 0, 0}, {0, 0, 0}]
  )
  |> Nx.add(k)
  |> Nx.product()
end)

I have composed a Livebook with some more detailed examples:
https://github.com/laszlokorte/nx-improvements-livebook/blob/main/ideas.livemd

I also noticed that there is https://github.com/elixir-nx/nx_signal that may also provide conv implementations (both fft based and time-domain based) but it is not quite clear how official/maintained nx_signal is and it would make Nx more approachable if either Nx' own conv implementation would be be more feature complete/consistent or if it would be clearly stated somewhere to not use Nx.conv but NxSignal.conv instead.

[polvalente]

For image signal processing capabilities, NxSignal's convolution features are likely closer to what you're looking for. You can consider anything under the elixir-nx org to be as official as the Nx library itself :)

@josevalim Do you think it's ok if we point users to NxSignal from the Nx docs?

[polvalente]

Regarding the padding options, I think adding Nx.cycle as a counterpart to reflect might be worth considering. An easy implementation could come as some combination of Nx.tile + Nx.slice + Nx.concatenate.

And for grad wrt k, this is likely just an oversight or some weird edge case in the grad engine, especially given that grad wrt {i, k} works.

All in all, each of the 3 points seems like we should tackle as separate issues!

If you'd like to, you can go ahead and break out 3 as a bug in itself directly.

2 is a feature proposal I'm particularly open to, but we should discuss whether to have Nx.cycle or to accept cyclic padding in Nx.pad itself. If Nx.cycle is implementable over other Nx functions, maybe we just provide it as a wrapper first. After we have #946 we can think about providing these as overridable functions without the need for a full-on Nx compiler. Let's also break it out as another issue? Then 1 is left for this one.

[laszlokorte]

Thanks for confirming this! I was confused because the nx_signal readme has no links to hexdoc and it says:

This library comes from the author's urge to experiment with audio processing in Elixir through Nx.

So I was not sure if nx_signal itself is just an experiment. But now I saw that there was actually a new release last week, that I had not noticed yet. Nice 🥳

I am fine with breaking the issue into 3 separate requests.

[polvalente]

@laszlokorte I'll try to update the readme there. If you want to submit a PR (or issues) to improve NxSignal docs, I'd be happy to review (or tackle, in the case of issues)

[josevalim]

Yeah, we can update the README to point to it.

1. We can support other convs as long as it is implemented on top of existing primitives
2. Ditto (i.e. backends should not need to handle new options/callbacks)
3. Looks like a bug, yes

[laszlokorte]

regarding (2) I just created a PR for the cyclic padding.

After knowing about nx_signal and spending some more time thinking about it I do now better understand that Nx.conv works only on 3+ dimensional tensors because it handles the batch dimension correctly and whereas NxSignal.conv works on any dimensional data but does not handle batch dimensions. I am no longer convinced that my proposal (1) makes sense.

If (3) is regarded as bug I will create a separate issue but I have no idea how to go about fixing it yet.

[polvalente]

@laszlokorte batch dimensions in NxSignal can be dealt with by vectorizing your input