ml_switcheroo.analysis.lifecycle

Lifecycle Analysis for Class Member Initialization.

This module provides the InitializationTracker, a static analysis tool that verifies that class members used in the forward pass (inference) are properly defined in the initialization phase (__init__).

This check helps detect: 1. Dynamic attribute definition (valid in Python, often invalid in static graph compilation like JAX/XLA). 2. Typos in member names between init and forward. 3. Implicit state that might be lost during transpilation if not explicitly declared.

Classes

InitializationTracker

Scans classes to ensure members used in forward are initialized in __init__.

Module Contents

class ml_switcheroo.analysis.lifecycle.InitializationTracker

Bases: libcst.CSTVisitor

Scans classes to ensure members used in forward are initialized in __init__.

It maintains a stack of Class Contexts to handle nested class definitions correctly.

warnings: List[str] = []

visit_ClassDef(node: libcst.ClassDef) → None

Enters a class definition. Pushes a new Context onto the stack.

Parameters:

node – The class definition node.

leave_ClassDef(node: libcst.ClassDef) → None

Exits a class definition and computes the difference between usages and inits. If discrepancies are found, they are recorded in self.warnings.

Parameters:

node – The class definition node.

visit_FunctionDef(node: libcst.FunctionDef) → None

Tracks entry into __init__ or forward/call methods. Sets context flags in_init or in_forward.

Parameters:

node – The function definition node.

leave_FunctionDef(node: libcst.FunctionDef) → None

Exits function scope. Resets context flags.

Parameters:

node – The function definition node.

visit_Assign(node: libcst.Assign) → None

Tracks assignments to self.x inside __init__.

Parameters:

node – The assignment node.

visit_AnnAssign(node: libcst.AnnAssign) → None

Tracks annotated assignments (self.x: int = …) inside __init__.

Parameters:

node – The annotated assignment node.

visit_Attribute(node: libcst.Attribute) → None

Tracks attribute access (self.x) inside forward.

Parameters:

node – The attribute access node.