Function fcompute

tests/python/te/test_te_create_primfunc.py:969–973

(*axes)

        x = te.placeholder([8, 8, 8, 8, 8], "float32", "x")

        def fcompute(*axes):
            r1 = te.reduce_axis(tvm.ir.Range.from_min_extent(0, axes[1]))
            r2 = te.reduce_axis(tvm.ir.Range.from_min_extent(0, r1))
            all_axes = [*axes, r1, r2]
            return te.sum(x(*all_axes), [r1, r2])

        y = te.compute([8, 8, 8], fcompute)
        f = te.create_prim_func([x, y])