diff --git a/tests/unit/test_d1_3_layout_diag.py b/tests/unit/test_d1_3_layout_diag.py index 30d85d63..76c148b7 100644 --- a/tests/unit/test_d1_3_layout_diag.py +++ b/tests/unit/test_d1_3_layout_diag.py @@ -81,7 +81,8 @@ def main(): # PV A-operand SMEM layout (sP) p_smem_s = utils.sm100.make_smem_layout_a(pv_mma, pv_mma_tiler, BFloat16, 1) - sP_alloc = cute.make_tensor(cute.find(cute.make_layout(p_smem_s.outer), 0), p_smem_s.outer) + # p_smem_s.outer IS a layout, not a shape + sP_alloc = cute.make_tensor(cute.find(p_smem_s.outer, 0), p_smem_s.outer) print(f" sP shape: {cute.shape(sP_alloc)}") sP_2d = cute.group_modes(sP_alloc, 0, 3) print(f" sP_2d shape: {cute.shape(sP_2d)} rank={len(cute.shape(sP_2d))}") @@ -89,7 +90,7 @@ def main(): # PV A-operand fragments tP_tmem = cute.make_tensor(tStS.iterator, p_smem_s.outer) # dummy tmem ptr - tP_smem = cute.make_tensor(cute.find(cute.make_layout(p_smem_s.outer), 0), p_smem_s.outer) + tP_smem = cute.make_tensor(cute.find(p_smem_s.outer, 0), p_smem_s.outer) tOrP_tmem = pv_thr.make_fragment_A(tP_tmem) tCrP_smem = pv_mma.make_fragment_A(tP_smem) tOrP_smem = pv_thr.make_fragment_A(tP_smem)