Genomics and RNA

Causal long convolution with FFT

MediumRaw PyTorchAbout 20 min2 public examples plus private tests

HyenaDNAEvo 1Evo 2

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Your task

Compute a linear causal convolution for each channel using FFTs. Return the original sequence length, including the optional diagonal skip term, without circular wraparound.

Function signature

fft_causal_long_convolution(
    u: torch.Tensor,
    kernel: torch.Tensor,
    options: dict,
    skip: torch.Tensor | None = None,
) -> torch.Tensor

Worked example

Past samples contribute to the current output

For one batch and one channel, let u = [1,2,3], kernel = [2,1], omit skip, and use fft_size = 4. The three returned positions are 2·1 = 2, 2·2 + 1·1 = 5, and 2·3 + 1·2 = 8, so y = [2,5,8] (up to FFT roundoff). The final full-convolution value 3 lies beyond the original sequence length and is cropped; it must not wrap into the start.

Parameters

Parameter	Shape and meaning
`u`torch.Tensor	$[B, D, L]$ real $\mathrm{float32}/\mathrm{float64} [B,D,L]$ , channels first; noncontiguous layouts permitted
`kernel`torch.Tensor	$[D, K]$ Lag-zero-first channelwise causal filters
`skip`torch.Tensor \| NoneDefault `None`	$[D]$ Learned diagonal residual coefficient
`options`dict	Scalar settings. Omitted keys use the defaults below.

Return

torch.Tensor · $[B, D, L]$

Return [B,D,L] causal linear convolution, including diagonal skip only when supplied.

Required behavior

Implement real-valued, channelwise causal LINEAR convolution with an FFT. $u$ is $[B,D,L]$ , kernel is $[D,K]$ , and optional skip is $[D]$ ; all share real $\mathrm{float32}/\mathrm{float64}$ dtype and device, with $B$ , $D$ , $L$ , $K$ positive. One kernel per channel is shared across the batch, and channels never mix.
Return $y [B,D,L]$ where $y[b,d,t]=\sum_{j=0..\operatorname{min}(t,K-1)} \mathrm{kernel}[d,j]\cdot u[b,d,t-j]$ + $\mathrm{skip}[d]\cdot u[b,d,t]$ if skip is supplied. Without skip, return only the convolution. With skip, include the stated diagonal residual once; there is no additional identity residual or bias.
kernel $[d,0]$ is the lag-zero coefficient: do not reverse kernel before an FFT. The only option is fft_size (integer, default $L+K-1$ ). It must be at least $L+K-1$ . Zero-pad BOTH input and kernel to that length, multiply their Fourier transforms with batch broadcasting, invert, and crop the first $L$ positions.
Shorter transforms cause circular wraparound and are invalid. Use rfft and irfft with norm='backward' and explicitly pass the same $n$ to both transforms (essential for odd transform lengths). Extra padding, unequal $L/K$ , kernels longer than the sequence, and singleton lengths are all valid.
Reject empty sequences/kernels and undersized fft_size with ValueError. The exact numerical exceptions torch.fft.rfft and torch.fft.irfft are allowed; upstream fftconv, torch.nn convolution, and filter-generation helpers are not. No masking, dropout, nonlinearity, normalization, or $\epsilon$ is applied.

Why this matters

Multiplication in frequency space implements convolution only after enough zero-padding prevents the tail from wrapping into early outputs.

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def fft_causal_long_convolution(
    u: torch.Tensor,  # (B, D, L)
    kernel: torch.Tensor,  # (D, K)
    options: dict,
    skip: torch.Tensor | None = None,  # (D,)
) -> torch.Tensor:
    """Return a tensor with shape (B, D, L)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement fft_causal_long_convolution")

Ready to implement it?

3 progressive hints, a solution walkthrough, and grading that checks values, shapes and gradients are waiting in the workspace.

Write and run your solution Sign in

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Sources and further reading

A primitive exercise, not an end-to-end reproduction of the named models. No affiliation or endorsement is implied.