aspire.flows.jax

Submodules

• aspire.flows.jax.flows
• aspire.flows.jax.utils

Functions

get_flow(*, key, dims[, flow_type, bijection_type, ...])

Package Contents

aspire.flows.jax.get_flow(*, key, dims, flow_type='masked_autoregressive_flow', bijection_type=None, bijection_kwargs=None, dtype=None, **kwargs)

Parameters:

• key (jax.Array)

• dims (int)

• flow_type (str | Callable)

• bijection_type (str | flowjax.bijections.AbstractBijection | None)

• bijection_kwargs (dict | None)

Return type:

flowjax.distributions.Transformed