@triton.autotune(configs=[triton.Config({}, num_warps=4), triton.Config({},
    num_warps=8), triton.Config({}, num_warps=16)], key=['C'])
@triton.jit
def _snake_fwd_triton(X, OUT, ALPHA, CR, X_stride1, X_stride2, X_stride3,
    OUT_stride1, OUT_stride2, OUT_stride3, A_stride, C_stride, C, N, CORR:
    tl.constexpr, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    batch_idx = pid // C
    channel_idx = pid % C
    block_start = tl.program_id(1) * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    X = X + batch_idx * X_stride1 + channel_idx * X_stride2
    x = tl.load(X + offsets * X_stride3, mask=offsets < N)
    alpha = tl.load(ALPHA + channel_idx * A_stride)
    sinax = tl.sin((alpha * x).to(tl.float32)).to(x.type)
    out = x + sinax * sinax / alpha
    if CORR:
        cr = tl.load(CR + channel_idx * C_stride)
        out = out / cr
    OUT = OUT + batch_idx * OUT_stride1 + channel_idx * OUT_stride2
    tl.store(OUT + offsets * OUT_stride3, out, mask=offsets < N)
