sequince_density.flows.flowjax

sequince_density.flows.flowjax#

FlowJAX-backed flow adapter.

Classes#

FlowJax

FlowJAX-backed normalizing flow.

Module Contents#

class sequince_density.flows.flowjax.FlowJax(dims, flow_type='masked_autoregressive_flow', *, key=None, seed=0, particle_transform=None, dtype=None, device=None, **kwargs)#

Bases: sequince_density.flows.base.Flow

FlowJAX-backed normalizing flow.

The adapter always exposes native JAX arrays. Conversion to the SMC particle namespace is handled by sequince.FlowProposal.

Parameters:

  • dims (int)

  • flow_type (str | Callable)

  • key (Any | None)

  • seed (int)

  • device (Any | None)

xp#
dtype#
key#
property flow#
fit(samples, context=None, *, rng=None, **kwargs)#
Return type:

FlowJax

state_dict()#
load_state_dict(state)#
forward(samples, context=None)#
inverse(samples, context=None)#
log_prob(samples, context=None)#
sample(n_samples, context=None, *, rng=None)#
Parameters:

n_samples (int)

sample_and_log_prob(n_samples, context=None, *, rng=None)#
Parameters:

n_samples (int)