@triton.heuristics({'EVEN_N': lambda *args, **meta: args[3] % meta[
    'BLOCK_COL'] == 0})
@triton.autotune(configs=[triton.Config({'BLOCK_COL': 32}, num_stages=5,
    num_warps=2), triton.Config({'BLOCK_COL': 64}, num_stages=5, num_warps=
    2), triton.Config({'BLOCK_COL': 128}, num_stages=3, num_warps=4),
    triton.Config({'BLOCK_COL': 256}, num_stages=3, num_warps=8), triton.
    Config({'BLOCK_COL': 512}, num_stages=3, num_warps=8), triton.Config({
    'BLOCK_COL': 1024}, num_stages=3, num_warps=16)], key=['N'])
@triton.jit
def kernel_bw(GRAD_ACT, GRAD_OUT, ACT_INPUTS, N, stride_gom, stride_aim, **META
    ):
    """
    Go over all the activation inputs, compute the corresponding gradient
    """
    BLOCK_N = META['BLOCK_COL']
    pid_m, pid_n = tl.program_id(axis=0), tl.program_id(axis=1)
    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    act_input_ptrs = ACT_INPUTS + pid_m * stride_aim + rn
    if META['EVEN_N']:
        act_in = tl.load(act_input_ptrs)
    else:
        act_in = tl.load(act_input_ptrs, mask=rn < N, other=0.0)
    grad_act = META['ACTIVATION_GRAD'](act_in)
    grad_out_ptrs = GRAD_OUT + pid_m * stride_gom + rn
    if META['EVEN_N']:
        grad_out = tl.load(grad_out_ptrs)
    else:
        grad_out = tl.load(grad_out_ptrs, mask=rn < N)
    grad_act *= grad_out
    grad_act_ptrs = GRAD_ACT + pid_m * stride_gom + rn
    tl.store(grad_act_ptrs, grad_act, mask=rn < N)
