pytorch-optimizer 3.8.1__tar.gz → 3.8.2__tar.gz

{pytorch_optimizer-3.8.1 → pytorch_optimizer-3.8.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytorch_optimizer
-Version: 3.8.1
+Version: 3.8.2
 Summary: optimizer & lr scheduler & objective function collections in PyTorch
 License: Apache-2.0
 License-File: LICENSE
@@ -31,8 +31,9 @@ Classifier: Topic :: Software Development :: Libraries
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Dist: numpy (<=1.24.4) ; python_version < "3.9"
 Requires-Dist: numpy (>1.24.4) ; python_version >= "3.9"
-Requires-Dist: torch (>=1.10) ; python_version >= "3.9"
+Requires-Dist: torch (>=1.10) ; python_version >= "3.10"
 Requires-Dist: torch (>=1.10,<2.5) ; python_version < "3.9"
+Requires-Dist: torch (>=1.10,<2.9) ; python_version == "3.9"
 Project-URL: Documentation, https://pytorch-optimizers.readthedocs.io/en/latest
 Project-URL: Homepage, https://github.com/kozistr/pytorch_optimizer
 Project-URL: Repository, https://github.com/kozistr/pytorch_optimizer

{pytorch_optimizer-3.8.1 → pytorch_optimizer-3.8.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "3.8.1"
+version = "3.8.2"
 description = "optimizer & lr scheduler & objective function collections in PyTorch"
 license = "Apache-2.0"
 authors = ["kozistr <kozistr@gmail.com>"]
@@ -37,6 +37,7 @@ classifiers = [
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
     "Operating System :: OS Independent",
     "Topic :: Scientific/Engineering",
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
@@ -53,7 +54,8 @@ numpy = [
 ]
 torch = [
     { version = ">=1.10,<2.5", python = "<3.9", source = "torch" },
-    { version = ">=1.10", python = ">=3.9", source = "torch" },
+    { version = ">=1.10,<2.9", python = ">=3.9,<3.10", source = "torch" },
+    { version = ">=1.10", python = ">=3.10", source = "torch" },
 ]
 
 [tool.poetry.group.dev.dependencies]

{pytorch_optimizer-3.8.1 → pytorch_optimizer-3.8.2}/pytorch_optimizer/base/exception.py

@@ -1,8 +1,9 @@
 class NoSparseGradientError(Exception):
-    """Raised when the gradient is sparse gradient.
+    r"""Raised when the gradient is sparse gradient.
 
-    :param optimizer_name: str. optimizer name.
-    :param note: str. special conditions to note (default '').
+    Args:
+        optimizer_name (str): The name of the optimizer where the error occurred.
+        note (str): Additional special conditions or notes (default is an empty string).
     """
 
     def __init__(self, optimizer_name: str, note: str = ''):
@@ -46,10 +47,11 @@ class NegativeStepError(Exception):
 
 
 class NoComplexParameterError(Exception):
-    """Raised when the dtype of the parameter is complex.
+    r"""Raised when the dtype of the parameter is complex.
 
-    :param optimizer_name: str. optimizer name.
-    :param note: str. special conditions to note (default '').
+    Args:
+        optimizer_name (str): The name of the optimizer where the error occurred.
+        note (str): Additional special conditions or notes (default is an empty string).
     """
 
     def __init__(self, optimizer_name: str, note: str = ''):

{pytorch_optimizer-3.8.1 → pytorch_optimizer-3.8.2}/pytorch_optimizer/base/optimizer.py

@@ -4,31 +4,30 @@ from typing import List, Optional, Tuple, Union
 
 import torch
 from torch.optim import Optimizer
-from torch.optim.optimizer import ParamsT
 
 from pytorch_optimizer.base.exception import NegativeLRError, NegativeStepError
 from pytorch_optimizer.base.type import (
-    BETAS,
-    CLOSURE,
-    DEFAULTS,
-    GROUP,
     HUTCHINSON_G,
-    LOSS,
     OPTIMIZER_INSTANCE_OR_CLASS,
-    PARAMETERS,
-    STATE,
+    Betas,
+    Closure,
+    Defaults,
+    Loss,
+    Parameters,
+    ParamGroup,
+    State,
 )
 
 
 class BaseOptimizer(ABC, Optimizer):
-    r"""Base optimizer class. Provides common functionalities for the optimizers."""
+    """Base optimizer class. Provides common functionalities for the optimizers."""
 
-    def __init__(self, params: Union[PARAMETERS, ParamsT], defaults: DEFAULTS) -> None:
+    def __init__(self, params: Parameters, defaults: Defaults) -> None:
         super().__init__(params, defaults)
 
     @staticmethod
     def load_optimizer(optimizer: OPTIMIZER_INSTANCE_OR_CLASS, **kwargs) -> Optimizer:
-        r"""Build torch.optim.Optimizer class."""
+        """Build torch.optim.Optimizer class."""
         if isinstance(optimizer, Optimizer):
             return optimizer
 
@@ -40,23 +39,23 @@ class BaseOptimizer(ABC, Optimizer):
 
     @staticmethod
     @torch.no_grad()
-    def set_hessian(param_groups: PARAMETERS, state: STATE, hessian: List[torch.Tensor]) -> None:
-        r"""Set hessian to state from external source. Generally useful when using functorch as a base.
+    def set_hessian(param_groups: Parameters, state: State, hessian: List[torch.Tensor]) -> None:
+        """Set hessian to state from external source. Generally useful when using functorch as a base.
+
+        Args:
+            param_groups: PARAMETERS. Parameter groups from optimizer.
+            state: STATE. Optimizer state dictionary.
+            hessian: List[torch.Tensor]. Sequence of Hessian tensors to set.
 
         Example:
-        -------
-            # Hutchinson's Estimator using HVP
-            noise = tree_map(lambda v: torch.randn_like(v), params)
-            loss_, hvp_est = jvp(grad(run_model_fn), (params,), (noise,))
-            hessian_diag_est  = tree_map(lambda a, b: a * b, hvp_est, noise)
+            # Hutchinson's Estimator using Hessian-vector product (HVP)
+            >>> noise = tree_map(lambda v: torch.randn_like(v), params)
+            >>> loss_, hvp_est = jvp(grad(run_model_fn), (params,), (noise,))
+            >>> hessian_diag_est = tree_map(lambda a, b: a * b, hvp_est, noise)
 
-            optimizer.set_hessian(hessian_diag_est)
+            >>> optimizer.set_hessian(hessian_diag_est)
             # OR
-            optimizer.step(hessian=hessian_diag_est)
-
-        :param param_groups: PARAMETERS. parameter groups.
-        :param state: STATE. optimizer state.
-        :param hessian: List[torch.Tensor]. sequence of hessian to set.
+            >>> optimizer.step(hessian=hessian_diag_est)
         """
         i: int = 0
         for group in param_groups:
@@ -70,12 +69,13 @@ class BaseOptimizer(ABC, Optimizer):
                 i += 1
 
     @staticmethod
-    def zero_hessian(param_groups: PARAMETERS, state: STATE, pre_zero: bool = True) -> None:
-        r"""Zero-out hessian.
+    def zero_hessian(param_groups: Parameters, state: State, pre_zero: bool = True) -> None:
+        """Zero-out Hessian.
 
-        :param param_groups: PARAMETERS. parameter groups.
-        :param state: STATE. optimizer state.
-        :param pre_zero: bool. zero-out hessian before computing the hessian.
+        Args:
+            param_groups (Parameters): Parameter groups from the optimizer.
+            state (State): Optimizer state dictionary.
+            pre_zero (bool): If True, zero-out the Hessian before computing/updating it.
         """
         for group in param_groups:
             for p in group['params']:
@@ -88,19 +88,20 @@ class BaseOptimizer(ABC, Optimizer):
     @staticmethod
     @torch.no_grad()
     def compute_hutchinson_hessian(
-        param_groups: PARAMETERS,
-        state: STATE,
+        param_groups: Parameters,
+        state: State,
         num_samples: int = 1,
         alpha: float = 1.0,
         distribution: HUTCHINSON_G = 'gaussian',
     ) -> None:
-        r"""Hutchinson's approximate hessian, added to the state under key `hessian`.
-
-        :param param_groups: PARAMETERS. parameter groups.
-        :param state: STATE. optimizer state.
-        :param num_samples: int. number of times to sample `z` for the approximation of the hessian trace.
-        :param alpha: float. alpha.
-        :param distribution: HUTCHINSON_G. type of distribution.
+        r"""Hutchinson's approximate Hessian, added to the state under key `hessian`.
+
+        Args:
+            param_groups (Parameters): Parameter groups from the optimizer.
+            state (State): Optimizer state dictionary.
+            num_samples (int): Number of times to sample noise vector `z` for the trace approximation.
+            alpha (float): Scaling factor for the Hessian estimate.
+            distribution (HUTCHINSON_G): Type of noise distribution used (e.g., Rademacher).
         """
         if distribution not in ('gaussian', 'rademacher'):
             raise NotImplementedError(f'hessian with distribution {distribution} is not implemented.')
@@ -136,15 +137,16 @@ class BaseOptimizer(ABC, Optimizer):
         fixed_decay: bool,
         ratio: Optional[float] = None,
     ) -> None:
-        r"""Apply weight decay.
-
-        :param p: torch.Tensor. parameter.
-        :param grad: torch.Tensor. gradient.
-        :param lr: float. learning rate.
-        :param weight_decay: float. weight decay (L2 penalty).
-        :param weight_decouple: bool. the optimizer uses decoupled weight decay as in AdamW.
-        :param fixed_decay: bool. fix weight decay.
-        :param ratio: Optional[float]. scale weight decay.
+        """Apply weight decay.
+
+        Args:
+            p (torch.Tensor): Parameter tensor to apply weight decay to.
+            grad (torch.Tensor): Gradient tensor of parameter p.
+            lr (float): Learning rate to scale the update.
+            weight_decay (float): Weight decay coefficient (L2 penalty).
+            weight_decouple (bool): If True, applies decoupled weight decay as in AdamW.
+            fixed_decay (bool): If True, fixes weight decay to not depend on learning rate.
+            ratio (Optional[float]): Optional scaling factor for weight decay.
         """
         if weight_decouple:
             p.mul_(1.0 - weight_decay * (1.0 if fixed_decay else lr) * (ratio if ratio is not None else 1.0))
@@ -159,13 +161,14 @@ class BaseOptimizer(ABC, Optimizer):
         eps: float,
         exp_avg_sq_eps: float = 1e-15,
     ) -> torch.Tensor:
-        r"""Apply AMSBound variant.
-
-        :param ams_bound: bool. whether to apply AMSBound.
-        :param exp_avg_sq: torch.Tensor. exp_avg_sq.
-        :param max_exp_avg_sq: Optional[torch.Tensor]. max_exp_avg_sq.
-        :param eps: float. epsilon.
-        :param exp_avg_sq_eps: float. eps value for numerical stability for exp_avg_sq.
+        """Apply AMSBound variant.
+
+        Args:
+            ams_bound (bool): Whether to apply the AMSBound variant.
+            exp_avg_sq (torch.Tensor): Exponential moving average of squared gradients.
+            max_exp_avg_sq (Optional[torch.Tensor]): Maximum of all exp_avg_sq elements, for AMSBound.
+            eps (float): Small epsilon value for numerical stability.
+            exp_avg_sq_eps (float): Epsilon used specifically for numerical stability in exp_avg_sq computations.
         """
         if ams_bound:
             if torch.is_complex(max_exp_avg_sq):
@@ -180,10 +183,11 @@ class BaseOptimizer(ABC, Optimizer):
 
     @staticmethod
     def debias(beta: float, step: int) -> float:
-        r"""Adam-style debias correction. Returns `1.0 - beta ** step`.
+        """Adam-style debias correction.
 
-        :param beta: float. beta.
-        :param step: int. number of step.
+        Args:
+            beta (float): Exponential decay rate for moment estimates.
+            step (int): Current optimization step number.
         """
         return 1.0 - math.pow(beta, step)  # fmt: skip
 
@@ -193,19 +197,21 @@ class BaseOptimizer(ABC, Optimizer):
 
         Simplified version of `\^{beta} = beta * (1.0 - beta ** (step - 1)) / (1.0 - beta ** step)`
 
-        :param beta: float. beta.
-        :param step: int. number of step.
+        Args:
+            beta (float): The original beta decay rate.
+            step (int): Current optimization step number.
         """
         beta_n: float = math.pow(beta, step)
         return (beta_n - beta) / (beta_n - 1.0)  # fmt: skip
 
     @staticmethod
     def apply_adam_debias(adam_debias: bool, step_size: float, bias_correction1: float) -> float:
-        r"""Apply AdamD variant.
+        """Apply AdamD variant.
 
-        :param adam_debias: bool. Only correct the denominator to avoid inflating step sizes early in training.
-        :param step_size: float. step size.
-        :param bias_correction1: float. bias_correction.
+        Args:
+            adam_debias (bool): If True, only corrects the denominator to avoid inflating step sizes early in training.
+            step_size (float): The step size for the update.
+            bias_correction1 (float): The bias correction factor for the first moment.
         """
         return step_size if adam_debias else step_size / bias_correction1
 
@@ -218,14 +224,15 @@ class BaseOptimizer(ABC, Optimizer):
         n_sma_threshold: int,
         degenerated_to_sgd: bool,
     ) -> Tuple[float, float]:
-        r"""Get step size for rectify optimizer.
-
-        :param is_rectify: bool. whether to apply rectify-variant.
-        :param step: int. number of steps.
-        :param lr: float. learning rate.
-        :param beta2: float. beta2.
-        :param n_sma_threshold: float. SMA threshold.
-        :param degenerated_to_sgd: bool. degenerated to SGD.
+        """Get step size for rectify optimizer.
+
+        Args:
+            is_rectify (bool): Whether to apply the rectify variant.
+            step (int): Current step number.
+            lr (float): Base learning rate.
+            beta2 (float): Beta2 parameter from optimizer (momentum term).
+            n_sma_threshold (float): Simple Moving Average (SMA) threshold for rectification.
+            degenerated_to_sgd (bool): Whether to degenerate to SGD if below threshold.
         """
         step_size: float = lr
         n_sma: float = 0.0
@@ -254,10 +261,11 @@ class BaseOptimizer(ABC, Optimizer):
     ) -> torch.Tensor:
         r"""Get AdaNorm gradient.
 
-        :param grad: torch.Tensor. gradient.
-        :param adanorm: bool. whether to use the AdaNorm variant.
-        :param exp_grad_norm: Optional[torch.Tensor]. exp_grad_norm.
-        :param r: Optional[float]. EMA factor. between 0.9 ~ 0.99 is preferred.
+        Args:
+            grad (torch.Tensor): Gradient.
+            adanorm (bool): Whether to use the AdaNorm variant.
+            exp_grad_norm (Optional[torch.Tensor]): Exponential moving average of gradient norm.
+            r (Optional[float]): EMA factor; between 0.9 and 0.99 is preferred.
         """
         if not adanorm or exp_grad_norm is None:
             return grad
@@ -273,7 +281,7 @@ class BaseOptimizer(ABC, Optimizer):
 
     @staticmethod
     def get_rms(x: torch.Tensor) -> torch.Tensor:
-        r"""Get RMS."""
+        """Get RMS."""
         return x.norm(2) / math.sqrt(x.numel())
 
     @staticmethod
@@ -282,17 +290,18 @@ class BaseOptimizer(ABC, Optimizer):
         exp_avg_sq_col: torch.Tensor,
         output: torch.Tensor,
     ) -> None:
-        r"""Get approximation of EMA of squared gradient."""
+        """Get approximation of EMA of squared gradient."""
         r_factor: torch.Tensor = (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1, keepdim=True)).rsqrt_().unsqueeze(-1)
         c_factor: torch.Tensor = exp_avg_sq_col.unsqueeze(-2).rsqrt()
         torch.mul(r_factor, c_factor, out=output)
 
     @staticmethod
     def apply_cautious(update: torch.Tensor, grad: torch.Tensor) -> None:
-        r"""Apply the Cautious Optimizer feature.
+        """Apply the Cautious Optimizer feature.
 
-        :param update: torch.Tensor. update. it'll be masked in in-place manner.
-        :param grad: torch.Tensor. gradient.
+        Args:
+            update (torch.Tensor): update. it'll be masked in in-place manner.
+            grad (torch.Tensor): gradient.
         """
         mask = (update * grad > 0).to(grad.dtype)
         mask.mul_(mask.numel() / (mask.sum() + 1))
@@ -300,11 +309,12 @@ class BaseOptimizer(ABC, Optimizer):
 
     @staticmethod
     def get_stable_adamw_rms(grad: torch.Tensor, exp_avg_sq: torch.Tensor, eps: float = 1e-16) -> float:
-        r"""Get StableAdamW RMS.
+        """Get StableAdamW RMS.
 
-        :param grad: torch.Tensor. gradient.
-        :param exp_avg_sq: torch.Tensor. exp_avg_sq.
-        :param eps: float. epsilon.
+        Args:
+            grad (torch.Tensor): gradient.
+            exp_avg_sq (torch.Tensor): Exponential moving average of squared gradient.
+            eps (float): Small value to prevent division by zero.
         """
         return grad.pow(2).div_(exp_avg_sq.clip(min=eps)).mean().sqrt_().clip_(min=1.0).item()
 
@@ -362,7 +372,7 @@ class BaseOptimizer(ABC, Optimizer):
         if x % y != 0:
             raise ValueError(f'{x} must be divisible by {y}')
 
-    def validate_betas(self, betas: BETAS, beta_range_type: str = '[)', beta3_range_type: str = '[]') -> None:
+    def validate_betas(self, betas: Betas, beta_range_type: str = '[)', beta3_range_type: str = '[]') -> None:
         if betas[0] is not None:
             self.validate_range(betas[0], 'beta1', 0.0, 1.0, range_type=beta_range_type)
 
@@ -382,13 +392,13 @@ class BaseOptimizer(ABC, Optimizer):
             self.validate_range(nus[1], 'nu2', 0.0, 1.0, range_type='[]')
 
     @abstractmethod
-    def init_group(self, group: GROUP, **kwargs) -> None:  # pragma: no cover
-        r"""Initialize the group of the optimizer and return is_complex."""
+    def init_group(self, group: ParamGroup, **kwargs) -> None:  # pragma: no cover
+        """Initialize the group of the optimizer and return is_complex."""
         return
 
     @staticmethod
     def view_as_real(param, *state_and_grads) -> tuple:
-        r"""View imaginary tensors as real tensors."""
+        """View imaginary tensors as real tensors."""
         if torch.is_complex(param):
             param = torch.view_as_real(param)
             state_and_grads = tuple(
@@ -400,9 +410,9 @@ class BaseOptimizer(ABC, Optimizer):
 
     @staticmethod
     def maximize_gradient(grad: torch.Tensor, maximize: bool = False) -> None:
-        r"""Maximize the objective with respect to the params, instead of minimizing."""
+        """Maximize the objective with respect to the params, instead of minimizing."""
         if maximize:
             grad.neg_()
 
-    def step(self, closure: CLOSURE = None) -> LOSS:  # pragma: no cover
+    def step(self, closure: Closure = None) -> Loss:  # pragma: no cover
         raise NotImplementedError

{pytorch_optimizer-3.8.1 → pytorch_optimizer-3.8.2}/pytorch_optimizer/base/scheduler.py

@@ -7,16 +7,18 @@ from pytorch_optimizer.base.exception import NegativeLRError, NegativeStepError
 
 
 class BaseLinearWarmupScheduler(ABC):
-    r"""BaseLinearWarmupScheduler class.
-
-        The LR Scheduler class based on this class has linear warmup strategy.
-
-    :param optimizer: Optimizer. It will set learning rate to all trainable parameters in optimizer.
-    :param t_max: int. total steps to train.
-    :param max_lr: float. maximum lr.
-    :param min_lr: float. minimum lr.
-    :param init_lr: float. initial lr.
-    :param warmup_steps: int. steps to warm-up.
+    """BaseLinearWarmupScheduler class.
+
+    A learning rate scheduler class that implements a linear warmup strategy.
+
+    Args:
+        optimizer (Optimizer): The optimizer whose learning rate will be scheduled.
+            It will set the learning rate to all trainable parameters in the optimizer.
+        t_max (int): Total number of training steps (epochs or iterations).
+        max_lr (float): The maximum learning rate after warmup.
+        min_lr (float): The minimum learning rate to decay to (or start from if warmup).
+        init_lr (float): Initial learning rate at the start of warmup.
+        warmup_steps (int): Number of steps to warm up linearly from init_lr to max_lr.
     """
 
     def __init__(

pytorch_optimizer-3.8.2/pytorch_optimizer/base/type.py

@@ -0,0 +1,27 @@
+from typing import Any, Callable, Dict, Iterable, Literal, Optional, Tuple, Type, Union
+
+import torch
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LRScheduler
+
+OPTIMIZER = Type[Optimizer]
+OPTIMIZER_INSTANCE_OR_CLASS = Union[OPTIMIZER, Optimizer]
+SCHEDULER = Type[LRScheduler]
+
+Defaults = Dict[str, Any]
+ParamGroup = Dict[str, Any]
+State = Dict
+Parameters = Optional[Union[Iterable[torch.Tensor], Iterable[ParamGroup]]]
+
+Closure = Optional[Callable[[], float]]
+Loss = Optional[float]
+Betas = Union[
+    Tuple[float, float],
+    Tuple[float, float, float],
+    Tuple[None, float],
+]
+
+HUTCHINSON_G = Literal['gaussian', 'rademacher']
+CLASS_MODE = Literal['binary', 'multiclass', 'multilabel']
+
+DATA_FORMAT = Literal['channels_first', 'channels_last']

{pytorch_optimizer-3.8.1 → pytorch_optimizer-3.8.2}/pytorch_optimizer/loss/bi_tempered.py

@@ -5,21 +5,22 @@ from torch import nn
 
 
 def log_t(u: torch.Tensor, t: float) -> torch.Tensor:
-    r"""Compute log_t for `u'."""
+    """Compute log_t for `u'."""
     return u.log() if t == 1.0 else (u.pow(1.0 - t) - 1.0) / (1.0 - t)
 
 
 def exp_t(u: torch.Tensor, t: float) -> torch.Tensor:
-    r"""Compute exp_t for `u'."""
+    """Compute exp_t for `u'."""
     return u.exp() if t == 1 else (1.0 + (1.0 - t) * u).relu().pow(1.0 / (1.0 - t))
 
 
 def compute_normalization_fixed_point(activations: torch.Tensor, t: float, num_iters: int) -> torch.Tensor:
     r"""Return the normalization value for each example (t > 1.0).
 
-    :param activations: torch.Tensor. A multidimensional tensor with last dimension `num_classes`.
-    :param t: float. Temperature (> 1.0 for tail heaviness).
-    :param num_iters: int. Number of iterations to run the method.
+    Args:
+        activations (torch.Tensor): A multi-dimensional tensor with the last dimension representing classes.
+        t (float): Temperature value (> 1.0 for tail heaviness).
+        num_iters (int): Number of iterations to run the method.
     """
     mu, _ = torch.max(activations, dim=-1, keepdim=True)
 
@@ -39,9 +40,10 @@ def compute_normalization_fixed_point(activations: torch.Tensor, t: float, num_i
 def compute_normalization_binary_search(activations: torch.Tensor, t: float, num_iters: int) -> torch.Tensor:
     """Compute normalization value for each example (t < 1.0).
 
-    :param activations: torch.Tensor. A multidimensional tensor with last dimension `num_classes`.
-    :param t: float. Temperature (> 1.0 for tail heaviness).
-    :param num_iters: int. Number of iterations to run the method.
+    Args:
+        activations (torch.Tensor): A multidimensional tensor with the last dimension `num_classes`.
+        t (float): Temperature parameter (< 1.0 for peak sharpening).
+        num_iters (int): Number of iterations to run the normalization.
     """
     mu, _ = torch.max(activations, dim=-1, keepdim=True)
     normalized_activations = activations - mu
@@ -68,7 +70,7 @@ def compute_normalization_binary_search(activations: torch.Tensor, t: float, num
 
 
 class ComputeNormalization(torch.autograd.Function):
-    r"""Custom backward pass for compute_normalization. See compute_normalization."""
+    """Custom backward pass for compute_normalization. See compute_normalization."""
 
     @staticmethod
     def forward(ctx, activations: torch.Tensor, t: float, num_iters: int) -> torch.Tensor:
@@ -100,19 +102,21 @@ class ComputeNormalization(torch.autograd.Function):
 def compute_normalization(activations: torch.Tensor, t: float, num_iters: int = 5) -> torch.Tensor:
     r"""Compute normalization value for each example.
 
-    :param activations: torch.Tensor. A multidimensional tensor with last dimension `num_classes`.
-    :param t: float. Temperature (> 1.0 for tail heaviness).
-    :param num_iters: int. Number of iterations to run the method.
+    Args:
+        activations (torch.Tensor): A multi-dimensional tensor with the last dimension `num_classes`.
+        t (float): Temperature parameter (> 1.0 for tail heaviness).
+        num_iters (int): Number of iterations to run the method.
     """
     return ComputeNormalization.apply(activations, t, num_iters)
 
 
 def tempered_softmax(activations: torch.Tensor, t: float, num_iters: int = 5) -> torch.Tensor:
-    r"""Tempered softmax function.
+    """Tempered softmax function.
 
-    :param activations: torch.Tensor. A multidimensional tensor with last dimension `num_classes`.
-    :param t: float. Temperature (> 1.0 for tail heaviness).
-    :param num_iters: int. Number of iterations to run the method.
+    Args:
+        activations (torch.Tensor): A multidimensional tensor with last dimension `num_classes`.
+        t (float): Temperature parameter (> 1.0 for tail heaviness).
+        num_iters (int): Number of iterations to run the method.
     """
     if t == 1.0:
         return activations.softmax(dim=-1)
@@ -133,14 +137,15 @@ def bi_tempered_logistic_loss(
 ) -> torch.Tensor:
     r"""Bi-Tempered Logistic Loss.
 
-    :param activations: torch.Tensor. A multidimensional tensor with last dimension `num_classes`.
-    :param labels: torch.Tensor. A tensor with shape and dtype as activations (onehot), or a long tensor of
-        one dimension less than activations (pytorch standard)
-    :param t1: float. Temperature 1 (< 1.0 for boundedness).
-    :param t2: float. Temperature 2 (> 1.0 for tail heaviness, < 1.0 for finite support).
-    :param label_smooth: float. Label smoothing parameter between [0, 1).
-    :param num_iters: int. Number of iterations to run the method.
-    :param reduction: str. type of reduction.
+    Args:
+        activations (torch.Tensor): A multidimensional tensor with last dimension `num_classes`.
+        labels (torch.Tensor): Tensor with the same shape and dtype as activations (one-hot encoded),
+            or a long tensor with one dimension less (class indices).
+        t1 (float): Temperature 1 (< 1.0 for boundedness of loss).
+        t2 (float): Temperature 2 (> 1.0 for tail heaviness, < 1.0 for finite support).
+        label_smooth (float): Label smoothing parameter, between 0 and 1.
+        num_iters (int): Number of iterations to run the normalization method.
+        reduction (str): Specifies reduction method to apply to output: 'none', 'mean', or 'sum'.
     """
     if len(labels.shape) < len(activations.shape):
         labels_onehot = torch.zeros_like(activations)
@@ -174,13 +179,15 @@ def bi_tempered_logistic_loss(
 class BiTemperedLogisticLoss(nn.Module):
     """Bi-Tempered Log Loss.
 
-    Reference : https://github.com/BloodAxe/pytorch-toolbelt/blob/develop/pytorch_toolbelt/losses/bitempered_loss.py
+    Reference:
+        https://github.com/BloodAxe/pytorch-toolbelt/blob/develop/pytorch_toolbelt/losses/bitempered_loss.py
 
-    :param t1: float. Temperature 1 (< 1.0 for boundedness).
-    :param t2: float. Temperature 2 (> 1.0 for tail heaviness, < 1.0 for finite support).
-    :param label_smooth: float. Label smoothing parameter between [0, 1).
-    :param ignore_index: Optional[int]. Index to ignore.
-    :param reduction: str. type of reduction.
+    Args:
+        t1 (float): Temperature 1 (< 1.0 for boundedness).
+        t2 (float): Temperature 2 (> 1.0 for tail heaviness, < 1.0 for finite support).
+        label_smooth (float): Label smoothing parameter between 0 and 1.
+        ignore_index (Optional[int]): Index to ignore during loss calculation.
+        reduction (str): Type of reduction to apply to output, e.g. 'mean', 'sum', or 'none'.
     """
 
     def __init__(
@@ -215,13 +222,15 @@ class BiTemperedLogisticLoss(nn.Module):
 
 
 class BinaryBiTemperedLogisticLoss(nn.Module):
-    """Modification of BiTemperedLogisticLoss for binary classification case.
-
-    :param t1: float. Temperature 1 (< 1.0 for boundedness).
-    :param t2: float. Temperature 2 (> 1.0 for tail heaviness, < 1.0 for finite support).
-    :param label_smooth: float. Label smoothing parameter between [0, 1).
-    :param ignore_index: Optional[int]. Index to ignore.
-    :param reduction: str. type of reduction.
+    """Bi-Tempered Logistic Loss for Binary Classification.
+
+    Args:
+        t1 (float): Temperature 1 (< 1.0 for boundedness of the loss).
+        t2 (float): Temperature 2 (> 1.0 for tail heaviness, < 1.0 for finite support).
+        label_smooth (float): Label smoothing parameter between 0 and 1.
+        ignore_index (Optional[int]): Specifies a target value that is ignored and does not contribute
+            to the input gradient.
+        reduction (str): Specifies the reduction to apply to the output: 'none', 'mean', or 'sum'.
     """
 
     def __init__(