ml_switcheroo.core.rewriter.passes.auxiliary¶
Auxiliary Logic Pass.
This module consolidates transformations that are adjacent to the core API logic, specifically: 1. Decorators: Rewriting or stripping function decorators (e.g., @torch.jit.script). 2. Control Flow: Analyzing and transforming loops (for, while) to ensure safety
in functional frameworks (e.g., JAX/XLA compatibility checks).
It merges logic previously found in decorators.py and control_flow.py.
Classes¶
Pass dealing with auxiliary syntax constructs: decorators and control flow. |
|
LibCST Transformer for auxiliary constructs. |
Module Contents¶
- class ml_switcheroo.core.rewriter.passes.auxiliary.AuxiliaryPass[source]¶
Bases:
ml_switcheroo.core.rewriter.interface.RewriterPassPass dealing with auxiliary syntax constructs: decorators and control flow.
- transform(module: libcst.Module, context: ml_switcheroo.core.rewriter.context.RewriterContext) libcst.Module[source]¶
Executes the auxiliary transformation logic.
- Parameters:
module – The source CST.
context – Shared state.
- Returns:
The transformed CST.
- class ml_switcheroo.core.rewriter.passes.auxiliary.AuxiliaryTransformer(context: ml_switcheroo.core.rewriter.context.RewriterContext)[source]¶
Bases:
libcst.CSTTransformerLibCST Transformer for auxiliary constructs.
- context¶
- visit_SimpleStatementLine(node: libcst.SimpleStatementLine) bool | None[source]¶
Reset statement buffers.
- leave_SimpleStatementLine(original_node: libcst.SimpleStatementLine, updated_node: libcst.SimpleStatementLine) libcst.SimpleStatementLine | libcst.FlattenSentinel[source]¶
Process statement errors.
- leave_Decorator(original_node: libcst.Decorator, updated_node: libcst.Decorator) libcst.Decorator | libcst.RemovalSentinel[source]¶
Rewrites decorators.
Logic: 1. Resolve decorator name (e.g. torch.jit.script). 2. Lookup semantics. 3. If target variant is None -> Remove. 4. If target variant has API -> Rename.