sequince_density.flows.zuko
===========================

.. py:module:: sequince_density.flows.zuko

.. autoapi-nested-parse::

   Zuko-backed proposal and flow adapter.



Attributes
----------

.. autoapisummary::

   sequince_density.flows.zuko.logger


Classes
-------

.. autoapisummary::

   sequince_density.flows.zuko.ZukoFlow


Module Contents
---------------

.. py:data:: logger

.. py:class:: ZukoFlow(dims, flow_class = 'MAF', particle_transform=None, seed = 1234, device = 'cpu', dtype = None, **kwargs)

   Bases: :py:obj:`sequince_density.flows.base.Flow`


   Zuko-backed flow adapter for Sequince.


   .. py:attribute:: xp


   .. py:attribute:: dtype


   .. py:attribute:: seed
      :value: 1234



   .. py:property:: flow


   .. py:method:: loss_fn(x)


   .. py:method:: 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.

      :param samples: Nonempty, finite training particles. Transformed particles must
                      also be finite and have the same shape.
      :type samples: array_like, shape (n_samples, dims)
      :param context: Context passed to the particle transform.
      :type context: object, optional
      :param n_epochs: Maximum number of training epochs.
      :type n_epochs: int, default=100
      :param lr: Initial Adam learning rate.
      :type lr: float, default=1e-3
      :param batch_size: Maximum number of particles per training or validation batch.
      :type batch_size: int, default=500
      :param validation_fraction: 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.
      :type validation_fraction: float, default=0.2
      :param clip_grad: Positive maximum total gradient norm, applied before each optimizer
                        step. None disables clipping.
      :type clip_grad: float or None, default=None
      :param lr_annealing: Apply cosine learning-rate annealing from lr to zero over n_epochs,
                           stepping the scheduler after each completed epoch.
      :type lr_annealing: bool, default=False
      :param patience: 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.
      :type patience: int or None, default=None
      :param show_progress: Show epoch progress and training and validation losses.
      :type show_progress: bool, default=False
      :param rng: RNG for fitting, isolated from global Torch state when supplied.
      :type rng: RandomGenerator or int or None, optional

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



   .. py:method:: state_dict()


   .. py:method:: load_state_dict(state)


   .. py:method:: sample_and_log_prob(n_samples, context=None, *, rng=None)


   .. py:method:: sample(n_samples, context=None, *, rng=None)


   .. py:method:: log_prob(x, context=None)


   .. py:method:: forward(x, context=None)


   .. py:method:: inverse(z, context=None)