ml_switcheroo.core.rewriter.calls

Calls Rewriter Package.

This package handles the transformation of function calls, assignment unwrapping, and complex API dispatching logic (infix operators, conditional rules).

Submodules

• ml_switcheroo.core.rewriter.calls.dispatch
• ml_switcheroo.core.rewriter.calls.mixer
• ml_switcheroo.core.rewriter.calls.transformers
• ml_switcheroo.core.rewriter.calls.utils

Classes

CallMixin

Mixin for transforming Call nodes and unpacking Assignments.

Package Contents

class ml_switcheroo.core.rewriter.calls.CallMixin(semantics: ml_switcheroo.semantics.manager.SemanticsManager, config: ml_switcheroo.config.RuntimeConfig)

Bases: ml_switcheroo.core.rewriter.normalization.NormalizationMixin, ml_switcheroo.core.rewriter.base.BaseRewriter

Mixin for transforming Call nodes and unpacking Assignments.

Responsible for: 1. Handling functional apply patterns (Flax). 2. Lifecycle method stripping (.to(), .cuda()). 3. Plugin dispatch. 4. Standard API pivoting (Lookup -> Normalize -> Rewrite). 5. Output Transformation (Indexing, Casting).

leave_Assign(original_node: libcst.Assign, updated_node: libcst.Assign) → libcst.Assign

Handles assignment unwrapping for Functional -> OOP transitions.

Scenario: y, updates = layer.apply(vars, x) Target: y = layer(x) (NNX/Torch style)

leave_Call(original: libcst.Call, updated: libcst.Call) → libcst.Call | libcst.BinaryOperation | libcst.UnaryOperation | libcst.CSTNode

Rewrites function calls with detailed Trace Logging. Implements Logic 4: Layer Init State Threading.