@triton.autotune(configs=_k_configs, key=['N'])
@triton.jit
def k_dropout_bw(GRAD_IN, GRAD_OUT, INPUTS, BIAS, SEEDS, stride_grad,
    stride_inputs, N, p, **META):
    """
    Apply dropout on an input tensor
    GRAD_OUT    (M, N)
    GRAD_IN     (M, N)
    BIAS        (N,)
    SEEDS       (M,)
    p : dropout probability
    """
    BLOCK_SIZE = META['BLOCK_SIZE']
    row = tl.program_id(axis=0)
    col = tl.program_id(axis=1)
    grad_offsets = row * stride_grad + col * BLOCK_SIZE + tl.arange(0,
        BLOCK_SIZE)
    mask = col * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) < N
    grad_out_ptrs = GRAD_OUT + grad_offsets
    grad_out = tl.load(grad_out_ptrs, mask=mask)
    if META['ACTIVATION_GRAD']:
        input_ptrs = (INPUTS + row * stride_inputs + col * BLOCK_SIZE + tl.
            arange(0, BLOCK_SIZE))
        inputs = tl.load(input_ptrs, mask=mask)
        if META['USE_BIAS']:
            b_ptrs = BIAS + col * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
            b = tl.load(b_ptrs, mask=mask)
            inputs += b
        act_grad = META['ACTIVATION_GRAD'](inputs)
        grad_out *= act_grad
    if p > 0.0:
        output = _drop_and_scale(SEEDS, row, p, grad_offsets, grad_out)
    else:
        output = grad_out
    y_ptrs = GRAD_IN + grad_offsets
    tl.store(y_ptrs, output, mask=mask)
