def verify_take(src_shape, indices_src, axis=None):
src_dtype = "float32"
indices_dtype = "int32"
indices_src = np.array(indices_src, dtype=indices_dtype)
a = sym.Variable("a")
indices = sym.Variable("indices")
if axis is None:
out = sym.take(a, indices)
else:
out = sym.take(a, indices, axis=axis)
for target, ctx in ctx_list():
# set input
shape_dict = {"a": src_shape, "indices": indices_src.shape}
type_dict = {"a": src_dtype, "indices": indices_dtype}
graph, lib, _ = nnvm.compiler.build(out,
target,
shape=shape_dict,
dtype=type_dict)
m = graph_runtime.create(graph, lib, ctx)
shape_size = 1
for i in range(len(src_shape)):
shape_size = shape_size * src_shape[i]
a_src = np.arange(shape_size, dtype=src_dtype).reshape((src_shape))
if axis is None:
out_np = np.take(a_src, indices_src)
else:
out_np = np.take(a_src, indices_src, axis=axis)
#print("out_np:", out_np.shape, "\n", out_np)
m.run(a=a_src, indices=indices_src)
out = m.get_output(0, tvm.nd.empty(out_np.shape, dtype=src_dtype))
#print("out:", out.shape, "\n", out.asnumpy())
np.testing.assert_allclose(out.asnumpy(), out_np, atol=1e-5, rtol=1e-5)
def verify_take(src_shape, indices_src, axis=None):
src_dtype = "float32"
indices_dtype = "int32"
indices_src = np.array(indices_src, dtype=indices_dtype)
a = sym.Variable("a", shape=src_shape)
indices = sym.Variable("indices", shape=indices_src.shape)
y = sym.take(a, indices, axis=axis)
def forward(a, indices):
return np.take(a, indices=indices, axis=axis)
a_src = np.arange(np.prod(src_shape), dtype=src_dtype).reshape(src_shape)
check_function(y, forward,
dtype={'a': src_dtype, 'indices': indices_dtype},
values={'a': a_src, 'indices': indices_src})
def verify_take(src_shape, indices_src, axis=None):
src_dtype = "float32"
indices_dtype = "int32"
indices_src = np.array(indices_src, dtype=indices_dtype)
a = sym.Variable("a", shape=src_shape)
indices = sym.Variable("indices", shape=indices_src.shape)
y = sym.take(a, indices, axis=axis)
def forward(a, indices):
return np.take(a, indices=indices, axis=axis)
a_src = np.arange(np.prod(src_shape), dtype=src_dtype).reshape(src_shape)
check_function(y, forward,
dtype={'a': src_dtype, 'indices': indices_dtype},
values={'a': a_src, 'indices': indices_src})