Round-Trip Contract🔗

The round trip has five steps:

1. Lean defines the expected model family.

2. Lean exports matching PyTorch code.

3. Python trains or modifies the weights.

4. Python writes a named tensor payload.

5. The chosen Lean family importer looks up its required names and parses every nested array at the statically expected shape.

If any check fails, TorchLean rejects the artifact before constructing a model.

The shape of the boundary is intentionally closer to a state_dict contract than to Python object serialization. PyTorch users are used to seeing names such as:

layer1.weight
layer1.bias
layer2.weight
layer2.bias

Each name must map to a particular parameter in the Lean family record.

The bridge delegates loading PyTorch's own serialization format to PyTorch, then exports the required tensor payload as JSON. Nested numeric arrays keep the Lean parser small and make the name-to-parameter mapping inspectable. The file is larger than a binary payload, however, and its numbers are decimal text rather than stored bit patterns. The float chapter is where exact bit-level claims live; this chapter is about names and shapes.

The expected architecture gives each name a job. For the 2 → 3 → 1 MLP, the first weight must contain three output rows with two input coefficients each; its bias contains three values. The second weight has one row of three coefficients, followed by one output bias. A file with four arrays of the right total sizes could still attach them to the wrong meanings. Looking up names and parsing their particular nested shapes establishes a more useful contract than flattening all values and consuming them in file order.

The contract ends at model state. Learning rate, optimizer momentum, random-generator state, and the training data are not reconstructed by these four arrays. That is enough to evaluate the family's forward pass. Resuming the same training trajectory would require a broader artifact and an explicit account of how each additional state component is restored.

MLP Parameter Round Trip🔗

The MLP family exposes the full exchange in four tensors. Export its Python modules from the repository root:

# Generate the MLP family files expected by the companion
# Python script.
lake exe torchlean pytorch_roundtrip --model mlp --action export

Exported MLP PyTorch files under NN/Examples/Interop/PyTorch/MLP/.

That step writes two readable reference modules, TestMLP_PyTorch.py and TestMLP_WithWeights.py, whose class matches what the Lean importer expects to see back: Linear(2, 3), ReLU, Linear(3, 1).

Step three and four are one Python script, train_mlp.py:

# Train the one-pair MLP example and write its named
# parameter payload.
python3 NN/Examples/Interop/PyTorch/MLP/train_mlp.py

Model info: TestMLP(2→3→1, ReLU)
Initial output: tensor([[-0.0945]], grad_fn=<AddmmBackward0>)

Starting training...
Epoch 0, Loss: 1.197846
Epoch 50, Loss: 0.000000
Epoch 100, Loss: 0.000000
Epoch 150, Loss: 0.000000
Final output: tensor([[1.0000]], grad_fn=<AddmmBackward0>)
Saved trained weights to .../NN/Examples/Interop/PyTorch/MLP/mlp.json
Original (fresh init) output: tensor([[-0.0945]], grad_fn=<AddmmBackward0>)
Training improved output by: 1.094462

This run fits one input-target pair, x=(0.5,0.8) and y=1, using two hundred SGD steps at learning rate 0.1. The fit supplies a modified parameter payload whose prediction can be checked after import; it provides no evidence of generalization. The script calls torch.manual_seed(0) before initialization, and repeating it in the recorded environment regenerates the same mlp.json.

Step five reads the file back:

# Load the exported parameter values into the expected Lean
# MLP family.
lake exe torchlean pytorch_roundtrip --model mlp --action import

== MLP import example ==
Loaded: NN/Examples/Interop/PyTorch/MLP/mlp.json
Output (native tensor operations, Float):
[1.000000]

PyTorch finished training at tensor([[1.0000]]). Lean, reading only the JSON file and running its own forward pass, prints [1.000000]. The imported parameters reproduce the observed output at the displayed precision. One rounded output is not a proof that the programs compute the same function on every input.

The complete exchanged payload is:

{
  "params": {
    "layers.0.weight": [
      [0.028833430260419846, 0.10097555816173553],
      [-0.08230451494455338, -0.07359390705823898],
      [-0.03851543739438057, 0.026815736666321754]
    ],
    "layers.0.bias": [
      0.057182908058166504, 0.07928895205259323, -0.00887440424412489
    ],
    "layers.2.weight": [
      [0.11024449020624161, -0.030221307650208473, -0.01965653896331787]
    ],
    "layers.2.bias": [0.9832008481025696]
  },
  "meta": {
    "format": "TorchLean.MLP",
    "pytorch_keys": "layers.*",
    "dtype": "float32"
  }
}

The output bias layers.2.bias is about 0.983. With one training point, that bias alone could fit the target while the other parameters remained fixed; the fit does not identify a unique parameter vector.

