Repository: Lightning-AI/pytorch-lightning
Pull request: Make the batch_idx argument optional in all step methods
## What does this PR do?

<!--
Please include a summary of the change and which issue is fixed.
Please also include relevant motivation and context.
List any dependencies that are required for this change.

If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

The following links the related issue to the PR (https://docs.github.com/en/free-pro-team@latest/github/managing-your-work-on-github/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword)
-->

make `batch_idx` argument optional in `test_step`, `validation_step` and `predict_step`.

Fixes #18506

<!-- Does your PR introduce any breaking changes? If yes, please list them. -->

<details>
  <summary><b>Before submitting</b></summary>

- Was this **discussed/agreed** via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/Lightning-AI/lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- Did you make sure to **update the documentation** with your changes? (if necessary)
- Did you write any **new necessary tests**? (not for typos and docs)
- [x] Did you verify new and **existing tests pass** locally with your changes?
- Did you list all the **breaking changes** introduced by this pull request?
- Did you **update the CHANGELOG**? (not for typos, docs, test updates, or minor internal changes/refactors)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

</details>

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the [review guidelines](https://github.com/Lightning-AI/lightning/wiki/Review-guidelines). In short, see the following bullet-list:

<details>
  <summary>Reviewer checklist</summary>

- [ ] Is this pull request ready for review? (if not, please submit in draft mode)
- [ ] Check that all items from **Before submitting** are resolved
- [ ] Make sure the title is self-explanatory and the description concisely explains the PR
- [ ] Add labels and milestones (and optionally projects) to the PR so it can be classified

</details>

<!--

Did you have fun?

Make sure you had fun coding 🙃

-->


<!-- readthedocs-preview pytorch-lightning start -->
----
:books: Documentation preview :books:: https://pytorch-lightning--18512.org.readthedocs.build/en/18512/

<!-- readthedocs-preview pytorch-lightning end -->

Problem statement:
Make the batch_idx argument optional in all step methods
### Bug description

The training_step method does not require declaring a batch_idx argument:

```py
def training_step(self, batch, batch_idx):
    pass

# or
def training_step(self, batch):
    pass
```

But for validation/test/predict steps, the following raises an error:


```py
def validation_step(self, batch):
    pass
```

Is this intentional or a bug? 

### What version are you seeing the problem on?

v2.0, master

### How to reproduce the bug

```python
import os

import torch
from torch.utils.data import DataLoader, Dataset

from lightning.pytorch import LightningModule, Trainer


class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def training_step(self, batch, batch_idx):
        return self(batch).sum()

    def validation_step(self, batch):
        self(batch).sum()

    def configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)


def run():
    train_data = DataLoader(RandomDataset(32, 64), batch_size=2)
    val_data = DataLoader(RandomDataset(32, 64), batch_size=2)

    model = BoringModel()
    trainer = Trainer(
        default_root_dir=os.getcwd(),
        limit_train_batches=1,
        limit_val_batches=1,
        limit_test_batches=1,
        num_sanity_val_steps=0,
        max_epochs=1,
        enable_model_summary=False,
    )
    trainer.fit(model, train_dataloaders=train_data, val_dataloaders=val_data)


if __name__ == "__main__":
    run()
```


### Error messages and logs

```
  File "/Users/adrian/repositories/lightning/examples/pytorch/bug_report/bug_report_model.py", line 57, in <module>
    run()
  File "/Users/adrian/repositories/lightning/examples/pytorch/bug_report/bug_report_model.py", line 53, in run
    trainer.fit(model, train_dataloaders=train_data, val_dataloaders=val_data)
  File "/Users/adrian/repositories/lightning/src/lightning/pytorch/trainer/trainer.py", line 542, in fit
    call._call_and_handle_interrupt(
  File "/Users/adrian/repositories/lightning/src/lightning/pytorch/trainer/call.py", line 44, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/Users/adrian/repositories/lightning/src/lightning/pytorch/trainer/trainer.py", line 581, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/Users/adrian/repositories/lightning/src/lightning/pytorch/trainer/trainer.py", line 990, in _run
    results = self._run_stage()
  File "/Users/adrian/repositories/lightning/src/lightning/pytorch/trainer/trainer.py", line 1033, in _run_stage
    self.fit_loop.run()
  File "/Users/adrian/repositories/lightning/src/lightning/pytorch/loops/fit_loop.py", line 202, in run
    self.advance()
  File "/Users/adrian/repositories/lightning/src/lightning/pytorch/loops/fit_loop.py", line 363, in advance
    self.epoch_loop.run(self._data_fetcher)
  File "/Users/adrian/repositories/lightning/src/lightning/pytorch/loops/training_epoch_loop.py", line 137, in run
    self.on_advance_end(data_fetcher)
  File "/Users/adrian/repositories/lightning/src/lightning/pytorch/loops/training_epoch_loop.py", line 280, in on_advance_end
    self.val_loop.run()
  File "/Users/adrian/repositories/lightning/src/lightning/pytorch/loops/utilities.py", line 181, in _decorator
    return loop_run(self, *args, **kwargs)
  File "/Users/adrian/repositories/lightning/src/lightning/pytorch/loops/evaluation_loop.py", line 137, in run
    self._evaluation_step(batch, batch_idx, dataloader_idx, dataloader_iter)
  File "/Users/adrian/repositories/lightning/src/lightning/pytorch/loops/evaluation_loop.py", line 391, in _evaluation_step
    output = call._call_strategy_hook(trainer, hook_name, *step_args)
  File "/Users/adrian/repositories/lightning/src/lightning/pytorch/trainer/call.py", line 309, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/Users/adrian/repositories/lightning/src/lightning/pytorch/strategies/strategy.py", line 401, in validation_step
    return self.lightning_module.validation_step(*args, **kwargs)
TypeError: BoringModel.validation_step() takes 2 positional arguments but 3 were given
```

### Environment

<details>
  <summary>Current environment</summary>

```
#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0): 2.1.0dev
#- Lightning App Version (e.g., 0.5.2): -
#- PyTorch Version (e.g., 2.0): 2.0.1
#- Python version (e.g., 3.9): 3.10
#- OS (e.g., Linux): MacOs
#- CUDA/cuDNN version: -
#- GPU models and configuration: -
#- How you installed Lightning(`conda`, `pip`, source): source
#- Running environment of LightningApp (e.g. local, cloud):
```

</details>


### More info

_No response_

cc @carmocca @justusschock @awaelchli @borda

Hints from issue discussion:
Its neither intentional nor a bug. It was just never added to the other loops. You could consider it a missing feature there. The relevant piece of code is:

https://github.com/Lightning-AI/lightning/blob/a3f6e98e7c8a12e677bb9d10bd4f61ba716e0cad/src/lightning/pytorch/loops/training_epoch_loop.py#L449-L452

which is missing from https://github.com/Lightning-AI/lightning/blob/a3f6e98e7c8a12e677bb9d10bd4f61ba716e0cad/src/lightning/pytorch/loops/evaluation_loop.py#L423 and https://github.com/Lightning-AI/lightning/blob/a3f6e98e7c8a12e677bb9d10bd4f61ba716e0cad/src/lightning/pytorch/loops/prediction_loop.py#L268


hi @awaelchli , can I work on this? I would like to give it a try.

Commit message: update changelog

Patch to review:
diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md
index edc019115446b..0508f84db647e 100644
--- a/src/lightning/pytorch/CHANGELOG.md
+++ b/src/lightning/pytorch/CHANGELOG.md
@@ -185,6 +185,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Due to lack of reliability, Trainer now only runs on one GPU instead of all GPUs in a Jupyter notebook if `devices="auto"` (default) ([#18291](https://github.com/Lightning-AI/lightning/pull/18291))
 
 
+- make the `batch_idx` argument optional in `validation_step`, `test_step` and `predict_step` to maintain consistency with `training_step` ([#18512](https://github.com/Lightning-AI/lightning/pull/18512))
+
+
 ### Deprecated
 
 - Deprecated the `SingleTPUStrategy` (`strategy="single_tpu"`) in favor of `SingleDeviceXLAStrategy` (`strategy="single_xla"`) ([#17383](https://github.com/Lightning-AI/lightning/pull/17383))
diff --git a/src/lightning/pytorch/loops/evaluation_loop.py b/src/lightning/pytorch/loops/evaluation_loop.py
index fc214774559e1..cb7ad6c5fc159 100644
--- a/src/lightning/pytorch/loops/evaluation_loop.py
+++ b/src/lightning/pytorch/loops/evaluation_loop.py
@@ -42,6 +42,7 @@
 from lightning.pytorch.utilities.data import has_len_all_ranks
 from lightning.pytorch.utilities.exceptions import SIGTERMException
 from lightning.pytorch.utilities.model_helpers import is_overridden
+from lightning.pytorch.utilities.signature_utils import is_param_in_hook_signature
 
 if _RICH_AVAILABLE:
     from rich import get_console
@@ -373,7 +374,6 @@ def _evaluation_step(
         hook_kwargs = self._build_kwargs(
             batch, batch_idx, dataloader_idx if self._is_sequential and self.num_dataloaders > 1 else None
         )
-        step_args = hook_kwargs.values() if not using_dataloader_iter else (dataloader_iter,)
 
         self.batch_progress.increment_ready()
 
@@ -388,6 +388,11 @@ def _evaluation_step(
         self.batch_progress.increment_started()
 
         hook_name = "test_step" if trainer.testing else "validation_step"
+        step_args = (
+            self._build_step_args_from_hook_kwargs(hook_kwargs, hook_name)
+            if not using_dataloader_iter
+            else (dataloader_iter,)
+        )
         output = call._call_strategy_hook(trainer, hook_name, *step_args)
 
         self.batch_progress.increment_processed()
@@ -438,6 +443,14 @@ def _build_kwargs(self, batch: Any, batch_idx: int, dataloader_idx: Optional[int
             step_kwargs["dataloader_idx"] = dataloader_idx
         return step_kwargs
 
+    def _build_step_args_from_hook_kwargs(self, hook_kwargs: OrderedDict, step_hook_name: str) -> tuple:
+        """Helper method to build args for `test_step` or `validation_step`."""
+        kwargs = hook_kwargs.copy()
+        step_hook_fx = getattr(self.trainer.lightning_module, step_hook_name)
+        if not is_param_in_hook_signature(step_hook_fx, "batch_idx", min_args=2):
+            kwargs.pop("batch_idx", None)
+        return tuple(kwargs.values())
+
     def _verify_dataloader_idx_requirement(self) -> None:
         trainer = self.trainer
         step_hook = "test_step" if trainer.testing else "validation_step"
diff --git a/src/lightning/pytorch/loops/prediction_loop.py b/src/lightning/pytorch/loops/prediction_loop.py
index b6e9fcff1045a..6f649b0fe2139 100644
--- a/src/lightning/pytorch/loops/prediction_loop.py
+++ b/src/lightning/pytorch/loops/prediction_loop.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from collections import OrderedDict
-from typing import Any, Dict, Iterator, List, Optional, Union
+from typing import Any, Iterator, List, Optional, Union
 
 import torch
 from lightning_utilities import WarningCache
@@ -38,6 +38,7 @@
 from lightning.pytorch.utilities.combined_loader import CombinedLoader
 from lightning.pytorch.utilities.data import has_len_all_ranks
 from lightning.pytorch.utilities.exceptions import MisconfigurationException
+from lightning.pytorch.utilities.signature_utils import is_param_in_hook_signature
 from lightning.pytorch.utilities.types import _PREDICT_OUTPUT
 
 
@@ -236,7 +237,6 @@ def _predict_step(
         # the `_step` methods don't take a batch_idx when `dataloader_iter` is used, but all other hooks still do,
         # so we need different kwargs
         hook_kwargs = self._build_kwargs(batch, batch_idx, dataloader_idx if self.num_dataloaders > 1 else None)
-        step_args = hook_kwargs.values() if not using_dataloader_iter else (dataloader_iter,)
 
         call._call_callback_hooks(trainer, "on_predict_batch_start", *hook_kwargs.values())
         call._call_lightning_module_hook(trainer, "on_predict_batch_start", *hook_kwargs.values())
@@ -244,6 +244,11 @@ def _predict_step(
         self.batch_progress.increment_started()
 
         # configure step_kwargs
+        step_args = (
+            self._build_step_args_from_hook_kwargs(hook_kwargs, "predict_step")
+            if not using_dataloader_iter
+            else (dataloader_iter,)
+        )
         predictions = call._call_strategy_hook(trainer, "predict_step", *step_args)
         if predictions is None:
             self._warning_cache.warn("predict returned None if it was on purpose, ignore this warning...")
@@ -265,7 +270,7 @@ def _predict_step(
         if self._return_predictions or any_on_epoch:
             self._predictions[dataloader_idx].append(move_data_to_device(predictions, torch.device("cpu")))
 
-    def _build_kwargs(self, batch: Any, batch_idx: int, dataloader_idx: Optional[int]) -> Dict[str, Any]:
+    def _build_kwargs(self, batch: Any, batch_idx: int, dataloader_idx: Optional[int]) -> OrderedDict:
         """Assembles the keyword arguments for the ``predict_step``
 
         Args:
@@ -282,6 +287,14 @@ def _build_kwargs(self, batch: Any, batch_idx: int, dataloader_idx: Optional[int
             step_kwargs["dataloader_idx"] = dataloader_idx
         return step_kwargs
 
+    def _build_step_args_from_hook_kwargs(self, hook_kwargs: OrderedDict, step_hook_name: str) -> tuple:
+        """Helper method to build args for `predict_step`."""
+        kwargs = hook_kwargs.copy()
+        step_hook_fx = getattr(self.trainer.lightning_module, step_hook_name)
+        if not is_param_in_hook_signature(step_hook_fx, "batch_idx", min_args=2):
+            kwargs.pop("batch_idx", None)
+        return tuple(kwargs.values())
+
     def _get_batch_indices(self, dataloader: object) -> List[List[int]]:  # batches x samples
         """Returns a reference to the seen batch indices if the dataloader has a batch sampler wrapped by our
         :class:`~lightning.pytorch.overrides.distributed._IndexBatchSamplerWrapper`."""
diff --git a/tests/tests_pytorch/loops/test_evaluation_loop.py b/tests/tests_pytorch/loops/test_evaluation_loop.py
index b4c6ca4e462fd..d1235e2c85a60 100644
--- a/tests/tests_pytorch/loops/test_evaluation_loop.py
+++ b/tests/tests_pytorch/loops/test_evaluation_loop.py
@@ -496,3 +496,28 @@ def test_step(self, batch, batch_idx):
     actual = trainer.num_val_batches if fn == "validate" else trainer.num_test_batches
     assert actual == [3, 2]
     assert seen == expected
+
+
+def test_evaluation_loop_when_batch_idx_argument_is_not_given(tmpdir):
+    class TestModel(BoringModel):
+        def __init__(self) -> None:
+            super().__init__()
+            self.validation_step_called = False
+            self.test_step_called = False
+
+        def validation_step(self, batch):
+            self.validation_step_called = True
+            return {"x": self.step(batch)}
+
+        def test_step(self, batch):
+            self.test_step_called = True
+            return {"y": self.step(batch)}
+
+    trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=1)
+    model = TestModel()
+
+    trainer.fit(model)
+    assert model.validation_step_called
+
+    trainer.test(model)
+    assert model.test_step_called
diff --git a/tests/tests_pytorch/loops/test_prediction_loop.py b/tests/tests_pytorch/loops/test_prediction_loop.py
index 412154884a4d4..de98d3853b3a9 100644
--- a/tests/tests_pytorch/loops/test_prediction_loop.py
+++ b/tests/tests_pytorch/loops/test_prediction_loop.py
@@ -300,3 +300,20 @@ def on_predict_batch_end(self, outputs, batch, batch_idx, **_):
     model = IgnoringModel2()
     with pytest.raises(RuntimeError, match="no `dataloader_idx` argument in `IgnoringModel2.on_predict_batch_end"):
         trainer.predict(model)
+
+
+def test_prediction_loop_when_batch_idx_argument_is_not_given(tmpdir):
+    class TestModel(BoringModel):
+        def __init__(self) -> None:
+            super().__init__()
+            self.predict_step_called = False
+
+        def predict_step(self, batch):
+            self.predict_step_called = True
+            return self.step(batch)
+
+    trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=1)
+    model = TestModel()
+
+    trainer.predict(model)
+    assert model.predict_step_called


Code review comments:
diff_hunk: @@ -496,3 +496,28 @@ def test_step(self, batch, batch_idx):
     actual = trainer.num_val_batches if fn == "validate" else trainer.num_test_batches
     assert actual == [3, 2]
     assert seen == expected
+
+
+def test_evaluation_loop_when_batch_idx_argument_is_not_given(tmpdir):
+    class TestModel(BoringModel):
+        def __init__(self) -> None:
+            super().__init__()
+            self.validation_step_called = False
+            self.test_step_called = False
+
+        def validation_step(self, batch):
+            self.validation_step_called = True
+            return {"x": self.step(batch)}
+
+        def test_step(self, batch):
+            self.test_step_called = True
+            return {"y": self.step(batch)}
+
+    trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=1)
path: tests/tests_pytorch/loops/test_evaluation_loop.py
text: @user1:
We should probably set `logger=False`, `enable_checkpointing=False`, `enable_progress_bar=False` in these tests

@author:
> Revised, I feel the reason for doing this is to reduce the execution of unrelated code during testing. Are there any other reasons why we need to do this?

@user2:
That's right. It'll make the test run faster too

diff_hunk: @@ -496,3 +496,28 @@ def test_step(self, batch, batch_idx):
     actual = trainer.num_val_batches if fn == "validate" else trainer.num_test_batches
     assert actual == [3, 2]
     assert seen == expected
+
+
+def test_evaluation_loop_when_batch_idx_argument_is_not_given(tmpdir):
+    class TestModel(BoringModel):
+        def __init__(self) -> None:
+            super().__init__()
+            self.validation_step_called = False
+            self.test_step_called = False
+
+        def validation_step(self, batch):
+            self.validation_step_called = True
+            return {"x": self.step(batch)}
+
+        def test_step(self, batch):
+            self.test_step_called = True
+            return {"y": self.step(batch)}
+
+    trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=1)
+    model = TestModel()
+
+    trainer.fit(model)
path: tests/tests_pytorch/loops/test_evaluation_loop.py
text: @user1:
```suggestion
    trainer.validate(model)
```

diff_hunk: @@ -185,6 +185,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Due to lack of reliability, Trainer now only runs on one GPU instead of all GPUs in a Jupyter notebook if `devices="auto"` (default) ([#18291](https://github.com/Lightning-AI/lightning/pull/18291))
 
 
+- make the `batch_idx` argument optional in `validation_step`, `test_step` and `predict_step` to maintain consistency with `training_step` ([#18512](https://github.com/Lightning-AI/lightning/pull/18512))
path: src/lightning/pytorch/CHANGELOG.md
text: @user1:
```suggestion
- Made the `batch_idx` argument optional in `validation_step`, `test_step` and `predict_step` to maintain consistency with `training_step` ([#18512](https://github.com/Lightning-AI/lightning/pull/18512))
```