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, *, rng=None)


   .. 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)