ml_switcheroo.frameworks.mlx

Apple MLX Framework Adapter.

This module provides the adapter for Apple’s MLX array framework. It supports: 1. Unified Memory math: Mapping mlx.core operations. 2. Neural Networks: Mapping mlx.nn layers and containers. 3. Discovery: Runtime introspection of the MLX API surface. 4. Test Config: Provides generation templates for gen-tests.

Refactor: Distributed definitions populated for MLX specific Layers, Ops, and Optimizers.

Classes

MLXAdapter

Adapter for Apple MLX (Silicon-optimized tensor framework).

Module Contents

class ml_switcheroo.frameworks.mlx.MLXAdapter

Adapter for Apple MLX (Silicon-optimized tensor framework).

display_name: str = 'Apple MLX'

inherits_from: str | None = None

ui_priority: int = 50

property search_modules: List[str]

Returns list of MLX submodules to scan during discovery.

property import_alias: Tuple[str, str]

import mlx.core as mx.

Type:

Default alias for core array operations

property import_namespaces: Dict[str, Dict[str, str]]

Defines namespace mapping for source-to-target imports. Maps torch.nn -> mlx.nn (aliased as nn) and mlx.core mappings.

property discovery_heuristics: Dict[str, List[str]]

Regex patterns for categorizing discovered APIs into Tiers.

property test_config: Dict[str, str]

Templates for generating physical verification tests.

property supported_tiers: List[ml_switcheroo.enums.SemanticTier]

Returns supported semantic tiers (Array, Neural, Extras).

property structural_traits: ml_switcheroo.frameworks.base.StructuralTraits

Defines structural rewriting rules (Classes, Methods, Init).

property definitions: Dict[str, ml_switcheroo.frameworks.base.StandardMap]

Static definitions for MLX mappings. Covers Optimization, Math, Layers, and Compilation.

property rng_seed_methods: List[str]

Returns list of global seed setters.

collect_api(category: ml_switcheroo.frameworks.base.StandardCategory) → List[ml_switcheroo.core.ghost.GhostRef]

Performs runtime introspection to discover available MLX APIs.

convert(data: Any) → Any

Converts input data (NumPy/List) to MLX Tensor for verification.

classmethod get_example_code() → str

Returns default example.

get_tiered_examples() → Dict[str, str]

Returns MLX idiomatic examples used for validity testing.

get_device_syntax(device_type: str, device_index: str | None = None) → str

Returns device constructor syntax.

get_serialization_imports() → List[str]

Returns imports for serialization.

get_serialization_syntax(op: str, file_arg: str, object_arg: str | None = None) → str

Returns save/load syntax.

apply_wiring(snapshot: Dict[str, Any]) → None

Applies manual wiring for MLX. Overrides/Patches snapshot items that cannot be statically defined.