@triton.autotune(configs=[triton.Config({}, num_warps=4), triton.Config({},
    num_warps=8), triton.Config({}, num_warps=16)], reset_to_zero=['DYDA',
    'DYDC'], key=['C'])
@triton.jit
def _snake_bwd_triton(X, OUT, ALPHA, CR, GRAD, DYDX, DYDA, DYDC, X_stride1,
    X_stride2, X_stride3, OUT_stride1, OUT_stride2, OUT_stride3,
    GRAD_stride1, GRAD_stride2, GRAD_stride3, DYDX_stride1, DYDX_stride2,
    DYDX_stride3, DYDA_stride, DYDC_stride, ALPHA_stride, CR_stride, C, N,
    CORR: tl.constexpr, X_NEEDS_GRAD: tl.constexpr, ALPHA_NEEDS_GRAD: tl.
    constexpr, CR_NEEDS_GRAD: 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)
    GRAD = GRAD + batch_idx * GRAD_stride1 + channel_idx * GRAD_stride2
    grad = tl.load(GRAD + offsets * GRAD_stride3, mask=offsets < N, other=0)
    if CORR:
        cr = tl.load(CR + channel_idx * CR_stride)
    if ALPHA_NEEDS_GRAD | CR_NEEDS_GRAD:
        OUT = OUT + batch_idx * OUT_stride1 + channel_idx * OUT_stride2
        out = tl.load(OUT + offsets * OUT_stride3, mask=offsets < N, other=0)
        outgrad = tl.sum(out * grad, axis=0)
    if X_NEEDS_GRAD | ALPHA_NEEDS_GRAD:
        X = X + batch_idx * X_stride1 + channel_idx * X_stride2
        x = tl.load(X + offsets * X_stride3, mask=offsets < N, other=0)
        alpha = tl.load(ALPHA + channel_idx * ALPHA_stride)
        sin2ax = tl.sin((2 * alpha * x).to(tl.float32)).to(x.type)
        dydx = (sin2ax + 1) * grad
        if CORR:
            dydx = dydx / cr
    if X_NEEDS_GRAD:
        DYDX = DYDX + batch_idx * DYDX_stride1 + channel_idx * DYDX_stride2
        tl.store(DYDX + offsets * DYDX_stride3, dydx, mask=offsets < N)
    if ALPHA_NEEDS_GRAD:
        dyda = (tl.sum(x * dydx, axis=0) - outgrad) / alpha
        tl.atomic_add(DYDA + channel_idx * DYDA_stride, dyda)
    if CR_NEEDS_GRAD:
        dydc = -outgrad / cr
        tl.atomic_add(DYDC + channel_idx * DYDC_stride, dydc)
