ml_switcheroo.analysis.audit

Audit functionality to determine coverage gaps for a source codebase.

This module provides the CoverageScanner, a static analysis visitor that inspects source code to identify API usage (calls and attributes) relative to specific frameworks (e.g., PyTorch, JAX). It resolves import aliases to fully qualified names (FQNs) and checks them against the SemanticsManager to determine if they are currently supported by the transpiler.

Classes

CoverageScanner

Scans a file to identify API calls and checks if they exist in the Semantics Manager.

Module Contents

class ml_switcheroo.analysis.audit.CoverageScanner(semantics: ml_switcheroo.semantics.manager.SemanticsManager, allowed_roots: Set[str])

Bases: libcst.CSTVisitor

Scans a file to identify API calls and checks if they exist in the Semantics Manager.

It tracks import aliases to correctly resolve calls like jnp.sum back to jax.numpy.sum before querying the knowledge base.

results: Dict[str, Tuple[bool, str]]

Dictionary mapping Fully Qualified Names (FQN) to a tuple of (is_supported, framework_key).

semantics

allowed_roots

visit_Import(node: libcst.Import) → None

Visits import … statements to populate the alias map.

Parameters:

node – The CST import node.

visit_ImportFrom(node: libcst.ImportFrom) → None

Visits from … import … statements to populate the alias map.

Parameters:

node – The CST import-from node.

visit_Call(node: libcst.Call) → None

Visits function call nodes to check the invoked function.

Parameters:

node – The CST call node.

visit_Attribute(node: libcst.Attribute) → None

Visits attributes to check for framework constants.

This allows detecting usage like torch.float32 which are not calls but still require support. Checks are idempotent via the results dictionary.

Parameters:

node – The CST attribute node.