MCPcopy Create free account
hub / github.com/apache/tvm / test_tensor

Function test_tensor

tests/python/disco/test_session.py:177–186  ·  view source on GitHub ↗
(session_kind)

Source from the content-addressed store, hash-verified

175
176@pytest.mark.parametrize("session_kind", _all_session_kinds)
177def test_tensor(session_kind):
178 num_workers = 4
179 sess = session_kind(num_workers=num_workers)
180 device = tvm.cpu(0)
181 x_np = np.arange(6).astype("float32").reshape([2, 3])
182 y_np = np.arange(6).astype("float32").reshape([2, 3]) + 1
183 x_disc = _numpy_to_worker_0(sess, x_np, device=device)
184 y_disc = sess.get_global_func("tests.disco.add_one_tensor")(x_disc)
185 y_nd = _numpy_from_worker_0(sess, y_disc, shape=y_np.shape, dtype=y_np.dtype)
186 np.testing.assert_equal(y_nd, y_np)
187
188
189@pytest.mark.parametrize("session_kind", _all_session_kinds)

Callers

nothing calls this directly

Calls 7

_numpy_to_worker_0Function · 0.85
_numpy_from_worker_0Function · 0.85
get_global_funcMethod · 0.80
cpuMethod · 0.45
reshapeMethod · 0.45
astypeMethod · 0.45
arangeMethod · 0.45

Tested by

no test coverage detected

Used in the wild real call sites across dependent graphs

searching dependent graphs…