ml_switcheroo.cli.handlers.convert_weights

Weight Migration Script Generation.

This module provides the WeightScriptGenerator, a utility that creates standalone Python scripts for migrating model checkpoints (weights) between deep learning frameworks.

It analyzes the source AST to identify model layers, queries the Semantic Knowledge Base for parameter mapping rules (e.g., weight -> kernel), and determines tensor layout permutations (e.g., NCHW -> NHWC).

Supported Directions: - PyTorch -> JAX (Flax) - JAX (Flax) -> PyTorch

Unsupported combinations degrade gracefully by returning False.

Classes

WeightScriptGenerator

Generates a Python script to migrate weights between frameworks.

Module Contents

class ml_switcheroo.cli.handlers.convert_weights.WeightScriptGenerator(semantics: ml_switcheroo.semantics.manager.SemanticsManager, config: ml_switcheroo.config.RuntimeConfig)

Generates a Python script to migrate weights between frameworks.

semantics

source_fw

target_fw

source_adapter

target_adapter

generate(source_path: pathlib.Path, output_script: pathlib.Path) → bool

Main entry point to generate the migration script.

Parameters:

• source_path – Path to the Python file containing the source model definition.

• output_script – Path where the generated script will be written.

Returns:

True if generation was successful.

Return type:

bool