Consistent imports in docs for core APIs (#18869)

Co-authored-by: Sebastian Raschka <[email protected]>

(cherry picked from commit f6a36cf)

Author: awaelchli

_notebooks

@@ -88,10 +88,10 @@ How to setup the debug mode for Training on TPUs?
 
 .. code-block:: python
 
-    import lightning.pytorch as pl
+    import lightning as L
 
     my_model = MyLightningModule()
-    trainer = pl.Trainer(accelerator="tpu", devices=8, strategy="xla_debug")
+    trainer = L.Trainer(accelerator="tpu", devices=8, strategy="xla_debug")
     trainer.fit(my_model)
 
 Example Metrics report:

docs/source-pytorch/accelerators/tpu_faq.rst

@@ -44,10 +44,10 @@ To use a full TPU pod skip to the TPU pod section.
 
 .. code-block:: python
 
-    import lightning.pytorch as pl
+    import lightning as L
 
     my_model = MyLightningModule()
-    trainer = pl.Trainer(accelerator="tpu", devices=8)
+    trainer = L.Trainer(accelerator="tpu", devices=8)
     trainer.fit(my_model)
 
 That's it! Your model will train on all 8 TPU cores.
@@ -113,10 +113,10 @@ By default, TPU training will use 32-bit precision. To enable it, do
 
 .. code-block:: python
 
-    import lightning.pytorch as pl
+    import lightning as L
 
     my_model = MyLightningModule()
-    trainer = pl.Trainer(accelerator="tpu", precision="16-true")
+    trainer = L.Trainer(accelerator="tpu", precision="16-true")
     trainer.fit(my_model)
 
 Under the hood the xla library will use the `bfloat16 type <https://en.wikipedia.org/wiki/Bfloat16_floating-point_format>`_.

docs/source-pytorch/accelerators/tpu_intermediate.rst

@@ -132,12 +132,11 @@ For even more speed benefit, DeepSpeed offers an optimized CPU version of ADAM c
 
 .. code-block:: python
 
-    import lightning.pytorch
-    from lightning.pytorch import Trainer
+    from lightning.pytorch import LightningModule, Trainer
     from deepspeed.ops.adam import DeepSpeedCPUAdam
 
 
-    class MyModel(pl.LightningModule):
+    class MyModel(LightningModule):
         ...
 
         def configure_optimizers(self):
@@ -180,7 +179,7 @@ Also please have a look at our :ref:`deepspeed-zero-stage-3-tips` which contains
     from deepspeed.ops.adam import FusedAdam
 
 
-    class MyModel(pl.LightningModule):
+    class MyModel(LightningModule):
         ...
 
         def configure_optimizers(self):
@@ -202,7 +201,7 @@ You can also use the Lightning Trainer to run predict or evaluate with DeepSpeed
     from lightning.pytorch import Trainer
 
 
-    class MyModel(pl.LightningModule):
+    class MyModel(LightningModule):
         ...
 
 
@@ -228,7 +227,7 @@ This reduces the time taken to initialize very large models, as well as ensure w
     from deepspeed.ops.adam import FusedAdam
 
 
-    class MyModel(pl.LightningModule):
+    class MyModel(LightningModule):
         ...
 
         def configure_model(self):
@@ -367,7 +366,7 @@ This saves memory when training larger models, however requires using a checkpoi
     import deepspeed
 
 
-    class MyModel(pl.LightningModule):
+    class MyModel(LightningModule):
         ...
 
         def configure_model(self):

docs/source-pytorch/advanced/model_parallel/deepspeed.rst

@@ -398,7 +398,7 @@ The :class:`~lightning.pytorch.core.datamodule.LightningDataModule` class provid
 
 .. code-block:: python
 
-    class MNISTDataModule(pl.LightningDataModule):
+    class MNISTDataModule(L.LightningDataModule):
         def prepare_data(self):
             MNIST(self.data_dir, download=True)
 
@@ -421,7 +421,7 @@ For this, all data pre-loading should be done on the main process inside :meth:`
 
 .. code-block:: python
 
-    class MNISTDataModule(pl.LightningDataModule):
+    class MNISTDataModule(L.LightningDataModule):
         def __init__(self, data_dir: str):
             self.mnist = MNIST(data_dir, download=True, transform=T.ToTensor())
 

docs/source-pytorch/advanced/training_tricks.rst

@@ -164,7 +164,7 @@ to the class constructor. For example, your model is defined as:
 .. code:: python
 
     # model.py
-    class MyModel(pl.LightningModule):
+    class MyModel(L.LightningModule):
         def __init__(self, criterion: torch.nn.Module):
             self.criterion = criterion
 

docs/source-pytorch/cli/lightning_cli_advanced.rst

@@ -54,9 +54,9 @@ Modify a checkpoint anywhere
 ****************************
 When you need to change the components of a checkpoint before saving or loading, use the :meth:`~lightning.pytorch.core.hooks.CheckpointHooks.on_save_checkpoint` and :meth:`~lightning.pytorch.core.hooks.CheckpointHooks.on_load_checkpoint` of your ``LightningModule``.
 
-.. code:: python
+.. code-block:: python
 
-    class LitModel(pl.LightningModule):
+    class LitModel(L.LightningModule):
         def on_save_checkpoint(self, checkpoint):
             checkpoint["something_cool_i_want_to_save"] = my_cool_pickable_object
 
@@ -65,9 +65,12 @@ When you need to change the components of a checkpoint before saving or loading,
 
 Use the above approach when you need to couple this behavior to your LightningModule for reproducibility reasons. Otherwise, Callbacks also have the :meth:`~lightning.pytorch.callbacks.callback.Callback.on_save_checkpoint` and :meth:`~lightning.pytorch.callbacks.callback.Callback.on_load_checkpoint` which you should use instead:
 
-.. code:: python
+.. code-block:: python
+
+    import lightning as L
+
 
-    class LitCallback(pl.Callback):
+    class LitCallback(L.Callback):
         def on_save_checkpoint(self, checkpoint):
             checkpoint["something_cool_i_want_to_save"] = my_cool_pickable_object
 

docs/source-pytorch/common/checkpointing_advanced.rst

@@ -127,7 +127,7 @@ In some cases, we may also pass entire PyTorch modules to the ``__init__`` metho
 
 .. code-block:: python
 
-    class LitAutoencoder(pl.LightningModule):
+    class LitAutoencoder(L.LightningModule):
         def __init__(self, encoder, decoder):
             ...
 
@@ -160,7 +160,7 @@ For example, let's pretend we created a LightningModule like so:
         ...
 
 
-    class Autoencoder(pl.LightningModule):
+    class Autoencoder(L.LightningModule):
         def __init__(self, encoder, decoder, *args, **kwargs):
             ...
 

docs/source-pytorch/common/checkpointing_basic.rst

@@ -27,7 +27,7 @@ Any value that has been logged via *self.log* in the LightningModule can be moni
 
 .. code-block:: python
 
-        class LitModel(pl.LightningModule):
+        class LitModel(L.LightningModule):
             def training_step(self, batch, batch_idx):
                 self.log("my_metric", x)
 

docs/source-pytorch/common/checkpointing_intermediate.rst

@@ -39,7 +39,7 @@ To add a test loop, implement the **test_step** method of the LightningModule
 
 .. code:: python
 
-    class LitAutoEncoder(pl.LightningModule):
+    class LitAutoEncoder(L.LightningModule):
         def training_step(self, batch, batch_idx):
             ...
 
@@ -99,7 +99,7 @@ To add a validation loop, implement the **validation_step** method of the Lightn
 
 .. code:: python
 
-    class LitAutoEncoder(pl.LightningModule):
+    class LitAutoEncoder(L.LightningModule):
         def training_step(self, batch, batch_idx):
             ...
 
@@ -126,5 +126,5 @@ To run the validation loop, pass in the validation set to **.fit**
    valid_loader = DataLoader(valid_set)
 
    # train with both splits
-   trainer = Trainer()
+   trainer = L.Trainer()
    trainer.fit(model, train_loader, valid_loader)