ml_switcheroo.core.rewriter.control_flow¶

Control Flow Rewriting Logic.

Handles structural control flow transformations, specifically for loops. Unlike API calls which map via Semantics JSON, control flow hooks are dispatched via reserved system triggers (e.g. transform_for_loop).

This logic enables a prioritization chain for loop handling:

  1. Static Unrolling: Optimization strategy. If a loop is determined to be static (e.g., range(5)), it is unrolled into flat code.

  2. Safety Transformation: Compliance strategy. If the loop remains dynamic and the target framework (like JAX) enforces functional purity, this ensures the loop is flagged or converted to a scan/while construct.

Classes¶

ControlFlowMixin

Mixin for visiting Control Flow nodes (For, While, If).

Module Contents¶

class ml_switcheroo.core.rewriter.control_flow.ControlFlowMixin(semantics: ml_switcheroo.semantics.manager.SemanticsManager, config: ml_switcheroo.config.RuntimeConfig)¶

Bases: ml_switcheroo.core.rewriter.base.BaseRewriter

Mixin for visiting Control Flow nodes (For, While, If).

leave_For(original_node: libcst.For, updated_node: libcst.For) → libcst.For | libcst.CSTNode¶

Invokes loop transformation logic.

Implements a priority chain: 1. Static Unroll: Checks transform_for_loop_static. If the loop index

is constant, unrolls it (Optimization).

  1. General Transform: Checks transform_for_loop. Handles general case logic or applies safety warnings (e.g., Escape Hatch for JAX).

Parameters:

  • original_node (cst.For) – The node before transformation.

  • updated_node (cst.For) – The node after child visitors have run.

Returns:

The transformed node (potentially a FlattenSentinel of unrolled statements) or the original node if untouched.

Return type:

Union[cst.For, cst.CSTNode]