TVM Parse Modul

tvm_bug_11_03.py

import numpy as np

import tvm
from tvm import relay
import tvm.relay.testing
import pytest
from numpy import isclose
from typing import Union

from tvm import relay, te, auto_scheduler

data = relay.var("data", relay.TensorType((64, 784), "float32"))
# "linear1.weight" will lead to bugs
# "linear_weight" can run without error.
w1 = relay.var("linear1.weight", relay.TensorType((784, 64), "float32"))

net = relay.nn.matmul(data, w1)
net = relay.nn.relu(net)

mod = tvm.IRModule.from_expr(net)

print("==" * 25, "TVM IRs", "==" * 25)
print(mod)


SEMVER = '#[version = "0.0.5"]\n'

def assert_graph_equal(lhs, rhs):
    tvm.ir.assert_structural_equal(lhs, rhs, map_free_vars=True)

def roundtrip(expr):
    x = tvm.parser.fromtext(expr.astext())
    assert_graph_equal(x, expr)

# Testing Utilities for full modules.
def parse_module(code):
    mod = tvm.parser.parse(SEMVER + code)
    roundtrip(mod)
    return mod

program = str(mod)
parsed = parse_module(program)
print(parsed)