text: import numpy as np
  
  import megengine as mge
  import megengine.autodiff as ad
  import megengine.functional as F
  import megengine.module as M
  import megengine.optimizer as optim
  from megengine.jit import trace
  
  
  def minibatch_generator(batch_size):
      while True:
          inp_data = np.zeros((batch_size, 2))
          label = np.zeros(batch_size, dtype=np.int32)
          for i in range(batch_size):
              inp_data[i, :] = np.random.rand(2) * 2 - 1
              label[i] = 1 if np.prod(inp_data[i]) < 0 else 0
          yield {"data": inp_data.astype(np.float32), "label": label.astype(np.int32)}
  
  
  class XORNet(M.Module):
      def __init__(self):
          self.mid_dim = 14
          self.num_class = 2
          super().__init__()
          self.fc0 = M.Linear(self.num_class, self.mid_dim, bias=True)
          self.fc1 = M.Linear(self.mid_dim, self.mid_dim, bias=True)
          self.fc2 = M.Linear(self.mid_dim, self.num_class, bias=True)
  
      def forward(self, x):
          x = self.fc0(x)
          x = F.tanh(x)
          x = self.fc1(x)
          x = F.tanh(x)
          x = self.fc2(x)
          return x
  
  
  def main():
  
      if not mge.is_cuda_available():
          mge.set_default_device("cpux")
  
      net = XORNet()
      gm = ad.GradManager().attach(net.parameters())
      opt = optim.SGD(net.parameters(), lr=0.01, momentum=0.9)
      batch_size = 64
      train_dataset = minibatch_generator(batch_size)
      val_dataset = minibatch_generator(batch_size)
  
      def train_fun(data, label):
          opt.clear_grad()
          with gm:
              pred = net(data)
              loss = F.loss.cross_entropy(pred, label)
              gm.backward(loss)
          opt.step()
          return pred, loss
  
      def val_fun(data, label):
          pred = net(data)
          loss = F.loss.cross_entropy(pred, label)
          return pred, loss
  
      @trace(symbolic=True, capture_as_const=True)
      def pred_fun(data):
          pred = net(data)
          pred_normalized = F.softmax(pred)
          return pred_normalized
  
      data = np.random.random((batch_size, 2)).astype(np.float32)
      label = np.zeros((batch_size,)).astype(np.int32)
      train_loss = []
      val_loss = []
      for step, minibatch in enumerate(train_dataset):
          if step > 1000:
              break
          data = minibatch["data"]
          label = minibatch["label"]
          net.train()
          _, loss = train_fun(data, label)
          train_loss.append((step, loss.numpy()))
          if step % 50 == 0:
              minibatch = next(val_dataset)
              net.eval()
              _, loss = val_fun(data, label)
              loss = loss.numpy()
              val_loss.append((step, loss))
              print("Step: {} loss={}".format(step, loss))
          opt.step()
  
      test_data = np.array(
          [
              (0.5, 0.5),
              (0.3, 0.7),
              (0.1, 0.9),
              (-0.5, -0.5),
              (-0.3, -0.7),
              (-0.9, -0.1),
              (0.5, -0.5),
              (0.3, -0.7),
              (0.9, -0.1),
              (-0.5, 0.5),
              (-0.3, 0.7),
              (-0.1, 0.9),
          ]
      )
  
      # tracing only accepts tensor as input
      data = mge.tensor(test_data, dtype=np.float32)
      net.eval()
      out = pred_fun(data)
      pred_output = out.numpy()
      pred_label = np.argmax(pred_output, 1)
  
      print("Test data")
      print(test_data)
  
      with np.printoptions(precision=4, suppress=True):
          print("Predicated probability:")
          print(pred_output)
  
      print("Predicated label")
      print(pred_label)
  
      model_name = "xornet_deploy.mge"
  
      print("Dump model as {}".format(model_name))
      pred_fun.dump(model_name, arg_names=["data"])
  
  
  if __name__ == "__main__":
      main()
  