The decimal strings retain enough significant digits for finite binary32 values because Python's json.dump prints the shortest decimal that round-trips the underlying binary64 value that tolist() produced from a binary32 tensor. This serialization can preserve those numerical values, while meta.dtype separately records their source dtype. The importer still constructs host Float tensors; source dtype metadata does not select the arithmetic of the Lean forward pass.

The retained transcript's last label says Training improved output by. The current script labels that same absolute change Output change from initialization. It is a change in a prediction, not a validation score: moving a prediction by a larger amount is only helpful relative to the chosen target. The initial and final values above make that interpretation possible without relying on the older label.

Import Validation🔗

The importer is strict about required tensors, and it checks them before constructing the typed family record used by the example.

The checks actually performed are:

• the root is a JSON object, optionally with a params object;

• every required family-specific key exists under one of the accepted naming conventions;

• every scalar leaf is a JSON number;

• every nested array has exactly the length required by the expected Lean shape.

JSON object order is irrelevant: the importer looks up names and constructs a fixed typed record. The family loaders validate a supplied meta.dtype or per-tensor dtype metadata as float32. meta.format is not an architecture proof, and extra parameter keys are ignored. Absent dtype metadata is accepted for the older bare-payload format. Imported numbers become host Float values, so this format does not preserve exact binary32 payload bits.

parseTensor : (s : Shape) → Lean.Json → Option (Tensor Float s)

getTensor? : StateDict → String → (s : Shape) → Option (Tensor Float s)

some [[1.000000, 2.000000], [3.000000, 4.000000], [5.000000, 6.000000]]

some [0.500000, -0.500000, 0.250000]

false

false

false

error: Failed to load MLP state dict

-- Try each complete naming convention as one alternative.
tryKeys "fc1.weight" "fc1.bias" "fc2.weight" "fc2.bias" <|>
  tryKeys "layers.0.weight" "layers.0.bias"
          "layers.2.weight" "layers.2.bias"

RuntimeError: Error(s) in loading state_dict for Sequential:
  size mismatch for 0.weight: copying a param with shape
  torch.Size([2, 3]) from checkpoint, the shape in current model
  is torch.Size([3, 2]).

RuntimeError: Error(s) in loading state_dict for Sequential:
  Missing key(s) in state_dict: "0.bias".

Serialized Tensor Layouts🔗

The imported record supplies this example's family-specific forward pass. Using it through another model API requires an explicit adapter. Even before that step, the names must agree with one of the chosen family's accepted conventions.

The MLP loader accepts two key conventions, fc1.*/fc2.* and layers.0.*/layers.2.*, because PyTorch itself produces both. A module with named attributes gives fc1.weight; the same architecture written as nn.Sequential gives 0.weight, which becomes layers.0.weight once the sequential block is held in an attribute called layers. The index gap between 0 and 2 is not an error: position 1 is the ReLU, and activations have no parameters. The CNN and transformer loaders have their own fixed key sets.

Here are both conventions carrying identical weights, plus a third payload that names one layer each way:

some [5.500000]

some [5.500000]

none

The identity weights and the [1, 1] output row determine the result directly: the hidden vector is [2, 3], the ReLU leaves it alone, and the sum plus the 0.5 bias is 5.5. Both naming conventions construct the same tensors and produce this value.

The mixed payload fails because load is written as tryKeys "fc1.weight" ... <|> tryKeys "layers.0.weight" ..., so each convention is tried as a complete set. A payload that names its first layer fc1 and its second layers.2.weight satisfies neither alternative and is rejected, even though every tensor it needs is present under some name. This loader requires a consistent complete naming convention. Supporting mixed conventions would need an explicit conflict policy; shape checking cannot identify the intended parameter when two candidate keys both have the expected shape.

This small exchange format handles known TorchLean families. Keep the full object in Python when a training pipeline needs PyTorch's complete serialization graph, and export only the checked payload that TorchLean understands.

The identity weights make the naming test easy to check by hand. They would not reveal an accidental transpose of the first weight matrix, since transposing an identity changes nothing. The transformer example below uses an asymmetric matrix to expose that mistake.

Model Families🔗

The three family examples exercise different parameter layouts and forward paths.

# Run the complete MLP exchange using one family-specific
# schema.
lake exe torchlean pytorch_roundtrip --model mlp --action export
python3 NN/Examples/Interop/PyTorch/MLP/train_mlp.py
lake exe torchlean pytorch_roundtrip --model mlp --action import

# Exercise convolutional parameter layouts through the CNN
# exchange.
lake exe torchlean pytorch_roundtrip --model cnn --action export
python3 NN/Examples/Interop/PyTorch/CNN/train_cnn.py
lake exe torchlean pytorch_roundtrip --model cnn --action import

