ml_switcheroo.core.rewriter.decorators¶

Decorator Rewriting Mixin.

This module provides logic to handle Python decorators (e.g., @torch.jit.script). It enables: 1. Renaming: Mapping decorators between frameworks (e.g., @torch.jit.script -> @jax.jit). 2. Removal: stripping decorators that have no equivalent in the target framework (if mapped to null).

Logic robustness handles conflict with CallMixin. Since leave_Decorator runs after child nodes (the decorator expression) are visited, CallMixin might have already transformed calls inside the decorator (e.g. @torch.jit.script(optk=True)). To perform accurate lookups, we inspect the original_node (preserving Source semantics) before applying target logic.

Classes¶

DecoratorMixin

Mixin for transforming Decorator nodes.

Module Contents¶

class ml_switcheroo.core.rewriter.decorators.DecoratorMixin(semantics: ml_switcheroo.semantics.manager.SemanticsManager, config: ml_switcheroo.config.RuntimeConfig)¶

Bases: ml_switcheroo.core.rewriter.base.BaseRewriter

Mixin for transforming Decorator nodes. Part of PivotRewriter.

leave_Decorator(original_node: libcst.Decorator, updated_node: libcst.Decorator) → libcst.Decorator | libcst.RemovalSentinel¶

Processes decorators attached to functions or classes.

Logic: 1. Identifies the decorator name from original_node to ensure we key off the

Source Framework API, even if CallMixin modified the children in updated_node.

  1. Looks up the semantic definition.

  2. If the target variant is explicitly null (None), removes the decorator.

  3. If the target variant specifies a new API, rewrites the name.