ml_switcheroo.frameworks.paxml

PaxML (Praxis) Framework Adapter (Level 2).

This adapter specializes the core JAX stack for Google’s PaxML framework. It inherits Level 0 (Core JAX) and Level 1 (Optax/Orbax) capabilities from JAXStackMixin but implements the unique structural traits of the Praxis library.

Refactor: Populates definitions for PaxML layers and namespaces.

Attributes

praxis

Classes

PaxmlAdapter

Adapter for PaxML (Praxis Layers) running on JAX.

Module Contents

ml_switcheroo.frameworks.paxml.praxis = None

class ml_switcheroo.frameworks.paxml.PaxmlAdapter

Bases: ml_switcheroo.frameworks.common.jax_stack.JAXStackMixin

Adapter for PaxML (Praxis Layers) running on JAX.

display_name: str = 'PaxML / Praxis'

inherits_from: str = 'jax'

ui_priority: int = 60

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

Scans API surface. Delegates Math/Opt to Jax Core, handles Layers (Praxis) dynamically.

property search_modules: List[str]

property import_alias: Tuple[str, str]

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

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

property test_config: Dict[str, str]

property supported_tiers: List[Any]

property structural_traits: ml_switcheroo.frameworks.base.StructuralTraits

Defines PaxML specific structural traits. Reflects that Praxis uses a setup() method for initialization rather than __init__.

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

Static definitions to ensure core functionality even if discovery fails.

property rng_seed_methods: List[str]

convert(data)

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

Applies JAX Stack wiring + PaxML specific logic.

classmethod get_example_code() → str

get_tiered_examples() → Dict[str, str]

Returns PaxML (Praxis) idiomatic examples.