ml_switcheroo.core.dsl

Operation Definition Language (ODL) Schema.

This module defines the Pydantic models used to validate the input (e.g. YAML) for defining new operations. It represents the source of truth for generating new operation definitions in the codebase, enforcing schema validity before AST injection occurs.

It defines the structure for: - Parameters (types, constraints, shapes, dtypes). - Variants (framework implementations). - Plugins (scaffolding logic). - Operations (abstract definitions).

Classes

OpType	Classification of the operation's syntactic usage.
ParameterDef	Definition of a single argument in an abstract operation signature.
Rule	Declarative rule for conditional logic within a variant or plugin.
ImportReq	Structured definition for a required import, supporting aliasing.
FrameworkVariant	Configuration for how a specific framework implements an operation.
PluginType	Enumeration of supported plugin structures to generate.
PluginScaffoldDef	Metadata for generating a new Python plugin file.
OperationDef	Top-level definition of a new Abstract Operation.

Module Contents

class ml_switcheroo.core.dsl.OpType

Bases: str, enum.Enum

Classification of the operation’s syntactic usage.

FUNCTION = 'function'

CONTEXT = 'context'

DECORATOR = 'decorator'

class ml_switcheroo.core.dsl.ParameterDef(/, **data: Any)

Bases: pydantic.BaseModel

Definition of a single argument in an abstract operation signature.

Updated to support semantic constraints (min, max, options), tensor properties (rank, dtype, shape), and function signature attributes (variadic, kind) to guide fuzzing and validation logic.

name: str = None

type: str | None = None

doc: str | None = None

default: str | None = None

min: float | None = None

max: float | None = None

options: List[str | int | float | bool] | None = None

rank: int | None = None

dtype: str | None = None

shape_spec: str | None = None

is_variadic: bool = None

kind: str = None

class ml_switcheroo.core.dsl.Rule(/, **data: Any)

Bases: pydantic.BaseModel

Declarative rule for conditional logic within a variant or plugin. Evaluated at runtime to dynamically switch APIs.

model_config

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

if_arg: str = None

op: ml_switcheroo.enums.LogicOp = None

is_val: str | int | float | bool | List[str | int | float] = None

use_api: str = None

class ml_switcheroo.core.dsl.ImportReq(/, **data: Any)

Bases: pydantic.BaseModel

Structured definition for a required import, supporting aliasing.

module: str = None

alias: str | None = None

class ml_switcheroo.core.dsl.FrameworkVariant(/, **data: Any)

Bases: pydantic.BaseModel

Configuration for how a specific framework implements an operation.

Defines the API endpoint, argument mappings, required imports, and versioning constraints.

api: str | None = None

args: Dict[str, str] | None = None

arg_values: Dict[str, Dict[str, str]] | None = None

kwargs_map: Dict[str, str] | None = None

inject_args: Dict[str, str | int | float | bool] | None = None

casts: Dict[str, str] | None = None

requires_plugin: str | None = None

dispatch_rules: List[Rule] = None

pack_to_tuple: str | None = None

layout_map: Dict[str, str] | None = None

transformation_type: str | None = None

operator: str | None = None

macro_template: str | None = None

output_select_index: int | None = None

output_adapter: str | None = None

output_cast: str | None = None

min_version: str | None = None

max_version: str | None = None

required_imports: List[str | ImportReq] = None

missing_message: str | None = None

class ml_switcheroo.core.dsl.PluginType

Bases: str, enum.Enum

Enumeration of supported plugin structures to generate.

CALL = 'call_transform'

BLOCK = 'block_transform'

class ml_switcheroo.core.dsl.PluginScaffoldDef(/, **data: Any)

Bases: pydantic.BaseModel

Metadata for generating a new Python plugin file.

name: str = None

type: PluginType = None

doc: str = None

rules: List[Rule] = None

class ml_switcheroo.core.dsl.OperationDef(/, **data: Any)

Bases: pydantic.BaseModel

Top-level definition of a new Abstract Operation.

Contains the semantic signature, documentation, framework implementations, and metadata for verification, fuzzing, and documentation generation.

operation: str = None

description: str = None

op_type: OpType = None

std_args: List[str | Tuple[str, str] | ParameterDef | Dict[str, Any]] = None

variants: Dict[str, FrameworkVariant | None] = None

scaffold_plugins: List[PluginScaffoldDef] = None

return_type: str | None = None

is_inplace: bool = None

test_rtol: float = None

test_atol: float = None

skip_fuzzing: bool = None

nondeterministic: bool = None

differentiable: bool = None

deprecated: bool = None

replaced_by: str | None = None

documentation_url: str | None = None

device_constraints: List[str] = None

output_shape_calc: str | None = None

complexity: str | None = None