ml_switcheroo.core.rewriter.calls.mixer
=======================================

.. py:module:: ml_switcheroo.core.rewriter.calls.mixer

.. autoapi-nested-parse::

   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
-------

.. autoapisummary::

   ml_switcheroo.core.rewriter.calls.mixer.CallMixin


Module Contents
---------------

.. py:class:: CallMixin(semantics: ml_switcheroo.semantics.manager.SemanticsManager, config: ml_switcheroo.config.RuntimeConfig)

   Bases: :py:obj:`ml_switcheroo.core.rewriter.normalization.NormalizationMixin`, :py:obj:`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).


   .. py:method:: 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)



   .. py:method:: leave_Call(original: libcst.Call, updated: libcst.Call) -> Union[libcst.Call, libcst.BinaryOperation, libcst.UnaryOperation, libcst.CSTNode]

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