from padiff import auto_diff
import torch
import paddle

class SimpleModule(torch.nn.Module):
  def __init__(self):
      super(SimpleModule, self).__init__()
      self.linear1 = torch.nn.Linear(100, 10)
  def forward(self, x):
      x = self.linear1(x)
      return x

class SimpleLayer(paddle.nn.Layer):
  def __init__(self):
      super(SimpleLayer, self).__init__()
      self.linear1 = paddle.nn.Linear(100, 10)
  def forward(self, x):
      x = self.linear1(x)
      return x

module = SimpleModule()
layer = SimpleLayer()

inp = paddle.rand((100, 100)).numpy().astype("float32")
inp = ({'x': torch.as_tensor(inp) },
     {'x': paddle.to_tensor(inp)})

auto_diff(module, layer, inp, atol=1e-4, auto_init=True)

############################
#      torch_model.py      #
############################

from padiff import *
import torch

class SimpleModule(torch.nn.Module):
  def __init__(self):
      super(SimpleModule, self).__init__()
      self.linear1 = torch.nn.Linear(100, 10)
  def forward(self, x):
      x = self.linear1(x)
      return x

module = SimpleModule()
module = create_model(module)

inp = paddle.ones((100, 100)).numpy().astype("float32")

for i in range(6):
    out = module(torch.as_tensor(inp))
    loss = out.mean()
    module.backward(loss)
    module.try_dump(2, f"./torch/step_{i}")


############################
#      paddle_model.py     #
############################

from padiff import *
import paddle

class SimpleLayer(paddle.nn.Layer):
  def __init__(self):
      super(SimpleLayer, self).__init__()
      self.linear1 = paddle.nn.Linear(100, 10)
  def forward(self, x):
      x = self.linear1(x)
      return x

# 此处需自行保证两个模型的初始权重以及输入数据是对齐的
layer = SimpleLayer()
layer = create_model(layer)

inp = paddle.rand((100, 100)).numpy().astype("float32")

for i in range(6):
    out = layer(paddle.to_tensor(inp))
    loss = out.mean()
    layer.backward(loss)
    layer.try_dump(2, f"./paddle/step_{i}")


############################
#         check.py        #
############################

from padiff import *

for i in range(6):
    if i % 2 == 0:
        assert check_report(f"./torch/step_{i}", f"./paddle/step_{i}") == True
        assert check_params(f"./torch/step_{i}", f"./paddle/step_{i}") == True