Function test_gemm_async_partial_scale_factor

tests/python/tirx/test_op.py:60–73

Regression test for F7: gemm_async must reject partial scale factors.

()

def test_gemm_async_partial_scale_factor():
    """Regression test for F7: gemm_async must reject partial scale factors."""
    from tvm.tirx.script.builder.tirx import gemm_async

    A = decl_buffer((64, 64), "float16", scope="shared")
    B = decl_buffer((64, 64), "float16", scope="shared")
    C = decl_buffer((64, 64), "float16", scope="shared")
    SF = decl_buffer((64,), "float16", scope="shared")

    with pytest.raises(ValueError, match="SFA and SFB must both be provided or both be None"):
        gemm_async(C[:, :], A[:, :], B[:, :], SFA=SF[:])

    with pytest.raises(ValueError, match="SFA and SFB must both be provided or both be None"):
        gemm_async(C[:, :], A[:, :], B[:, :], SFB=SF[:])


if __name__ == "__main__":