ml_switcheroo.plugins.data_loader

Plugin for transforming Data Loaders to Generic Shim.

Handles the mapping of torch.utils.data.DataLoader (or similar iterators) to the GenericDataLoader shim.

This plugin is blindly executed whenever the Semantic Knowledge Base maps an operation to requires_plugin: “convert_dataloader”. It does not check the target framework name.

Implementation Details:

  • Filters arguments to ensure compatibility with the Generic Shim.

  • Explicitly handles num_workers, pin_memory, and drop_last by mapping names and passing them (since the Shim accepts them as optional kwargs).

  • Injects the Shim class definition at the top of the file on first use via ctx.inject_preamble.

Functions

get_shim_code(→ str)

Returns the source code for the GenericDataLoader class.

transform_dataloader(→ libcst.CSTNode)

Middleware to rewrite DataLoader instantiation.

Module Contents

ml_switcheroo.plugins.data_loader.get_shim_code() → str[source]

Returns the source code for the GenericDataLoader class.

This code string is injected into the preamble of generated files by the convert_dataloader plugin when a target framework maps DataLoader to GenericDataLoader.

Capabilities: - Basic Batching - Shuffling - Dimension preservation - Validation stubbing (ignoring num_workers/pin_memory safely)

Returns:

Python source code for the shim class.

Return type:

str

ml_switcheroo.plugins.data_loader.transform_dataloader(node: libcst.Call, ctx: ml_switcheroo.core.hooks.HookContext) → libcst.CSTNode[source]

Middleware to rewrite DataLoader instantiation.

Triggers

  • Operations marked with requires_plugin: "convert_dataloader" in JSON.

Actions

  • Injects GenericDataLoader class definition into preamble (idempotent).

  • Rewrites function call to GenericDataLoader(…)

  • Filters arguments (Dataset as pos 0, preserves supported kwargs).