@triton.jit
def store_full_2d(vals, ptr, sz0: tl.constexpr, sz1: tl.constexpr, stride0=
    None, stride1=1):
    """Store 2d block into matrix (defined by ptr)"""
    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)
    tl.store(ptr + offs, vals, mask)
