ml_switcheroo¶

ml-switcheroo Package.

A deterministic AST transpiler for converting Deep Learning models between frameworks (e.g., PyTorch -> JAX/Flax).

This package exposes the core compilation engine and configuration utilities for programmatic usage.

Usage:

Simple String Conversion:

import ml_switcheroo as mls
# Auto-detects defaults based on installed libraries (e.g. Torch -> JAX)
code = "y = torch.abs(x)"
result = mls.convert(code)
print(result)

# Explicit conversion
result = mls.convert(code, source="torch", target="jax")

Advanced Usage (AST Engine):

from ml_switcheroo import ASTEngine, RuntimeConfig

config = RuntimeConfig(source_framework="torch", target_framework="jax", strict_mode=True)
engine = ASTEngine(config=config)
res = engine.run("y = torch.abs(x)")

if res.success:
    print(res.code)
else:
    print(f"Errors: {res.errors}")

Submodules¶

Attributes¶

__version__

Functions¶

convert(→ str)

Transpiles a string of Python code from one framework to another.

Package Contents¶

ml_switcheroo.__version__ = '0.0.1'¶

ml_switcheroo.convert(code: str, source: str | None = None, target: str | None = None, strict: bool = False, plugin_settings: Dict[str, Any] | None = None, semantics: semantics.manager.SemanticsManager | None = None) → str[source]¶

Transpiles a string of Python code from one framework to another.

This is a high-level convenience wrapper around the ASTEngine. For file-based conversions or batch processing, consider using ml_switcheroo.cli or using ASTEngine directly.

Parameters:

code – The source code to convert.
source – The source framework key (e.g., “torch”). If None, determined dynamically by priority via RuntimeConfig.
target – The target framework key (e.g., “jax”). If None, determined dynamically by priority via RuntimeConfig.
strict – If True, the engine will return an error if an API cannot be mapped. If False (default), the original code is preserved wrapped in escape hatch comments.
plugin_settings – Specific configuration flags passed to plugin hooks (e.g., {“rng_arg_name”: “seed”}).
semantics – An existing Knowledge Base instance. If None, a new one is initialized from disk.

Returns:

The transpiled source code string.

Raises:

ValueError – If the conversion fails (e.g. syntax errors or strict mode violations).