@triton.jit
def scan_kernel(x_ptr, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.
    constexpr, AXIS: tl.constexpr):
    range_m = tl.arange(0, BLOCK_SIZE_M)
    range_n = tl.arange(0, BLOCK_SIZE_N)
    x = tl.load(x_ptr + range_m[:, None] * BLOCK_SIZE_N + range_n[None, :])
    x = tl.cumsum(x, axis=AXIS)
    tl.store(x_ptr + range_m[:, None] * BLOCK_SIZE_N + range_n[None, :], x)
