explicit_conv1d#

explicit_conv1d(x, weight)#

Depthwise “same” convolution.

\[y[b, c, n] = \sum_{k=0}^{K-1} w[c, k] \, x[b, c, n + k - \lfloor K/2 \rfloor]\]

with zero padding outside \([0, L)\). The result matches fft_conv1d on the same inputs and is differentiable with respect to both arguments.

Parameters:

  • x (Tensor) – Input of shape (batch, channels, seq_len), bfloat16 or float16, on a CUDA device. seq_len must be a multiple of 128 in [128, 2048].

  • weight (Tensor) – Filters of shape (channels, filter_len) with filter_len <= 2 * seq_len; any floating-point dtype, converted to x’s dtype internally.

Returns:

(batch, channels, seq_len) in x’s dtype.

Return type:

Tensor

Note

Requires a GPU of compute capability 10.x (datacenter Blackwell and later); other devices raise NotImplementedError. The filter’s gradient is accumulated in float32.

Example

>>> x = torch.randn(8, 128, 1024, device="cuda", dtype=torch.bfloat16, requires_grad=True)
>>> w = torch.randn(128, 2047, device="cuda", dtype=torch.bfloat16, requires_grad=True)
>>> y = explicit_conv1d(x, w)
>>> y.sum().backward()
On this page