== CNN import example ==
Loaded: NN/Examples/Interop/PyTorch/CNN/cnn.json
Output (executable `nn` module on CPU, Float):
[0.999962, 0.500014]

# Exchange the encoder projections and normalization
# parameters.
lake exe torchlean pytorch_roundtrip --model transformer --action export
python3 NN/Examples/Interop/PyTorch/Transformer/train_transformer.py
lake exe torchlean pytorch_roundtrip --model transformer --action import

== Transformer import example ==
Loaded: NN/Examples/Interop/PyTorch/Transformer/transformer_encoder.json
Output (executable `nn` module on CPU, Float):
[[0.999998, -0.999998]]

true

true

some [1.000000, 3.000000]

some [1.000000, 2.000000]

PyTorch Code Generation🔗

The IR export example does the reverse kind of work. It lowers a TorchLean model to the shared IR and emits runnable PyTorch code for a curated set of architectures: linear, mlp, sum, autoencoder, mha, mha-mask, and transformer.

# Emit a standalone Python program for the seeded IR
# example.
lake exe torchlean torch_ir_pytorch --arch mlp > exported_model.py
python3 exported_model.py

The mlp architecture is three lines of TorchLean:

-- Fix the architecture before assigning seeded parameter
-- values.
def archMLP : nn.Builder (nn.Sequential [2] [1]) :=
  nn.Sequential![nn.linear 2 3, nn.relu, nn.linear 3 1]

and nn.build 0 archMLP turns that description into a model with concrete seeded parameters. The emitted Python has no builder and no seed; the parameters are baked in as literals, and the forward pass is the lowered IR written out one SSA binding at a time:

class TorchLeanMLP(nn.Module):
  def __init__(self):
    super().__init__()
    dtype = torch.float32
    self.const_2 = nn.Parameter(torch.tensor([
      float.fromhex('0x182a07c19a71ecp-54'),
      float.fromhex('0x1597e07783b3c0p-53'),
      float.fromhex('0x1f66b498260754p-53'),
      float.fromhex('0x1bf20b38c869acp-54'),
      float.fromhex('0x18795e185062c8p-53'),
      float.fromhex('-0x1ec171132474b8p-53')],
      dtype=dtype).reshape((3, 2)))
    self.register_buffer("const_5", torch.tensor(
      [0.0, 0.0, 0.0], dtype=dtype).reshape((3,)))
    self.const_11 = nn.Parameter(torch.tensor([
      float.fromhex('-0x10fbf4e603fdfdp-52'),
      float.fromhex('-0x170671f4bea0ecp-53'),
      float.fromhex('0x1accf23e183060p-53')],
      dtype=dtype).reshape((1, 3)))
    self.register_buffer("const_14", torch.tensor(
      [0.0], dtype=dtype).reshape((1,)))

  def forward(self, x):
    v0 = x
    v1 = v0.reshape((1, 2))
    v2 = self.const_2
    v3 = v2.transpose(0, 1)
    v4 = torch.matmul(v1, v3)
    v5 = self.const_5
    v6 = torch.broadcast_to(v5, (1, 3))
    v7 = v4 + v6
    v8 = v7.reshape((3,))
    v9 = torch.relu(v8)
    v10 = v9.reshape((1, 3))
    v11 = self.const_11
    v12 = v11.transpose(0, 1)
    v13 = torch.matmul(v10, v12)
    ...

The transpose(0, 1) calls are the orientation convention made explicit: TorchLean stores a linear layer's matrix as [out, in] and the emitter writes the transpose that x @ W.T needs, rather than silently reordering the numbers on the way out. The emitted file also trains itself when run, which exercises forward, backward, and an optimizer step. In the excerpt, weight constants become nn.Parameters while the zero biases become registered buffers. A matching forward value therefore does not establish that subsequent training updates the same set of parameters. A CPU run (CUDA_VISIBLE_DEVICES=\"\" python3 exported_model.py) reports:

loss 3.9950685501098633

The generated bindings let us follow one linear layer without guessing how Python interprets a matrix. The input becomes a single row, the stored [3, 2] weight is transposed to [2, 3], and matrix multiplication produces three hidden features. Broadcasting the [3] bias adds one value per feature before the row is reshaped back to a vector. Those transformations preserve the intended feature order while adapting it to PyTorch's matrix interface.

That printed loss comes from the emitted file's training example. To compare its initial forward pass with Lean, I'll use the same input and seeded parameters on both sides.

[[3, 2]:
 [[0.377565, 0.674790], [0.981287, 0.436648], [0.764815, -0.961113]],
 [3]:
 [0.000000, 0.000000, 0.000000],
 [1, 3]:
 [[-1.061513, -0.719537, 0.837518]],
 [1]:
 [0.000000]]

