ml_switcheroo.core.rewriter.calls.mixer¶
CallMixin Definition.
The primary visitor logic for handling Call and Assign nodes in the AST. This module orchestrates the transformation pipeline by delegating specific tasks (dispatch, reshaping, lifecycle management) to helper modules.
Classes¶
Mixin for transforming Call nodes and unpacking Assignments. |
Module Contents¶
- class ml_switcheroo.core.rewriter.calls.mixer.CallMixin(semantics: ml_switcheroo.semantics.manager.SemanticsManager, config: ml_switcheroo.config.RuntimeConfig)¶
Bases:
ml_switcheroo.core.rewriter.normalization.NormalizationMixin,ml_switcheroo.core.rewriter.base.BaseRewriterMixin 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.