Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/core/lightning.py (Python)

Code before:
                def test_epoch_end(self, outputs):
                    final_value = 0
                    for dataloader_outputs in outputs:
                        for test_step_out in dataloader_outputs:
                            # do something
                            final_value += test_step_out

                    self.log('final_metric', final_value)
        """

    def configure_optimizers(self):
        r"""
        Choose what optimizers and learning-rate schedulers to use in your optimization.
        Normally you'd need one. But in the case of GANs or similar you might have multiple.

        Return:
            Any of these 6 options.

            - Single optimizer.
            - List or Tuple - List of optimizers.
            - Two lists - The first list has multiple optimizers, the second a list of LR schedulers (or lr_dict).
            - Dictionary, with an 'optimizer' key, and (optionally) a 'lr_scheduler'
              key which value is a single LR scheduler or lr_dict.
            - Tuple of dictionaries as described, with an optional 'frequency' key.
            - None - Fit will run without any optimizer.

        Note:
            The ``frequency`` value specified in a dict along with the ``optimizer`` key is an int corresponding
            to the number of sequential batches optimized with the specific optimizer.
            It should be given to none or to all of the optimizers.
            There is a difference between passing multiple optimizers in a list,
            and passing multiple optimizers in dictionaries with a frequency of 1:
            In the former case, all optimizers will operate on the given batch in each optimization step.
            In the latter, only one optimizer will operate on the given batch at every step.
            This is different from the ``frequency`` value specified in the lr_dict mentioned below.

            .. code-block:: python

                def configure_optimizers(self):
                    optimizer_one = torch.optim.SGD(self.model.parameters(), lr=0.01)
                    optimizer_two = torch.optim.SGD(self.model.parameters(), lr=0.01)
                    return [
                        {
                            'optimizer': optimizer_one,
                            'frequency': 5
                        },
                        {
                            'optimizer': optimizer_two,
                            'frequency': 10
                    }]


Review comment (performance):
mind update optimizer, lr_scheduler, and frequency with
```
``param``
```
above

Code after:
            so that you don't have to change your code.

        .. code-block:: python

            # pseudocode
            sub_batches = split_batches_for_dp(batch)
            batch_parts_outputs = [test_step(sub_batch) for sub_batch in sub_batches]
            test_step_end(batch_parts_outputs)

        Args:
            batch_parts_outputs: What you return in :meth:`test_step` for each batch part.

        Return:
            None or anything

        .. code-block:: python

            # WITHOUT test_step_end
            # if used in DP or DDP2, this batch is 1/num_gpus large
            def test_step(self, batch, batch_idx):
                # batch is 1/num_gpus big
                x, y = batch

                out = self(x)
                loss = self.softmax(out)
                self.log('test_loss', loss)

            # --------------
            # with test_step_end to do softmax over the full batch
            def test_step(self, batch, batch_idx):
                # batch is 1/num_gpus big
                x, y = batch

                out = self.encoder(x)
                return out

            def test_step_end(self, output_results):
                # this out is now the full size of the batch
                all_test_step_outs = output_results.out
                loss = nce_loss(all_test_step_outs)
                self.log('test_loss', loss)

        See Also:
            See the :ref:`advanced/multi_gpu:Multi-GPU training` guide for more details.
        """

    def test_epoch_end(self, outputs: EPOCH_OUTPUT) -> None:
        """
        Called at the end of a test epoch with the output of all test steps.

        .. code-block:: python