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, clip_grad=None, lr_annealing=False, patience=None, show_progress=False, *, rng=None)

Fit by maximum likelihood, restoring the best epoch when validating.

Parameters:

• samples (array_like, shape (n_samples, dims)) – Nonempty, finite training particles. Transformed particles must also be finite and have the same shape.

• context (object, optional) – Context passed to the particle transform.

• n_epochs (int, default=100) – Maximum number of training epochs.

• lr (float, default=1e-3) – Initial Adam learning rate.

• batch_size (int, default=500) – Maximum number of particles per training or validation batch.

• validation_fraction (float, default=0.2) – Fraction held out for validation. Zero skips validation and retains the final parameters. For nonzero fractions, an empty split after rounding uses all particles for both training and validation.

• clip_grad (float or None, default=None) – Positive maximum total gradient norm, applied before each optimizer step. None disables clipping.

• lr_annealing (bool, default=False) – Apply cosine learning-rate annealing from lr to zero over n_epochs, stepping the scheduler after each completed epoch.

• patience (int or None, default=None) – Positive number of consecutive epochs without a strictly lower validation loss before stopping. Improvement resets the counter; ties count as unsuccessful epochs. Requires validation; None disables early stopping.

• show_progress (bool, default=False) – Show epoch progress and training and validation losses.

• rng (RandomGenerator or int or None, optional) – RNG for fitting, isolated from global Torch state when supplied.

Returns:

This fitted adapter. fit_history contains losses for every completed epoch; validation_loss is empty when validation is off.

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)