ml_switcheroo.plugins.device_allocator

Plugin for translating Device Allocation logic.

Delegates syntax generation to the target FrameworkAdapter to ensure logic remains decoupled from the core transpiler.

Supported: - Extracts device type and index from torch.device(…). - Handles cuda:0 string parsing. - Calls adapter.get_device_syntax() for target code generation.

Functions

transform_device_allocator(→ libcst.BaseExpression)

Plugin Hook: Transforms device construction calls via Adapter delegation.

Module Contents

ml_switcheroo.plugins.device_allocator.transform_device_allocator(node: libcst.Call, ctx: ml_switcheroo.core.hooks.HookContext) → libcst.BaseExpression

Plugin Hook: Transforms device construction calls via Adapter delegation.

Triggers:

Operations marked with requires_plugin: “device_allocator” (e.g., torch.device).

Parameters:

  • node – The original CST Call node.

  • ctx – HookContext for target framework access.

Returns:

A CST Expression representing the target device access.