tvm sum of rows reduction example https://tvm.apache.org/docs/v0.13.0/how_to/work_with_schedules/reduction.html

example.py

# apache-tvm==0.14.dev273
# numpy==1.26.4
# pytest==8.3.5

from __future__ import absolute_import, print_function
import tvm
import tvm.testing
from tvm import te
import numpy as np

n = te.var("n")
m = te.var("m")
A = te.placeholder((n, m), name="A")
k = te.reduce_axis((0, m), "k")
B = te.compute((n,), lambda i: te.sum(A[i, k], axis=k), name="B")

# Create schedule
s = te.create_schedule(B.op)

# Build and execute
mod = tvm.build(s, [A, B])
ctx = tvm.cpu()

# set up input and output array
n, m = 10, 20  # Example dimensions
a = tvm.nd.array(np.arange(n*m, dtype=np.float32).reshape((n, m)))
b = tvm.nd.array(np.zeros((n,),dtype=np.float32))

# Execute
mod(a, b)
print(b)