import numpy as np
import onnx
import onnxruntime as ort
import torch


class SmallClassifier(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = torch.nn.Sequential(
            torch.nn.Linear(4, 8),
            torch.nn.ReLU(),
            torch.nn.Linear(8, 3),
        )

    def forward(self, features):
        return self.layers(features)


torch.manual_seed(7)
model = SmallClassifier().eval()
example_input = torch.randn(1, 4)

onnx_program = torch.onnx.export(
    model,
    (example_input,),
    dynamo=True,
    input_names=["features"],
    output_names=["scores"],
)
onnx_program.save("small-classifier.onnx")

onnx_model = onnx.load("small-classifier.onnx")
onnx.checker.check_model(onnx_model)

session = ort.InferenceSession(
    "small-classifier.onnx",
    providers=["CPUExecutionProvider"],
)
ort_inputs = {session.get_inputs()[0].name: example_input.numpy()}
ort_output = session.run(None, ort_inputs)[0]

with torch.no_grad():
    torch_output = model(example_input).numpy()

max_difference = np.max(np.abs(torch_output - ort_output))

print("exported: small-classifier.onnx")
print(f"checked: opset={onnx_model.opset_import[0].version}")
print(f"runtime output shape: {ort_output.shape}")
print(f"max difference: {max_difference:.8f}")