ml_switcheroo.plugins.context_to_function_wrap

Plugin for handling Context Manager rewriting.

This module addresses the impedance mismatch between PyTorch’s global state context managers (like torch.no_grad()) and JAX’s functional, explicit gradient handling.

It provides a transformation hook that:

1. Detects usage of context managers flagged with context_to_function_wrap.

2. Injects a nullcontext shim into the function preamble.

3. Rewrites the specific API call to use this shim, ensuring the with …: block remains valid Python syntax while effectively disabling gradient tracking semantics.

Functions

transform_context_manager(→ libcst.Call)

Plugin Hook: Transforms valid Source context managers into JAX-compatible shims.

Module Contents

ml_switcheroo.plugins.context_to_function_wrap.transform_context_manager(node: libcst.Call, ctx: ml_switcheroo.core.hooks.HookContext) → libcst.Call

Plugin Hook: Transforms valid Source context managers into JAX-compatible shims.

Triggers

• Operations marked with requires_plugin: "context_to_function_wrap" in Semantic JSONs.

• Primarily targets torch.no_grad and torch.enable_grad.

Parameters:

• node – The original CST Call node (e.g., torch.no_grad()).

• ctx – The HookContext providing injection capabilities.

Returns:

The transformed CST Call node pointing to the injected shim.