@triton.jit
def load_full_2d(ptr, sz0: tl.constexpr, sz1: tl.constexpr, stride0=None,
    stride1=1):
    """Load 2d block [0,...,sz0-1] x [0,...,sz1-1]"""
    stride0 = stride0 or sz1
    offs = get_2d_offset(tl.arange(0, sz0), tl.arange(0, sz1), stride0, stride1
        )
    mask = get_2d_mask(tl.arange(0, sz0), tl.arange(0, sz1), sz0, sz1)
    return tl.load(ptr + offs, mask)
