(nthd)
| 356 | return fun |
| 357 | |
| 358 | def verify(nthd): |
| 359 | func = sched(nthd) |
| 360 | nn = 3 |
| 361 | # checks three typical cases |
| 362 | vals = [nthd - 1, nthd, nthd + 1] |
| 363 | for kk in [x for x in vals]: |
| 364 | size = (nn, kk) |
| 365 | a = tvm.runtime.tensor(np.random.uniform(size=size).astype("float32"), dev) |
| 366 | b = tvm.runtime.tensor(np.zeros(nn, dtype="float32"), dev) |
| 367 | func(a, b) |
| 368 | tvm.testing.assert_allclose(b.numpy(), np.sum(a.numpy(), axis=1), rtol=1e-3) |
| 369 | |
| 370 | verify(16) |
| 371 | verify(32) |