Function print_matrix

cutlass.py/mma.py:171–178

(mtx, rows, cols, func=lambda x: x, prompt="")

    frag = wgmma_op.get_accumulator_matrix()
    
    def print_matrix(mtx, rows, cols, func=lambda x: x, prompt=""):
        print(prompt)
        for x in range(rows):
            for y in range(cols):
                item = mtx[x][y]
                item = func(item)
                print(item, end=" ")
            print()
    
    print_matrix(frag, 64, 128, func=lambda x: f"T{x[2]}({x[0]},{x[1]})", prompt="Fragment WgMMA M64N128K16:")