sequince_density.flows.flowjax
==============================

.. py:module:: sequince_density.flows.flowjax

.. autoapi-nested-parse::

   FlowJAX-backed flow adapter.



Classes
-------

.. autoapisummary::

   sequince_density.flows.flowjax.FlowJax


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

.. py:class:: FlowJax(dims, flow_type = 'masked_autoregressive_flow', *, key = None, seed = 0, particle_transform=None, dtype=None, device = None, **kwargs)

   Bases: :py:obj:`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 :class:`sequince.FlowProposal`.


   .. py:attribute:: xp


   .. py:attribute:: dtype


   .. py:attribute:: key


   .. py:property:: flow


   .. py:method:: fit(samples, context=None, *, rng=None, **kwargs)


   .. py:method:: state_dict()


   .. py:method:: load_state_dict(state)


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


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


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


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


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