Convolutions?

[Nin17]

What is the recommended array API standard compatible way to compute a convolution?

To me it seems the options are to either use a for loop over the kernel, use ffts or use scipy.signal.oaconvolve which seems to be compatible.

None of these options are particularly performant for my use case.
Personally I'd be interested in something with which I could implement some of the filters found in scipy.ndimage and discrete wavelet transforms in PyWavelets.

There doesn't seem to have been any interest in convolutions being in the standard directly - I could only find this comment:

• I'll mention cov, corrcoef, convolve, convolve2d, correlate, and correlate2d for completeness: happy to leave them out.

Many libraries include similarly named convolution functions:

• pytorch
  • torch.nn.functional.conv1d
  • torch.nn.functional.conv2d
  • torch.nn.functional.conv3d
• tensorflow
  • tf.nn.conv1d
  • tf.nn.conv2d
  • tf.nn.conv3d
• mlx
  • mlx.core.conv1d
  • mlx.core.conv2d
  • mlx.core.conv3d
  • mlx.core.conv_general
• paddle
  • paddle.conv1d
  • paddle.conv2d
  • paddle.conv3d

jax has jax.lax.conv_general_dilated for which the above functions seem to be special cases.

Is this in scope for inclusion in the standard, potentially as an extension module?

[ev-br]

There's an issue about adding it to array-api-extra: data-apis/array-api-extra#286

As for the "recommended"/performant way, SciPy roadmap has the convolution item (for a decade or so, predating anything array api ): https://docs.scipy.org/doc/scipy/dev/roadmap-detailed.html#signal
Not sure whether a single best performance function is even possible in the array API world,

This is to say, I think 1) some array agnostic function is a good fit for array-api-extra, and 2) an investigation of performance landscape across array sizes, array backends and hardware types would be a nice exercise.

[rgommers]

Convolution functions were discussed early on as a candidate for an extension. I think it'd be useful for someone to write a good summary of the current set of convolution functions out there and their overlap.

I think the conclusion will be that the overlap is small, or more precisely, that there are different sets of commonality:

• PyTorch, MLX, jax.lax, (and TF, not so relevant anymore): conv1d & co for deep learning focused convolution. Output size == input size, cross-correlation (no kernel flip) mandatory batch parameter, allows striding. Different convention between PyTorch and MLX for 3-D (NHWC vs. NHWC), JAX can do both layouts.
• NumPy, CuPy, Dask (and SciPy): 1-D signal processing style regular correlation (kernel is flipped) only , no striding
• CuPy (in cupyx.scipy.ndimage) and SciPy: 2-D image processing style convolution, more similar to conv2d in deep learning frameworks but no batching/striding.

I think the deep learning style functions are potentially most interesting for an extension. I'd be interested to see how scipy.ndimage.convolve2d maps onto conv2d (is it a strict subset?) and whether one can write a conv*d style function for array-api-extra.