Function abs

tests/python/tvmscript/test_tvmscript_roundtrip.py:1998–2008

()

def abs():
    @T.prim_func(s_tir=True)
    def abs(a: T.handle) -> None:
        A = T.match_buffer(a, (128, 128), "float32")

        for i, j in T.grid(128, 128):
            with T.sblock("A"):
                vi, vj = T.axis.remap("SS", [i, j])
                A[vi, vj] = T.abs(A[vi, vj])

    return abs


def constant_folding():