ml_switcheroo.compiler.backends.python_snippet¶

Python Snippet Emitter.

This module provides the PythonSnippetEmitter, a specialized backend component designed to generate atomic Python statements (CST nodes) from LogicalNodes.

Unlike PythonBackend which synthesizes entire modules/classes, this emitter focuses on generating individual: 1. Initialization Statements: self.layer = LayerClass(…) 2. Execution Statements: output = self.layer(input) 3. Expressions: self.layer(input) or func(input)

This is used by the GraphPatcher to surgically insert new code for nodes created during Graph Optimization (e.g. Fused Operations).

Classes¶

PythonSnippetEmitter

Generates isolated Python statements from LogicalNodes.

Module Contents¶

class ml_switcheroo.compiler.backends.python_snippet.PythonSnippetEmitter(framework: str = 'torch')¶

Generates isolated Python statements from LogicalNodes.

Adapts generation logic based on the target framework to ensure idiomatic code (e.g., injecting rngs for Flax, nn. prefix for Torch).

framework = 'torch'¶
emit_init(node: ml_switcheroo.compiler.ir.LogicalNode) → libcst.SimpleStatementLine¶

Generates the initialization statement for a stateful layer.

Parameters:

node – The logical node describing the layer.

Returns:

A LibCST statement node representing self.id = Kind(…).

emit_call(node: ml_switcheroo.compiler.ir.LogicalNode, input_vars: List[str], output_var: str) → libcst.SimpleStatementLine¶

Generates the execution call statement (Assignment).

Parameters:

  • node – The logical node.

  • input_vars – List of variable names to pass as arguments.

  • output_var – Variable name to assign the result to.

Returns:

A LibCST statement node (Assign).

emit_expression(node: ml_switcheroo.compiler.ir.LogicalNode, input_vars: List[str]) → libcst.BaseExpression¶

Generates the function call expression (without assignment).

Parameters:

  • node – The logical node.

  • input_vars – List of input variable names.

Returns:

A LibCST expression node.