Import Core

PyTorch import core (JSON parsing).

The Python side of TorchLean round-trips usually writes a JSON object containing nested arrays of floats (a Lean-readable projection of a PyTorch state_dict). Callers supply the Python script that loads .pt / .pth checkpoints and writes this JSON; the export helpers save and load PyTorch state dictionaries without converting them to JSON.

Design note:

In typical PyTorch workflows, weights are often serialized via torch.save(model.state_dict(), ...) or related checkpoint wrappers. TorchLean avoids parsing those PyTorch binary formats directly in Lean. Instead, PyTorch loads the checkpoint and emits a small JSON representation that is easy to validate against a Lean Shape and easy to diff in tests.

This importer is about weights only. Importing a captured graph (e.g. ONNX or torch.export) is a separate problem and lives at a different abstraction layer than this JSON state_dict adapter.

This module is where we keep the shared logic that most PyTorch → TorchLean importers need:

• parse nested JSON arrays into shape-checked Tensor Float s,
• handle a small amount of “state_dict ergonomics” (key lookup, optional wrappers, index parsing),
• keep everything model-agnostic, so the model-specific code can stay small and readable.

The public helpers are organized as follows:

• parseTensor is the core JSON-to-tensor conversion.
• loadWeights? and unwrapParams handle the two JSON layouts we accept.
• getTensor? and getTensorFirst? are the main lookup helpers used by the model-specific importers.

@[reducible, inline]

abbrev Import.PyTorch.StateDict : Type

A PyTorch-style state_dict encoded as a JSON object.

def Import.PyTorch.parseTensor (s : Spec.Shape) : Lean.Json → Option (TorchLean.Tensor Float s)

Parse a JSON value into a Tensor Float s.

The JSON encoding follows the tensor shape:

• scalars are JSON numbers,
• Shape.dim n s is a JSON array of length n whose entries recursively encode s.

If the JSON payload does not match the expected shape, we return none.

state_dict helpers

We use JSON objects keyed by strings because that mirrors PyTorch’s state_dict convention. Some TorchLean Python scripts wrap the object as { "params": { ... } }; loadWeights? accepts both formats.

def Import.PyTorch.unwrapParams (o : StateDict) : StateDict

If the object contains a "params" field that is itself an object, unwrap it.

We also merge any other top-level fields (e.g. "meta") into the returned dictionary so model importers can still read them. Parameter entries take precedence: wrapper metadata must never replace a tensor with the same key.

def Import.PyTorch.supportedDType : Lean.Json → Bool

Whether a JSON string names the float32 format accepted by this importer.

def Import.PyTorch.supportedWrapperDTypes (o : StateDict) : Bool

Whether optional state-dict metadata describes the numeric format supported by this importer.

Per-parameter metadata records one object per parameter. Older checked-in examples use one model-level "dtype": "float32" entry instead. Both formats are accepted, but an explicit bf16, float64, or integer dtype is rejected rather than silently reinterpreted as Tensor Float.

def Import.PyTorch.loadWeights? (j : Lean.Json) : Option StateDict

Load weights from JSON, accepting either:

• { ...state_dict... }, or
• { "params": { ...state_dict... } }.

def Import.PyTorch.getTensor? (o : StateDict) (k : String) (s : Spec.Shape) : Option (TorchLean.Tensor Float s)

Look up a key and parse it as a tensor of a given expected shape.

This is the helper most model-specific importers use to keep the “key wiring” readable.

def Import.PyTorch.getTensorFirst? (o : StateDict) (keys : List String) (s : Spec.Shape) : Option (TorchLean.Tensor Float s)

Try a list of state-dict keys in order and return the first tensor matching the expected shape.

This supports importers that accept both a compact interchange name and a framework-native module path for the same parameter.

Small parsing helpers used by shape-inferring importers

Some importers allow variable-width stacks (e.g. a PINN that learns its hidden widths from the checkpoint). Those cases need a little help to infer indices and matrix dimensions from JSON.

def Import.PyTorch.parseIndexedKey (pref suff key : String) : Option ℕ

Parse keys of the form prefix ++ <nat> ++ suffix.

Example: parseIndexedKey "layers." ".weight" "layers.3.weight" = some 3.

def Import.PyTorch.inferMatrixDims : Lean.Json → Option (ℕ × ℕ)

Infer (rows, cols) for a JSON matrix encoded as an array of arrays.

This helper infers dimensions from the outer length and first row length. Call sites that need stronger validation (all rows same length) should add an explicit check.

Convenience parsers for function-based constructors

Some call sites build tensors via Spec.vector_tensor / Spec.matrix_tensor, whose inputs are functions (Fin n → Float and Fin m → Fin n → Float).

parseFloatVec and parseFloatMatrix keep those call sites readable without duplicating JSON parsing logic outside this core module.

def Import.PyTorch.parseFloatVec (n : ℕ) (j : Lean.Json) : Option (Fin n → Float)

Parse a JSON array into a Fin n → Float function.

def Import.PyTorch.parseFloatMatrix (rows cols : ℕ) (j : Lean.Json) : Option (Fin rows → Fin cols → Float)

Parse a JSON matrix into a Fin rows → Fin cols → Float function.