Function check_assign

tests/python/runtime/test_runtime_trace.py:39–57

(dtype)

        return

    def check_assign(dtype):
        n = 4
        x = te.placeholder((n, n, n), name="X", dtype=dtype)
        y = te.compute(
            x.shape, lambda i, j, k: tvm.tirx.trace([x[i][j][k]], "tvm.tirx.trace_callback2")
        )
        z = te.compute(
            x.shape, lambda i, j, k: tvm.tirx.trace([y[i][j][k]], "tvm.tirx.trace_callback2")
        )
        f = tvm.compile(te.create_prim_func([x, y, z]), "llvm")

        xnd = tvm.runtime.tensor(np.ones((n, n, n), dtype=x.dtype))
        ynd = tvm.runtime.tensor(np.zeros((n, n, n), dtype=y.dtype))
        znd = tvm.runtime.tensor(np.zeros((n, n, n), dtype=z.dtype))
        f(xnd, ynd, znd)

        assert np.array_equal(xnd.numpy(), np.ones((n, n, n)))
        assert np.array_equal(ynd.numpy(), np.ones((n, n, n)))
        assert np.array_equal(znd.numpy(), np.ones((n, n, n)))

    for t in ["float64", "float32", "int64", "int32"]:
        check_assign(t)