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. Typoes 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 of discovered issues (e.g. “Member ‘conv1’ used but not initialized”).

Type:

List[str]

_scope_stack

Stack tracking nested class contexts.

Type:

List[_ClassContext]

warnings: List[str] = []

visit_ClassDef(node: libcst.ClassDef) → None

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

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.

visit_FunctionDef(node: libcst.FunctionDef) → None

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

leave_FunctionDef(node: libcst.FunctionDef) → None

Exits function scope. Resets context flags.

visit_Assign(node: libcst.Assign) → None

Tracks assignments to self.x inside __init__.

visit_AnnAssign(node: libcst.AnnAssign) → None

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

visit_Attribute(node: libcst.Attribute) → None

Tracks attribute access (self.x) inside forward.