ml_switcheroo.core.rewriter.structure_class
===========================================

.. py:module:: ml_switcheroo.core.rewriter.structure_class

.. autoapi-nested-parse::

   Class Structure Rewriting Logic.

   Handles transformation of Class Definitions.
   Logic 1: Swaps Module Base (e.g. `torch.nn.Module` -> `flax.nnx.Module`).



Classes
-------

.. autoapisummary::

   ml_switcheroo.core.rewriter.structure_class.ClassStructureMixin


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

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

   Bases: :py:obj:`ml_switcheroo.core.rewriter.base.BaseRewriter`


   Mixin for transforming ClassDef nodes (Logic 1).


   .. py:method:: visit_ClassDef(node: libcst.ClassDef) -> Optional[bool]

      Detects if we are entering a generic ML Module class.



   .. py:method:: leave_ClassDef(original_node: libcst.ClassDef, updated_node: libcst.ClassDef) -> libcst.CSTNode

      Performs the Base Class Swap logic.