xlstm_jax.distributed.single_gpu
================================

.. py:module:: xlstm_jax.distributed.single_gpu


Functions
---------

.. autoapisummary::

   xlstm_jax.distributed.single_gpu.accumulate_gradients_loop
   xlstm_jax.distributed.single_gpu.accumulate_gradients_scan
   xlstm_jax.distributed.single_gpu.accumulate_gradients


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

.. py:function:: accumulate_gradients_loop(state, batch, rng, num_minibatches, loss_fn)

   Calculate gradients and metrics for a batch using gradient accumulation.

   :param state: Current training state.
   :param batch: Full training batch.
   :param rng: Random number generator to use.
   :param num_minibatches: Number of mini-batches to split the batch into. Equal to the number of gradient accumulation
                           steps.
   :param loss_fn: Loss function to calculate gradients and metrics.

   :returns: Tuple with accumulated gradients, metrics, and collected mutable variables over the mini-batches.


.. py:function:: accumulate_gradients_scan(state, batch, rng, num_minibatches, loss_fn)

   Calculate gradients and metrics for a batch using gradient accumulation.

   In this version, we use `jax.lax.scan` to loop over the mini-batches. This is more efficient in terms of compilation
   time.

   :param state: Current training state.
   :param batch: Full training batch.
   :param rng: Random number generator to use.
   :param num_minibatches: Number of mini-batches to split the batch into. Equal to the number of gradient accumulation
                           steps.
   :param loss_fn: Loss function to calculate gradients and metrics.

   :returns: Tuple with accumulated gradients, metrics, and collected mutable variables over the mini-batches.


.. py:function:: accumulate_gradients(state, batch, rng, num_minibatches, loss_fn, use_scan = False)

   Calculate gradients and metrics for a batch using gradient accumulation.

   This function supports scanning over the mini-batches using `jax.lax.scan` or using a for loop.

   :param state: Current training state.
   :param batch: Full training batch.
   :param rng: Random number generator to use.
   :param num_minibatches: Number of mini-batches to split the batch into. Equal to the number of gradient accumulation
                           steps.
   :param loss_fn: Loss function to calculate gradients and metrics.
   :param use_scan: Whether to use `jax.lax.scan` for looping over the mini-batches.

   :returns: Tuple with accumulated gradients, metrics, and collected mutable variables over the mini-batches.