Crate axonml_serialize 

Model serialization for AxonML.

Save/load in .axonml (binary), .json, and SafeTensors formats. StateDict (PyTorch-compatible named parameter collection), Checkpoint with TrainingState (loss/val-loss/metric history, epoch/step, config map), format auto-detection (by extension and magic bytes), PyTorch weight conversion (convert_pytorch_key, convert_full_dict, linear weight transposition), metadata embedding, ONNX shape utilities.

File

crates/axonml-serialize/src/lib.rs

Author

Andrew Jewell Sr. — AutomataNexus LLC ORCID: 0009-0005-2158-7060

Updated

April 14, 2026 11:15 PM EST

Disclaimer

Use at own risk. This software is provided “as is”, without warranty of any kind, express or implied. The author and AutomataNexus shall not be held liable for any damages arising from the use of this software.

Structs

BundleHeader

Lightweight header decoded before the weights blob. Kept in sync with the Python converter’s parse_axonml header JSON (must round-trip verbatim).

Checkpoint

A complete training checkpoint.

CheckpointBuilder

Builder for creating checkpoints.

ModelBundle

Full bundle payload carrying everything the ONNX converter needs.

StateDict

State dictionary for storing model parameters.

StateDictEntry

An entry in the state dictionary.

TensorData

Serializable tensor data.

TrainingState

Training state for checkpointing.

Enums

BundleError

Errors produced by bundle save/load.

Format

Supported serialization formats.

OnnxOpType

ONNX operator type mapping.

Constants

AXONML_BUNDLE_VERSION

Current bundle format version.

AXONML_MAGIC

Magic bytes at the start of every .axonml bundle file.

Functions

convert_from_pytorch

Convert a state dict from PyTorch naming conventions.

detect_format

Detect the format from a file path based on extension.

detect_format_from_bytes

Detect format from file contents (magic bytes).

from_onnx_shape

Convert from ONNX shape (handling -1 for dynamic dimensions).

from_pytorch_key

Convert a PyTorch-style key to Axonml format.

load_bundle

Read a bundle from disk, returning both the header (decoded eagerly) and the full payload.

load_bundle_from_bytes

In-memory variant of load_bundle — useful for HTTP handlers or tests.

load_checkpoint

Load a training checkpoint.

load_header

Decode only the header of a .axonml file without reading the weights blob.

load_model

Load a model’s parameters from a saved state dictionary file.

load_state_dict

Load a state dictionary from a file.

pytorch_layer_mapping

Map of PyTorch layer names to Axonml equivalents.

save_bundle

Write a bundle to disk in the .axonml container format.

save_checkpoint

Save a complete training checkpoint.

save_model

Save a model’s state dictionary to a file.

save_state_dict

Save a state dictionary to a file with specified format.

to_onnx_shape

Convert a shape to ONNX format (with batch dimension handling).

to_pytorch_key

Convert an AxonML key to PyTorch format.

transpose_linear_weights

Transpose weights if needed for format conversion.

Type Aliases

BundleResult

Convenience alias for bundle operation results.