aspire.flows
============

.. py:module:: aspire.flows


Submodules
----------

.. toctree::
   :maxdepth: 1

   /autoapi/aspire/flows/base/index
   /autoapi/aspire/flows/jax/index
   /autoapi/aspire/flows/torch/index


Attributes
----------

.. autoapisummary::

   aspire.flows.logger


Functions
---------

.. autoapisummary::

   aspire.flows.get_flow_wrapper


Package Contents
----------------

.. py:data:: logger

.. py:function:: get_flow_wrapper(backend = 'zuko', flow_matching = False)

   Get the wrapper for the flow implementation.

   :param backend: The backend to use. Options are "zuko" (PyTorch), "flowjax" (JAX), or
                   any other registered flow class via entry points. Default is "zuko".
   :type backend: :py:class:`str`
   :param flow_matching: Whether to use flow matching variant of the flow. Default is False.
   :type flow_matching: :py:class:`bool`, *optional*

   :returns: * **FlowClass** (:py:class:`type`) -- The flow class corresponding to the specified backend.
             * **xp** (:py:class:`module`) -- The array API module corresponding to the specified backend.