@triton.jit
def _dropout(X, O, stride_x1, stride_x2, stride_o1, stride_o2, dropout_prob,
    dropout_seed, M, N, BLOCK: tl.constexpr):
    offs_m = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    offs_n = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
    X = X + offs_m[:, None] * stride_x1 + offs_n[None, :] * stride_x2
    x = tl.load(X, mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))
    offsets = offs_m[:, None] * M + offs_n[None, :]
    x = apply_dropout(x, offsets, dropout_prob, dropout_seed)
    O = O + offs_m[:, None] * stride_o1 + offs_n[None, :] * stride_o2
    tl.store(O, x, mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))
