ml_switcheroo.plugins.method_property

Plugin for transforming Method calls to Property attributes.

This module handles cases where one framework uses a method calls (e.g. x.size()) while the target framework uses properties (e.g. x.shape).

Common usage: - PyTorch: x.size() -> JAX/NumPy: x.shape - PyTorch: x.size(dim) -> JAX/NumPy: x.shape[dim]

Functions

transform_method_to_property(→ Union[libcst.Attribute, ...)

Plugin Hook: Transforms a method call into an attribute access or subscript.

Module Contents

ml_switcheroo.plugins.method_property.transform_method_to_property(node: libcst.Call, ctx: ml_switcheroo.core.hooks.HookContext) libcst.Attribute | libcst.Subscript | libcst.Call

Plugin Hook: Transforms a method call into an attribute access or subscript.

Triggers:

Operations like size mapped with requires_plugin: “method_to_property”.

Parameters:
  • node – The original CST Call (e.g., x.size()).

  • ctx – HookContext for API lookup.

Returns:

cst.Attribute (x.shape), cst.Subscript (x.shape[0]), or original node.