aspire.flows.jax.utils

aspire.flows.jax.utils#

Functions#

get_flow_function_class(name)

get_bijection_class(name)

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

Module Contents#

aspire.flows.jax.utils.get_flow_function_class(name)[source]#
Parameters:

name (str)

Return type:

Callable

aspire.flows.jax.utils.get_bijection_class(name)[source]#
Parameters:

name (str)

Return type:

Callable

aspire.flows.jax.utils.get_flow(*, key, dims, flow_type='masked_autoregressive_flow', bijection_type=None, bijection_kwargs=None, dtype=None, **kwargs)[source]#
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