ml_switcheroo.frameworks.paxml
==============================

.. py:module:: ml_switcheroo.frameworks.paxml

.. autoapi-nested-parse::

   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
----------

.. autoapisummary::

   ml_switcheroo.frameworks.paxml.praxis


Classes
-------

.. autoapisummary::

   ml_switcheroo.frameworks.paxml.PaxmlAdapter


Module Contents
---------------

.. py:data:: praxis
   :value: None


.. py:class:: PaxmlAdapter

   Bases: :py:obj:`ml_switcheroo.frameworks.common.jax_stack.JAXStackMixin`


   Adapter for PaxML (Praxis Layers) running on JAX.


   .. py:attribute:: display_name
      :type:  str
      :value: 'PaxML / Praxis'



   .. py:attribute:: inherits_from
      :type:  str
      :value: 'jax'



   .. py:attribute:: ui_priority
      :type:  int
      :value: 60



   .. py:method:: 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.



   .. py:property:: search_modules
      :type: List[str]



   .. py:property:: import_alias
      :type: Tuple[str, str]



   .. py:property:: import_namespaces
      :type: Dict[str, Dict[str, str]]



   .. py:property:: discovery_heuristics
      :type: Dict[str, List[str]]



   .. py:property:: test_config
      :type: Dict[str, str]



   .. py:property:: supported_tiers
      :type: List[Any]



   .. py:property:: structural_traits
      :type: ml_switcheroo.frameworks.base.StructuralTraits


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


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


      Static definitions to ensure core functionality even if discovery fails.


   .. py:property:: rng_seed_methods
      :type: List[str]



   .. py:method:: convert(data)


   .. py:method:: apply_wiring(snapshot: Dict[str, Any]) -> None

      Applies JAX Stack wiring + PaxML specific logic.



   .. py:method:: get_example_code() -> str
      :classmethod:



   .. py:method:: get_tiered_examples() -> Dict[str, str]

      Returns PaxML (Praxis) idiomatic examples.