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)