[-1.377817]

"0.000000"

"float.fromhex('0x112e0be826d695p-82')"

"0.500000"

# Keep the parameter payload fixed for the two arithmetic
# evaluations.
torch.manual_seed(7)
m = torch.nn.Sequential(torch.nn.Linear(2, 3), torch.nn.ReLU(), torch.nn.Linear(3, 1))
sd = m.state_dict()
payload = {"params": {
  "layers.0.weight": sd["0.weight"].tolist(),
  "layers.0.bias":   sd["0.bias"].tolist(),
  "layers.2.weight": sd["2.weight"].tolist(),
  "layers.2.bias":   sd["2.bias"].tolist()}}

float32 forward : -0.052048876881599426 0xbfaaa62680000000
float64 forward : -0.05204887814426984  0xbfaaa6268ad8a440

some [-0.052049]

some "bfaaa6268ad8a440"

Python Training And Parameter Export🔗

The training loop in train_mlp.py uses PyTorch's standard loss, backward pass, and optimizer:

# Clear each iteration's gradients before differentiating
# the current loss.
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.1)

x_train = torch.tensor([[0.5, 0.8]], dtype=torch.float32)
y_train = torch.tensor([[1.0]], dtype=torch.float32)

for epoch in range(200):
    optimizer.zero_grad()
    outputs = model(x_train)
    loss = criterion(outputs, y_train)
    loss.backward()
    optimizer.step()

The only TorchLean-aware part is the export, and it is a dictionary rebuild:

# Rename learned tensors into the exact keys accepted by the
# Lean family loader.
state_dict = model.state_dict()
new_format = {}
new_format['layers.0.weight'] = state_dict['fc1.weight'].tolist()
new_format['layers.0.bias'] = state_dict['fc1.bias'].tolist()
new_format['layers.2.weight'] = state_dict['fc2.weight'].tolist()
new_format['layers.2.bias'] = state_dict['fc2.bias'].tolist()

Those four lines implement the name mapping. A missing required key is rejected; a wrong mapping that still supplies every required key at the expected shape may be accepted. Keep numerical probes alongside the schema checks.

The optimizer loop clears accumulated gradients before each backward pass. Without that step, the next update would use contributions retained from previous iterations as well as the current loss. The export block runs after training and reads the resulting parameter values; it does not serialize the autograd graph that produced them. Renaming fc1.weight to layers.0.weight changes the transport key while preserving the nested numeric array. That explicit mapping is the place to review which learned tensor fills each Lean field.

Tensor Exchange Versus Model Import🔗

It is useful to distinguish three interop layers:

• tensor exchange moves arrays between frameworks;

• parameter import fills a known TorchLean model family with checked weights;

• semantic import claims that a foreign program has the same meaning as a TorchLean graph.

DLPack belongs mostly to the first layer: it is a standard in-memory tensor exchange format used by array and tensor libraries. It can help avoid extra copies, but it does not by itself tell Lean that a Python model has the same architecture, parameter names, or proof semantics.

TorchLean's current round trip implements the second layer for three known families: the selected loader checks required names and tensor shapes. Semantic import is stronger and requires an IR denotation and a proof relating the foreign program to it. The numeric comparison in the previous section checks one input; the verification chapters are where a claim about all inputs in a region gets established.

Interop Boundary Selection🔗

An imported model can subsequently be lowered to IR or exported again; these choices can be combined.

Guarantees And Limits🔗

The family importers return typed parameters for a forward pass. When preparing an exchange, keep these limits in view:

• arbitrary PyTorch training needs its own semantic or artifact bridge;

• the importer covers the supported artifact formats rather than the full PyTorch ecosystem;

• dtype metadata is validated when supplied, but architecture metadata, optimizer state, unlisted buffers, extra keys, and exact float bit patterns are not certified by these family loaders;

• malformed required tensors are rejected without a localized diagnostic, because the loader uses Option;

• emitted constants preserve finite binary64 values before the requested tensor dtype conversion;

• shape checking cannot distinguish two conventions that agree on shape, such as a transposed square matrix;

• binary32 behavior is handled by the floating point bridge rather than by the round trip format.

References🔗

• George et al. (2026) for the project overview and shared IR architecture.

• Paszke et al. (2019) for PyTorch's design, including state_dict conventions.

• Vaswani et al. (2017) for the encoder architecture used by the third family.

• PyTorch, serialization notes, for state_dict-style save and load workflows.

• PyTorch, nn.Module, for parameter naming and module conventions.

• DLPack documentation, for tensor exchange as distinct from model-family import.