Crate axonml_serialize 

Axonml Serialize - Model Serialization for Axonml ML Framework

This crate provides functionality for saving and loading trained models, including state dictionaries, model checkpoints, and format conversion.

Supported Formats

• Axonml Native (.axonml) - Efficient binary format
• JSON (.json) - Human-readable format for debugging
• SafeTensors (.safetensors) - Safe, fast format (optional feature)

Example

use axonml_serialize::{save_model, load_model, StateDict};

// Save model
save_model(&model, "model.axonml")?;

// Load model
let state_dict = load_state_dict("model.axonml")?;
model.load_state_dict(&state_dict)?;

@version 0.1.0 @author AutomataNexus Development Team

Structs

Checkpoint

A complete training checkpoint.

CheckpointBuilder

Builder for creating checkpoints.

StateDict

State dictionary for storing model parameters.

StateDictEntry

An entry in the state dictionary.

TensorData

Serializable tensor data.

TrainingState

Training state for checkpointing.

Enums

Format

Supported serialization formats.

OnnxOpType

ONNX operator type mapping.

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_checkpoint

Load a training checkpoint.

load_state_dict

Load a state dictionary from a file.

pytorch_layer_mapping

Map of PyTorch layer names to Axonml equivalents.

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 a Axonml key to PyTorch format.

transpose_linear_weights

Transpose weights if needed for format conversion.