Repository: Lightning-AI/pytorch-lightning
Pull request: Fix `TensorBoardLogger.log_graph` not recording the graph
## What does this PR do?


Fixes #17844 


<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>

- [x] Is this pull request ready for review? (if not, please submit in draft mode)
- [x] Check that all items from **Before submitting** are resolved
- [x] 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>



Problem statement:
`fabric.logger.log_graph()` seems to do nothing
### Bug description
fabric.logger.log_graph() seems to do nothing.
My code looks like below, but no graph record in tensorboard log.
```python
fabric = Fabric(loggers=TensorBoardLogger(root_dir="./logs", name=None))
model, optimizer = fabric.setup(model, optimizer)
fabric.logger.log_graph(model=model, input_array=next(iter(val_loader))[0])
```

I explored the source code of fabric.logger.log_graph()

```python
@rank_zero_only
    def log_graph(self, model: Module, input_array: Optional[Tensor] = None) -> None:
        model_example_input = getattr(model, "example_input_array", None)
        input_array = model_example_input if input_array is None else input_array

        if input_array is None:
            rank_zero_warn(
                "Could not log computational graph to TensorBoard: The `model.example_input_array` attribute"
                " is not set or `input_array` was not given."
            )
        elif not isinstance(input_array, (Tensor, tuple)):
            rank_zero_warn(
                "Could not log computational graph to TensorBoard: The `input_array` or `model.example_input_array`"
                f" has type {type(input_array)} which can't be traced by TensorBoard. Make the input array a tuple"
                f" representing the positional arguments to the model's `forward()` implementation."
            )
        elif callable(getattr(model, "_on_before_batch_transfer", None)) and callable(
            getattr(model, "_apply_batch_transfer_handler", None)
        ):
            # this is probably is a LightningModule
            input_array = model._on_before_batch_transfer(input_array)  # type: ignore[operator]
            input_array = model._apply_batch_transfer_handler(input_array)  # type: ignore[operator]
            self.experiment.add_graph(model, input_array)
```
It seems self.experiment.add_graph(model, input_array) should be reduced one indent, because that line no executed.


### Environment

<details>
  <summary>Current environment</summary>

```
#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version: 2.0.3
#- PyTorch Version: 2.0.1
#- Python version: 3.10
```

</details>


cc @awaelchli @borda @Blaizzy @carmocca @justusschock

Hints from issue discussion:
Yes @shihaoyin 
For Fabric it would make sense to move this outside the else block IMO. Contributions for this change are welcome :)

Commit message: fixes no graph recorded bug

Patch to review:
diff --git a/src/lightning/fabric/loggers/tensorboard.py b/src/lightning/fabric/loggers/tensorboard.py
index 3a5b468f3a6b7..07724dfa17a74 100644
--- a/src/lightning/fabric/loggers/tensorboard.py
+++ b/src/lightning/fabric/loggers/tensorboard.py
@@ -265,6 +265,11 @@ def log_graph(self, model: Module, input_array: Optional[Tensor] = None) -> None
             input_array = model._on_before_batch_transfer(input_array)  # type: ignore[operator]
             input_array = model._apply_batch_transfer_handler(input_array)  # type: ignore[operator]
             self.experiment.add_graph(model, input_array)
+        elif hasattr(model, "module") and callable(getattr(model, "_redirection_through_forward", None)):
+            # this is probably is a _FabricModule
+            self.experiment.add_graph(model.module, input_array)
+        else:
+            self.experiment.add_graph(model, input_array)
 
     @rank_zero_only
     def save(self) -> None:


Code review comments:
diff_hunk: @@ -265,6 +265,11 @@ def log_graph(self, model: Module, input_array: Optional[Tensor] = None) -> None
             input_array = model._on_before_batch_transfer(input_array)  # type: ignore[operator]
             input_array = model._apply_batch_transfer_handler(input_array)  # type: ignore[operator]
             self.experiment.add_graph(model, input_array)
+        elif hasattr(model, "module") and callable(getattr(model, "_redirection_through_forward", None)):
path: src/lightning/fabric/loggers/tensorboard.py
text: @user1:
I suggest to do this instead, which should be more robust:

```py
from lightning.fabric.wrappers import _unwrap_objects
model = _unwrap_objects(model)
```

We should probably extend this test here to address the TODO:
https://github.com/Lightning-AI/lightning/blob/66293758c0192955148e8b4bd9dd156927c23503/tests/tests_fabric/loggers/test_tensorboard.py#L149-L151

I can help with that if you wish.