sequince_density.flows.zuko

sequince_density.flows.zuko#

Zuko-backed proposal and flow adapter.

Attributes#

logger

Classes#

ZukoFlow

Zuko-backed flow adapter for Sequince.

Module Contents#

sequince_density.flows.zuko.logger#
class sequince_density.flows.zuko.ZukoFlow(dims, flow_class='MAF', particle_transform=None, seed=1234, device='cpu', dtype=None, **kwargs)#

Bases: sequince_density.flows.base.Flow

Zuko-backed flow adapter for Sequince.

Parameters:

  • dims (int)

  • flow_class (str | Callable)

  • seed (int)

  • device (str)

  • dtype (torch.dtype | None)

xp#
dtype#
seed = 1234#
property flow#
loss_fn(x)#
fit(samples, context=None, n_epochs=100, lr=0.001, batch_size=500, validation_fraction=0.2, *, rng=None)#
Parameters:
Return type:

ZukoFlow

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

n_samples (int)

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

n_samples (int)

log_prob(x, context=None)#
forward(x, context=None)#
inverse(z, context=None)#