def verify_any_ndarray_size(data_np_shape):
v = relay.var("v", shape=any_dims(len(data_np_shape)), dtype="float32")
n = relay.ndarray_size(v, dtype="int32")
mod = tvm.IRModule()
mod["main"] = relay.Function([v], n)
np_data = np.zeros(data_np_shape, dtype="float32")
ref_res = np.size(np_data)
check_result([np_data], mod, ref_res)
def verify_any_ndarray_size(data_np_shape):
v = relay.var("v", shape=any_dims(len(data_np_shape)), dtype='float32')
n = relay.ndarray_size(v, dtype='int32')
mod = tvm.IRModule()
mod['main'] = relay.Function([v], n)
np_data = np.zeros(data_np_shape, dtype='float32')
ref_res = np.size(np_data)
for kind in ["debug", "vm"]:
ex = relay.create_executor(kind, mod=mod, ctx=tvm.cpu(), target="llvm")
result = ex.evaluate()(np_data)
tvm.testing.assert_allclose(result.asnumpy(), ref_res)
def before(dtype):
x = relay.var("x", shape=c_shape, dtype="float32")
y = relay.var("y", shape=c_shape, dtype="float32")
z = relay.ndarray_size(x + y, dtype)
return relay.Function([x, y], z)