ml_switcheroo.core.rewriter.structure_class

Class Structure Rewriting Logic.

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

Classes

ClassStructureMixin

Mixin for transforming ClassDef nodes (Logic 1).

Module Contents

class ml_switcheroo.core.rewriter.structure_class.ClassStructureMixin(semantics: ml_switcheroo.semantics.manager.SemanticsManager, config: ml_switcheroo.config.RuntimeConfig)

Bases: ml_switcheroo.core.rewriter.base.BaseRewriter

Mixin for transforming ClassDef nodes (Logic 1).

visit_ClassDef(node: libcst.ClassDef) → bool | None

Detects if we are entering a generic ML Module class.

leave_ClassDef(original_node: libcst.ClassDef, updated_node: libcst.ClassDef) → libcst.CSTNode

Performs the Base Class Swap logic.