--unsafe-fixes

Author: Borda

.actions/assistant.py

@@ -22,7 +22,7 @@
 from itertools import chain
 from os.path import dirname, isfile
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Optional
 
 from packaging.requirements import Requirement
 from packaging.version import Version
@@ -128,7 +128,7 @@ def _parse_requirements(lines: Iterable[str]) -> Iterator[_RequirementWithCommen
         pip_argument = None
 
 
-def load_requirements(path_dir: str, file_name: str = "base.txt", unfreeze: str = "all") -> List[str]:
+def load_requirements(path_dir: str, file_name: str = "base.txt", unfreeze: str = "all") -> list[str]:
     """Loading requirements from a file.
 
     >>> path_req = os.path.join(_PROJECT_ROOT, "requirements")
@@ -223,7 +223,7 @@ def _load_aggregate_requirements(req_dir: str = "requirements", freeze_requireme
         fp.writelines([ln + os.linesep for ln in requires] + [os.linesep])
 
 
-def _retrieve_files(directory: str, *ext: str) -> List[str]:
+def _retrieve_files(directory: str, *ext: str) -> list[str]:
     all_files = []
     for root, _, files in os.walk(directory):
         for fname in files:
@@ -233,7 +233,7 @@ def _retrieve_files(directory: str, *ext: str) -> List[str]:
     return all_files
 
 
-def _replace_imports(lines: List[str], mapping: List[Tuple[str, str]], lightning_by: str = "") -> List[str]:
+def _replace_imports(lines: list[str], mapping: list[tuple[str, str]], lightning_by: str = "") -> list[str]:
     """Replace imports of standalone package to lightning.
 
     >>> lns = [
@@ -321,7 +321,7 @@ def copy_replace_imports(
             fo.writelines(lines)
 
 
-def create_mirror_package(source_dir: str, package_mapping: Dict[str, str]) -> None:
+def create_mirror_package(source_dir: str, package_mapping: dict[str, str]) -> None:
     """Create a mirror package with adjusted imports."""
     # replace imports and copy the code
     mapping = package_mapping.copy()

.pre-commit-config.yaml

@@ -74,7 +74,7 @@ repos:
     hooks:
       # try to fix what is possible
       - id: ruff
-        args: ["--fix"]
+        args: ["--fix", "--unsafe-fixes"]
       # perform formatting updates
       - id: ruff-format
       # validate if all is fine with preview mode

examples/fabric/build_your_own_trainer/trainer.py

@@ -1,7 +1,7 @@
 import os
 from collections.abc import Iterable, Mapping
 from functools import partial
-from typing import Any, List, Literal, Optional, Tuple, Union, cast
+from typing import Any, Literal, Optional, Union, cast
 
 import lightning as L
 import torch
@@ -19,11 +19,11 @@ def __init__(
         self,
         accelerator: Union[str, Accelerator] = "auto",
         strategy: Union[str, Strategy] = "auto",
-        devices: Union[List[int], str, int] = "auto",
+        devices: Union[list[int], str, int] = "auto",
         precision: Union[str, int] = "32-true",
         plugins: Optional[Union[str, Any]] = None,
-        callbacks: Optional[Union[List[Any], Any]] = None,
-        loggers: Optional[Union[Logger, List[Logger]]] = None,
+        callbacks: Optional[Union[list[Any], Any]] = None,
+        loggers: Optional[Union[Logger, list[Logger]]] = None,
         max_epochs: Optional[int] = 1000,
         max_steps: Optional[int] = None,
         grad_accum_steps: int = 1,
@@ -465,7 +465,7 @@ def get_latest_checkpoint(checkpoint_dir: str) -> Optional[str]:
 
     def _parse_optimizers_schedulers(
         self, configure_optim_output
-    ) -> Tuple[
+    ) -> tuple[
         Optional[L.fabric.utilities.types.Optimizable],
         Optional[Mapping[str, Union[L.fabric.utilities.types.LRScheduler, bool, str, int]]],
     ]:

examples/fabric/reinforcement_learning/rl/agent.py

@@ -1,5 +1,4 @@
 import math
-from typing import Dict, Tuple
 
 import gymnasium as gym
 import torch
@@ -43,7 +42,7 @@ def __init__(self, envs: gym.vector.SyncVectorEnv, act_fun: str = "relu", ortho_
             layer_init(torch.nn.Linear(64, envs.single_action_space.n), std=0.01, ortho_init=ortho_init),
         )
 
-    def get_action(self, x: Tensor, action: Tensor = None) -> Tuple[Tensor, Tensor, Tensor]:
+    def get_action(self, x: Tensor, action: Tensor = None) -> tuple[Tensor, Tensor, Tensor]:
         logits = self.actor(x)
         distribution = Categorical(logits=logits)
         if action is None:
@@ -58,12 +57,12 @@ def get_greedy_action(self, x: Tensor) -> Tensor:
     def get_value(self, x: Tensor) -> Tensor:
         return self.critic(x)
 
-    def get_action_and_value(self, x: Tensor, action: Tensor = None) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+    def get_action_and_value(self, x: Tensor, action: Tensor = None) -> tuple[Tensor, Tensor, Tensor, Tensor]:
         action, log_prob, entropy = self.get_action(x, action)
         value = self.get_value(x)
         return action, log_prob, entropy, value
 
-    def forward(self, x: Tensor, action: Tensor = None) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+    def forward(self, x: Tensor, action: Tensor = None) -> tuple[Tensor, Tensor, Tensor, Tensor]:
         return self.get_action_and_value(x, action)
 
     @torch.no_grad()
@@ -77,7 +76,7 @@ def estimate_returns_and_advantages(
         num_steps: int,
         gamma: float,
         gae_lambda: float,
-    ) -> Tuple[Tensor, Tensor]:
+    ) -> tuple[Tensor, Tensor]:
         next_value = self.get_value(next_obs).reshape(1, -1)
         advantages = torch.zeros_like(rewards)
         lastgaelam = 0
@@ -143,7 +142,7 @@ def __init__(
         self.avg_value_loss = MeanMetric(**torchmetrics_kwargs)
         self.avg_ent_loss = MeanMetric(**torchmetrics_kwargs)
 
-    def get_action(self, x: Tensor, action: Tensor = None) -> Tuple[Tensor, Tensor, Tensor]:
+    def get_action(self, x: Tensor, action: Tensor = None) -> tuple[Tensor, Tensor, Tensor]:
         logits = self.actor(x)
         distribution = Categorical(logits=logits)
         if action is None:
@@ -158,12 +157,12 @@ def get_greedy_action(self, x: Tensor) -> Tensor:
     def get_value(self, x: Tensor) -> Tensor:
         return self.critic(x)
 
-    def get_action_and_value(self, x: Tensor, action: Tensor = None) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+    def get_action_and_value(self, x: Tensor, action: Tensor = None) -> tuple[Tensor, Tensor, Tensor, Tensor]:
         action, log_prob, entropy = self.get_action(x, action)
         value = self.get_value(x)
         return action, log_prob, entropy, value
 
-    def forward(self, x: Tensor, action: Tensor = None) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+    def forward(self, x: Tensor, action: Tensor = None) -> tuple[Tensor, Tensor, Tensor, Tensor]:
         return self.get_action_and_value(x, action)
 
     @torch.no_grad()
@@ -177,7 +176,7 @@ def estimate_returns_and_advantages(
         num_steps: int,
         gamma: float,
         gae_lambda: float,
-    ) -> Tuple[Tensor, Tensor]:
+    ) -> tuple[Tensor, Tensor]:
         next_value = self.get_value(next_obs).reshape(1, -1)
         advantages = torch.zeros_like(rewards)
         lastgaelam = 0
@@ -193,7 +192,7 @@ def estimate_returns_and_advantages(
         returns = advantages + values
         return returns, advantages
 
-    def training_step(self, batch: Dict[str, Tensor]):
+    def training_step(self, batch: dict[str, Tensor]):
         # Get actions and values given the current observations
         _, newlogprob, entropy, newvalue = self(batch["obs"], batch["actions"].long())
         logratio = newlogprob - batch["logprobs"]

examples/fabric/reinforcement_learning/train_fabric.py

@@ -21,7 +21,6 @@
 import os
 import time
 from datetime import datetime
-from typing import Dict
 
 import gymnasium as gym
 import torch
@@ -38,7 +37,7 @@ def train(
     fabric: Fabric,
     agent: PPOLightningAgent,
     optimizer: torch.optim.Optimizer,
-    data: Dict[str, Tensor],
+    data: dict[str, Tensor],
     global_step: int,
     args: argparse.Namespace,
 ):

examples/fabric/reinforcement_learning/train_torch.py

@@ -22,7 +22,6 @@
 import random
 import time
 from datetime import datetime
-from typing import Dict
 
 import gymnasium as gym
 import torch
@@ -41,7 +40,7 @@
 def train(
     agent: PPOAgent,
     optimizer: torch.optim.Optimizer,
-    data: Dict[str, Tensor],
+    data: dict[str, Tensor],
     logger: SummaryWriter,
     global_step: int,
     args: argparse.Namespace,

examples/fabric/tensor_parallel/model.py

@@ -9,7 +9,7 @@
 
 
 from dataclasses import dataclass
-from typing import Optional, Tuple
+from typing import Optional
 
 import torch
 import torch.nn.functional as F
@@ -87,7 +87,7 @@ def apply_rotary_emb(
     xq: torch.Tensor,
     xk: torch.Tensor,
     freqs_cis: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor]:
     """Apply rotary embeddings to input tensors using the given frequency tensor.
 
     This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided

examples/pytorch/basics/autoencoder.py

@@ -18,7 +18,7 @@
 """
 
 from os import path
-from typing import Optional, Tuple
+from typing import Optional
 
 import torch
 import torch.nn.functional as F
@@ -45,7 +45,7 @@ def __init__(
         nrow: int = 8,
         padding: int = 2,
         normalize: bool = True,
-        value_range: Optional[Tuple[int, int]] = None,
+        value_range: Optional[tuple[int, int]] = None,
         scale_each: bool = False,
         pad_value: int = 0,
     ) -> None:

examples/pytorch/domain_templates/reinforce_learn_Qnet.py

@@ -36,7 +36,6 @@
 import random
 from collections import OrderedDict, deque, namedtuple
 from collections.abc import Iterator
-from typing import List, Tuple
 
 import gym
 import torch
@@ -103,7 +102,7 @@ def append(self, experience: Experience) -> None:
         """
         self.buffer.append(experience)
 
-    def sample(self, batch_size: int) -> Tuple:
+    def sample(self, batch_size: int) -> tuple:
         indices = random.sample(range(len(self.buffer)), batch_size)
         states, actions, rewards, dones, next_states = zip(*(self.buffer[idx] for idx in indices))
 
@@ -191,7 +190,7 @@ def get_action(self, net: nn.Module, epsilon: float, device: str) -> int:
         return action
 
     @torch.no_grad()
-    def play_step(self, net: nn.Module, epsilon: float = 0.0, device: str = "cpu") -> Tuple[float, bool]:
+    def play_step(self, net: nn.Module, epsilon: float = 0.0, device: str = "cpu") -> tuple[float, bool]:
         """Carries out a single interaction step between the agent and the environment.
 
         Args:
@@ -296,7 +295,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
         return self.net(x)
 
-    def dqn_mse_loss(self, batch: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
+    def dqn_mse_loss(self, batch: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
         """Calculates the mse loss using a mini batch from the replay buffer.
 
         Args:
@@ -319,7 +318,7 @@ def dqn_mse_loss(self, batch: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor
 
         return nn.MSELoss()(state_action_values, expected_state_action_values)
 
-    def training_step(self, batch: Tuple[torch.Tensor, torch.Tensor], nb_batch) -> OrderedDict:
+    def training_step(self, batch: tuple[torch.Tensor, torch.Tensor], nb_batch) -> OrderedDict:
         """Carries out a single step through the environment to update the replay buffer. Then calculates loss based on
         the minibatch received.
 
@@ -357,7 +356,7 @@ def training_step(self, batch: Tuple[torch.Tensor, torch.Tensor], nb_batch) -> O
 
         return OrderedDict({"loss": loss, "log": log, "progress_bar": log})
 
-    def configure_optimizers(self) -> List[Optimizer]:
+    def configure_optimizers(self) -> list[Optimizer]:
         """Initialize Adam optimizer."""
         optimizer = optim.Adam(self.net.parameters(), lr=self.lr)
         return [optimizer]

examples/pytorch/domain_templates/reinforce_learn_ppo.py

@@ -31,7 +31,7 @@
 
 import argparse
 from collections.abc import Iterator
-from typing import Callable, List, Tuple
+from typing import Callable
 
 import gym
 import torch
@@ -42,7 +42,7 @@
 from torch.utils.data import DataLoader, IterableDataset
 
 
-def create_mlp(input_shape: Tuple[int], n_actions: int, hidden_size: int = 128):
+def create_mlp(input_shape: tuple[int], n_actions: int, hidden_size: int = 128):
     """Simple Multi-Layer Perceptron network."""
     return nn.Sequential(
         nn.Linear(input_shape[0], hidden_size),
@@ -228,7 +228,7 @@ def __init__(
 
         self.state = torch.FloatTensor(self.env.reset())
 
-    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Passes in a state x through the network and returns the policy and a sampled action.
 
         Args:
@@ -243,7 +243,7 @@ def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Te
 
         return pi, action, value
 
-    def discount_rewards(self, rewards: List[float], discount: float) -> List[float]:
+    def discount_rewards(self, rewards: list[float], discount: float) -> list[float]:
         """Calculate the discounted rewards of all rewards in list.
 
         Args:
@@ -264,7 +264,7 @@ def discount_rewards(self, rewards: List[float], discount: float) -> List[float]
 
         return list(reversed(cumul_reward))
 
-    def calc_advantage(self, rewards: List[float], values: List[float], last_value: float) -> List[float]:
+    def calc_advantage(self, rewards: list[float], values: list[float], last_value: float) -> list[float]:
         """Calculate the advantage given rewards, state values, and the last value of episode.
 
         Args:
@@ -282,7 +282,7 @@ def calc_advantage(self, rewards: List[float], values: List[float], last_value:
         delta = [rews[i] + self.gamma * vals[i + 1] - vals[i] for i in range(len(rews) - 1)]
         return self.discount_rewards(delta, self.gamma * self.lam)
 
-    def generate_trajectory_samples(self) -> Tuple[List[torch.Tensor], List[torch.Tensor], List[torch.Tensor]]:
+    def generate_trajectory_samples(self) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]:
         """
         Contains the logic for generating trajectory data to train policy and value network
         Yield:
@@ -376,7 +376,7 @@ def critic_loss(self, state, action, logp_old, qval, adv) -> torch.Tensor:
         value = self.critic(state)
         return (qval - value).pow(2).mean()
 
-    def training_step(self, batch: Tuple[torch.Tensor, torch.Tensor]):
+    def training_step(self, batch: tuple[torch.Tensor, torch.Tensor]):
         """Carries out a single update to actor and critic network from a batch of replay buffer.
 
         Args:
@@ -407,7 +407,7 @@ def training_step(self, batch: Tuple[torch.Tensor, torch.Tensor]):
         self.log("loss_critic", loss_critic, on_step=False, on_epoch=True, prog_bar=False, logger=True)
         self.log("loss_actor", loss_actor, on_step=False, on_epoch=True, prog_bar=True, logger=True)
 
-    def configure_optimizers(self) -> List[Optimizer]:
+    def configure_optimizers(self) -> list[Optimizer]:
         """Initialize Adam optimizer."""
         optimizer_actor = torch.optim.Adam(self.actor.parameters(), lr=self.lr_actor)
         optimizer_critic = torch.optim.Adam(self.critic.parameters(), lr=self.lr_critic)