Pyro Primitives
===============

.. automodule:: numpyro.primitives

param
-----
.. autofunction:: numpyro.primitives.param

sample
------
.. autofunction:: numpyro.primitives.sample

plate
-----
.. autoclass:: numpyro.primitives.plate

plate_stack
-----------
.. autofunction:: numpyro.primitives.plate_stack

subsample
---------
.. autofunction:: numpyro.primitives.subsample

deterministic
-------------
.. autofunction:: numpyro.primitives.deterministic

prng_key
--------
.. autofunction:: numpyro.primitives.prng_key

factor
------
.. autofunction:: numpyro.primitives.factor

get_mask
--------
.. autofunction:: numpyro.primitives.get_mask

module
------
.. autofunction:: numpyro.primitives.module

mutable
-------
.. autofunction:: numpyro.primitives.mutable

flax_module
-----------
.. autofunction:: numpyro.contrib.module.flax_module

haiku_module
------------
.. autofunction:: numpyro.contrib.module.haiku_module

nnx_module
----------
.. autofunction:: numpyro.contrib.module.nnx_module

eqx_module
----------
.. autofunction:: numpyro.contrib.module.eqx_module

random_flax_module
------------------
.. autofunction:: numpyro.contrib.module.random_flax_module

random_haiku_module
-------------------
.. autofunction:: numpyro.contrib.module.random_haiku_module

random_nnx_module
-----------------
.. autofunction:: numpyro.contrib.module.random_nnx_module

random_eqx_module
-----------------
.. autofunction:: numpyro.contrib.module.random_eqx_module

scan
----
.. autofunction:: numpyro.contrib.control_flow.scan

cond
----
.. autofunction:: numpyro.contrib.control_flow.cond