import io
import pickle

import numpy as np

import megengine as mge
import megengine.functional as F
import megengine.module as M
import megengine.utils.comp_graph_tools as cgtools
from megengine.core._trace_option import set_symbolic_shape
from megengine.jit import trace
from megengine.traced_module import trace_module

set_symbolic_shape(True)


class Main(M.Module):
    def forward(self, x):
        return x["data"]


class PreProcess(M.Module):
    def __init__(self):
        super().__init__()
        self.A = F.zeros((1,))
        self.I = 
F.ones((1,))