@triton.jit
def load_1d(ptr, sz: tl.constexpr, n, max, stride=1):
    """Chunk 1d vector (defined by ptr) into 1d grid, where each chunk has size sz.
    Load the nth chunk. Ie, load [n*sz,...,(n+1)*sz-1]."""
    offs = get_1d_offest(sz, n)
    mask = get_1d_mask(offs, max)
    return tl.load(ptr + offs, mask)
