from pathlib import Path

import torch
from torch import nn
from torch.export import Dim


class ScoreModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(4, 8),
            nn.ReLU(),
            nn.Linear(8, 2),
        )

    def forward(self, features):
        return torch.softmax(self.layers(features), dim=-1)


torch.manual_seed(7)

model = ScoreModel().eval()
example_inputs = (torch.randn(2, 4),)
dynamic_shapes = {"features": {0: Dim("batch", min=1, max=8)}}
eager_output = model(*example_inputs)

exported = torch.export.export(
    model,
    example_inputs,
    dynamic_shapes=dynamic_shapes,
)

test_input = torch.randn(3, 4)
exported_output = exported.module()(test_input)

artifact = Path("score_model.pt2")
torch.export.save(exported, artifact)
loaded = torch.export.load(artifact)
loaded_output = loaded.module()(torch.randn(1, 4))

print(f"torch version: {torch.__version__}")
print(f"eager output shape: {tuple(eager_output.shape)}")
print(f"exported type: {type(exported).__name__}")
print(f"graph nodes: {len(list(exported.graph_module.graph.nodes))}")
print(f"range constraints: {exported.range_constraints}")
print(f"exported output shape: {tuple(exported_output.shape)}")
print(f"artifact: {artifact} ({artifact.stat().st_size} bytes)")
print(f"loaded output shape: {tuple(loaded_output.shape)}")