ml_switcheroo.core.engine¶

Orchestration Engine for AST Transformations.

This module provides the ASTEngine, the primary driver for the transpilation process. It coordinates the various analysis and transformation passes required to convert code from a source framework to a target framework.

The Engine pipeline consists of:

  1. Preprocessing: Parsing source code into a LibCST tree.

  2. Safety Analysis:

    • Purity: Checking for side effects if targeting functional frameworks (JAX).

    • Lifecycle: Verifying variable initialization.

    • Dependencies: Scanning for unmapped 3rd-party imports.

  3. Semantic Pivoting: Executing the PivotRewriter to map API calls and structure.

  4. Post-processing:

    • Import Fixer: Injecting necessary imports and removing unused source imports.

    • Structural Linting: Verifying no artifacts from the source framework remain.

The engine relies on RuntimeConfig to resolve ‘Flavours’ (e.g., distinguishing flax_nnx from generic jax) to load the correct structural traits.

Classes¶

ConversionResult

Structured result of a single file conversion.

ASTEngine

The main compilation unit.

Module Contents¶

class ml_switcheroo.core.engine.ConversionResult(/, **data: Any)¶

Bases: pydantic.BaseModel

Structured result of a single file conversion.

code: str = None¶
errors: List[str] = None¶
success: bool = None¶
trace_events: List[Dict[str, Any]] = None¶
property has_errors: bool¶

Returns True if any errors or warnings were recorded during conversion.

Returns:

True if errors list is non-empty.

Return type:

bool

class ml_switcheroo.core.engine.ASTEngine(semantics: ml_switcheroo.semantics.manager.SemanticsManager = None, config: ml_switcheroo.config.RuntimeConfig | None = None, source: str = 'torch', target: str = 'jax', strict_mode: bool = False, plugin_config: Dict[str, Any] | None = None)¶

The main compilation unit.

This class encapsulates the state and logic required to transpile a single unit of code. It manages the lifecycle of the LibCST tree and the invocation of visitor passes.

semantics¶
source¶
target¶
strict_mode¶
parse(code: str) → libcst.Module¶

Parses source string into a LibCST Module.

Parameters:

code (str) – Python source code.

Returns:

The parsed Abstract Syntax Tree.

Return type:

cst.Module

Raises:

libcst.ParserSyntaxError – If the input code is invalid Python.

to_source(tree: libcst.Module) → str¶

Converts CST back to source string.

Parameters:

tree (cst.Module) – The modified syntax tree.

Returns:

Generated Python code.

Return type:

str

run(code: str) → ConversionResult¶

Executes the full transpilation pipeline.

Passes performed: 1. Parse. 2. Purity Scan (if targeting JAX-like frameworks). 3. Lifecycle Analysis (Init/Forward mismatch). 4. Dependency Scan (Checking 3rd party libs). 5. Pivot Rewrite (The main transformation). 6. Import Fixer (Injecting new imports, pruning old ones). 7. Structural Linting (Verifying output cleanliness).

Parameters:

code (str) – The input source string.

Returns:

Object containing transformed code and error logs.

Return type:

ConversionResult