aspire.flows

Submodules

• aspire.flows.base
• aspire.flows.jax
• aspire.flows.torch

Attributes

logger

Functions

get_flow_wrapper([backend, flow_matching])

Get the wrapper for the flow implementation.

Package Contents

aspire.flows.logger

aspire.flows.get_flow_wrapper(backend='zuko', flow_matching=False)

Get the wrapper for the flow implementation.

Parameters:

• backend (str) – The backend to use. Options are “zuko” (PyTorch), “flowjax” (JAX), or any other registered flow class via entry points. Default is “zuko”.

• flow_matching (bool, optional) – Whether to use flow matching variant of the flow. Default is False.

Returns:

• FlowClass (type) – The flow class corresponding to the specified backend.

• xp (module) – The array API module corresponding to the specified backend.

Return type:

tuple[type, Any]