ml_switcheroo.core.rewriter.attributes

Attribute Rewriting Logic.

Handles transformation of attributes (e.g. torch.float32), and tracks assignment logic to detect stateful variables (e.g. self.layer = …).

Classes

AttributeMixin

Mixin for transforming attributes and tracking assignments.

Module Contents

class ml_switcheroo.core.rewriter.attributes.AttributeMixin(semantics: ml_switcheroo.semantics.manager.SemanticsManager, config: ml_switcheroo.config.RuntimeConfig)

Bases: ml_switcheroo.core.rewriter.base.BaseRewriter

Mixin for transforming attributes and tracking assignments.

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

Track stateful assignments (e.g. self.layer = Linear(…)) to determine if an object variable requires special handling later.

leave_Attribute(original: libcst.Attribute, updated: libcst.Attribute) → libcst.BaseExpression

Visits attributes (e.g. torch.float32).