scvi.train.JaxTrainingPlan
==========================

.. currentmodule:: scvi.train

.. add toctree option to make autodoc generate the pages

.. autoclass:: JaxTrainingPlan
   :show-inheritance:



Attributes table
~~~~~~~~~~~~~~~~

.. autosummary::

        ~scvi.train.JaxTrainingPlan.training






Methods table
~~~~~~~~~~~~~~

.. autosummary::

    ~scvi.train.JaxTrainingPlan.backward
    ~scvi.train.JaxTrainingPlan.configure_optimizers
    ~scvi.train.JaxTrainingPlan.forward
    ~scvi.train.JaxTrainingPlan.get_optimizer_creator
    ~scvi.train.JaxTrainingPlan.jit_training_step
    ~scvi.train.JaxTrainingPlan.jit_validation_step
    ~scvi.train.JaxTrainingPlan.optimizer_step
    ~scvi.train.JaxTrainingPlan.set_train_state
    ~scvi.train.JaxTrainingPlan.training_step
    ~scvi.train.JaxTrainingPlan.transfer_batch_to_device
    ~scvi.train.JaxTrainingPlan.validation_step





Attributes
~~~~~~~~~~



.. autoattribute:: JaxTrainingPlan.training






Methods
~~~~~~~



.. automethod:: JaxTrainingPlan.backward

.. automethod:: JaxTrainingPlan.configure_optimizers

.. automethod:: JaxTrainingPlan.forward

.. automethod:: JaxTrainingPlan.get_optimizer_creator

.. automethod:: JaxTrainingPlan.jit_training_step

.. automethod:: JaxTrainingPlan.jit_validation_step

.. automethod:: JaxTrainingPlan.optimizer_step

.. automethod:: JaxTrainingPlan.set_train_state

.. automethod:: JaxTrainingPlan.training_step

.. automethod:: JaxTrainingPlan.transfer_batch_to_device

.. automethod:: JaxTrainingPlan.validation_step