Refactor: allow reduction None (#891)

* allow reduction None
* Literal

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

torchmetrics/classification/jaccard.py

@@ -15,6 +15,7 @@
 
 import torch
 from torch import Tensor
+from typing_extensions import Literal
 
 from torchmetrics.classification.confusion_matrix import ConfusionMatrix
 from torchmetrics.functional.classification.jaccard import _jaccard_from_confmat
@@ -89,7 +90,7 @@ def __init__(
         absent_score: float = 0.0,
         threshold: float = 0.5,
         multilabel: bool = False,
-        reduction: str = "elementwise_mean",
+        reduction: Literal["elementwise_mean", "sum", "none", None] = "elementwise_mean",
         compute_on_step: Optional[bool] = None,
         **kwargs: Dict[str, Any],
     ) -> None:

torchmetrics/classification/kl_divergence.py

@@ -15,6 +15,7 @@
 
 import torch
 from torch import Tensor
+from typing_extensions import Literal
 
 from torchmetrics.functional.classification.kl_divergence import _kld_compute, _kld_update
 from torchmetrics.metric import Metric
@@ -79,7 +80,7 @@ class KLDivergence(Metric):
     def __init__(
         self,
         log_prob: bool = False,
-        reduction: Optional[str] = "mean",
+        reduction: Literal["mean", "sum", "none", None] = "mean",
         compute_on_step: Optional[bool] = None,
         **kwargs: Dict[str, Any],
     ) -> None:

torchmetrics/functional/classification/dice.py

@@ -15,6 +15,7 @@
 
 import torch
 from torch import Tensor
+from typing_extensions import Literal
 
 from torchmetrics.utilities.data import to_categorical
 from torchmetrics.utilities.distributed import reduce
@@ -64,7 +65,7 @@ def dice_score(
     bg: bool = False,
     nan_score: float = 0.0,
     no_fg_score: float = 0.0,
-    reduction: str = "elementwise_mean",
+    reduction: Literal["elementwise_mean", "sum", "none", None] = "elementwise_mean",
 ) -> Tensor:
     """Compute dice score from prediction scores.
 
@@ -78,7 +79,7 @@ def dice_score(
 
             - ``'elementwise_mean'``: takes the mean (default)
             - ``'sum'``: takes the sum
-            - ``'none'``: no reduction will be applied
+            - ``'none'`` or ``None``: no reduction will be applied
 
     Return:
         Tensor containing dice score

torchmetrics/functional/classification/jaccard.py

@@ -15,6 +15,7 @@
 
 import torch
 from torch import Tensor
+from typing_extensions import Literal
 
 from torchmetrics.functional.classification.confusion_matrix import _confusion_matrix_update
 from torchmetrics.utilities.data import get_num_classes
@@ -26,7 +27,7 @@ def _jaccard_from_confmat(
     num_classes: int,
     ignore_index: Optional[int] = None,
     absent_score: float = 0.0,
-    reduction: str = "elementwise_mean",
+    reduction: Literal["elementwise_mean", "sum", "none", None] = "elementwise_mean",
 ) -> Tensor:
     """Computes the intersection over union from confusion matrix.
 
@@ -41,7 +42,7 @@ def _jaccard_from_confmat(
 
             - ``'elementwise_mean'``: takes the mean (default)
             - ``'sum'``: takes the sum
-            - ``'none'``: no reduction will be applied
+            - ``'none'`` or ``None``: no reduction will be applied
     """
 
     # Remove the ignored class index from the scores.
@@ -73,7 +74,7 @@ def jaccard_index(
     ignore_index: Optional[int] = None,
     absent_score: float = 0.0,
     threshold: float = 0.5,
-    reduction: str = "elementwise_mean",
+    reduction: Literal["elementwise_mean", "sum", "none", None] = "elementwise_mean",
 ) -> Tensor:
     r"""
     Computes `Jaccard index`_
@@ -113,7 +114,7 @@ def jaccard_index(
 
             - ``'elementwise_mean'``: takes the mean (default)
             - ``'sum'``: takes the sum
-            - ``'none'``: no reduction will be applied
+            - ``'none'`` or ``None``: no reduction will be applied
 
     Return:
         IoU score: Tensor containing single value if reduction is

torchmetrics/functional/classification/kl_divergence.py

@@ -12,10 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Tuple
+from typing import Tuple
 
 import torch
 from torch import Tensor
+from typing_extensions import Literal
 
 from torchmetrics.utilities.checks import _check_same_shape
 from torchmetrics.utilities.data import METRIC_EPS
@@ -47,7 +48,7 @@ def _kld_update(p: Tensor, q: Tensor, log_prob: bool) -> Tuple[Tensor, int]:
     return measures, total
 
 
-def _kld_compute(measures: Tensor, total: Tensor, reduction: Optional[str] = "mean") -> Tensor:
+def _kld_compute(measures: Tensor, total: Tensor, reduction: Literal["mean", "sum", "none", None] = "mean") -> Tensor:
     """Computes the KL divergenece based on the type of reduction.
 
     Args:
@@ -77,7 +78,9 @@ def _kld_compute(measures: Tensor, total: Tensor, reduction: Optional[str] = "me
     return measures / total
 
 
-def kl_divergence(p: Tensor, q: Tensor, log_prob: bool = False, reduction: Optional[str] = "mean") -> Tensor:
+def kl_divergence(
+    p: Tensor, q: Tensor, log_prob: bool = False, reduction: Literal["mean", "sum", "none", None] = "mean"
+) -> Tensor:
     r"""Computes `KL divergence`_
 
     .. math::

torchmetrics/functional/image/psnr.py

@@ -15,6 +15,7 @@
 
 import torch
 from torch import Tensor, tensor
+from typing_extensions import Literal
 
 from torchmetrics.utilities import rank_zero_warn, reduce
 
@@ -24,7 +25,7 @@ def _psnr_compute(
     n_obs: Tensor,
     data_range: Tensor,
     base: float = 10.0,
-    reduction: str = "elementwise_mean",
+    reduction: Literal["elementwise_mean", "sum", "none", None] = "elementwise_mean",
 ) -> Tensor:
     """Computes peak signal-to-noise ratio.
 
@@ -39,7 +40,7 @@ def _psnr_compute(
 
             - ``'elementwise_mean'``: takes the mean (default)
             - ``'sum'``: takes the sum
-            - ``'none'``: no reduction will be applied
+            - ``'none'`` or ``None``: no reduction will be applied
 
     Example:
         >>> preds = torch.tensor([[0.0, 1.0], [2.0, 3.0]])
@@ -96,7 +97,7 @@ def peak_signal_noise_ratio(
     target: Tensor,
     data_range: Optional[float] = None,
     base: float = 10.0,
-    reduction: str = "elementwise_mean",
+    reduction: Literal["elementwise_mean", "sum", "none", None] = "elementwise_mean",
     dim: Optional[Union[int, Tuple[int, ...]]] = None,
 ) -> Tensor:
     """Computes the peak signal-to-noise ratio.
@@ -112,7 +113,7 @@ def peak_signal_noise_ratio(
 
             - ``'elementwise_mean'``: takes the mean (default)
             - ``'sum'``: takes the sum
-            - ``'none'``: no reduction will be applied
+            - ``'none'`` or None``: no reduction will be applied
 
         dim:
             Dimensions to reduce PSNR scores over provided as either an integer or a list of integers. Default is

torchmetrics/functional/image/ssim.py

@@ -51,7 +51,7 @@ def _ssim_compute(
     target: Tensor,
     kernel_size: Sequence[int] = (11, 11),
     sigma: Sequence[float] = (1.5, 1.5),
-    reduction: str = "elementwise_mean",
+    reduction: Literal["elementwise_mean", "sum", "none", None] = "elementwise_mean",
     data_range: Optional[float] = None,
     k1: float = 0.01,
     k2: float = 0.03,
@@ -68,7 +68,7 @@ def _ssim_compute(
 
             - ``'elementwise_mean'``: takes the mean (default)
             - ``'sum'``: takes the sum
-            - ``'none'``: no reduction will be applied
+            - ``'none'`` or ``None``: no reduction will be applied
 
         data_range: Range of the image. If ``None``, it is determined from the image (max - min)
         k1: Parameter of SSIM.
@@ -140,7 +140,7 @@ def structural_similarity_index_measure(
     target: Tensor,
     kernel_size: Sequence[int] = (11, 11),
     sigma: Sequence[float] = (1.5, 1.5),
-    reduction: str = "elementwise_mean",
+    reduction: Literal["elementwise_mean", "sum", "none", None] = "elementwise_mean",
     data_range: Optional[float] = None,
     k1: float = 0.01,
     k2: float = 0.03,
@@ -156,7 +156,7 @@ def structural_similarity_index_measure(
 
             - ``'elementwise_mean'``: takes the mean (default)
             - ``'sum'``: takes the sum
-            - ``'none'``: no reduction will be applied
+            - ``'none'`` or ``None``: no reduction will be applied
 
         data_range: Range of the image. If ``None``, it is determined from the image (max - min)
         k1: Parameter of SSIM.
@@ -193,7 +193,7 @@ def _get_normalized_sim_and_cs(
     target: Tensor,
     kernel_size: Sequence[int] = (11, 11),
     sigma: Sequence[float] = (1.5, 1.5),
-    reduction: str = "elementwise_mean",
+    reduction: Literal["elementwise_mean", "sum", "none", None] = "elementwise_mean",
     data_range: Optional[float] = None,
     k1: float = 0.01,
     k2: float = 0.03,
@@ -213,7 +213,7 @@ def _multiscale_ssim_compute(
     target: Tensor,
     kernel_size: Sequence[int] = (11, 11),
     sigma: Sequence[float] = (1.5, 1.5),
-    reduction: str = "elementwise_mean",
+    reduction: Literal["elementwise_mean", "sum", "none", None] = "elementwise_mean",
     data_range: Optional[float] = None,
     k1: float = 0.01,
     k2: float = 0.03,
@@ -238,7 +238,7 @@ def _multiscale_ssim_compute(
 
             - ``'elementwise_mean'``: takes the mean (default)
             - ``'sum'``: takes the sum
-            - ``'none'``: no reduction will be applied
+            - ``'none'`` or ``None``: no reduction will be applied
 
         data_range: Range of the image. If ``None``, it is determined from the image (max - min)
         k1: Parameter of structural similarity index measure.
@@ -304,7 +304,7 @@ def multiscale_structural_similarity_index_measure(
     target: Tensor,
     kernel_size: Sequence[int] = (11, 11),
     sigma: Sequence[float] = (1.5, 1.5),
-    reduction: str = "elementwise_mean",
+    reduction: Literal["elementwise_mean", "sum", "none", None] = "elementwise_mean",
     data_range: Optional[float] = None,
     k1: float = 0.01,
     k2: float = 0.03,
@@ -323,7 +323,7 @@ def multiscale_structural_similarity_index_measure(
 
             - ``'elementwise_mean'``: takes the mean (default)
             - ``'sum'``: takes the sum
-            - ``'none'``: no reduction will be applied
+            - ``'none'`` or ``None``: no reduction will be applied
 
         data_range: Range of the image. If ``None``, it is determined from the image (max - min)
         k1: Parameter of structural similarity index measure.

torchmetrics/functional/image/uqi.py

@@ -51,7 +51,7 @@ def _uqi_compute(
     target: Tensor,
     kernel_size: Sequence[int] = (11, 11),
     sigma: Sequence[float] = (1.5, 1.5),
-    reduction: Literal["elementwise_mean", "sum", "none"] = "elementwise_mean",
+    reduction: Optional[Literal["elementwise_mean", "sum", "none"]] = "elementwise_mean",
     data_range: Optional[float] = None,
     return_contrast_sensitivity: bool = False,
 ) -> Union[Tensor, Tuple[Tensor, Tensor]]:
@@ -66,7 +66,7 @@ def _uqi_compute(
 
             - ``'elementwise_mean'``: takes the mean (default)
             - ``'sum'``: takes the sum
-            - ``'none'``: no reduction will be applied
+            - ``'none'`` or ``None``: no reduction will be applied
 
         data_range: Range of the image. If ``None``, it is determined from the image (max - min)
 
@@ -128,7 +128,7 @@ def universal_image_quality_index(
     target: Tensor,
     kernel_size: Sequence[int] = (11, 11),
     sigma: Sequence[float] = (1.5, 1.5),
-    reduction: Literal["elementwise_mean", "sum", "none"] = "elementwise_mean",
+    reduction: Optional[Literal["elementwise_mean", "sum", "none"]] = "elementwise_mean",
     data_range: Optional[float] = None,
 ) -> Tensor:
     """Universal Image Quality Index.
@@ -142,7 +142,7 @@ def universal_image_quality_index(
 
             - ``'elementwise_mean'``: takes the mean (default)
             - ``'sum'``: takes the sum
-            - ``'none'``: no reduction will be applied
+            - ``'none'`` or ``None``: no reduction will be applied
 
         data_range: Range of the image. If ``None``, it is determined from the image (max - min)
 

torchmetrics/functional/pairwise/cosine.py

@@ -15,6 +15,7 @@
 
 import torch
 from torch import Tensor
+from typing_extensions import Literal
 
 from torchmetrics.functional.pairwise.helpers import _check_input, _reduce_distance_matrix
 
@@ -43,7 +44,10 @@ def _pairwise_cosine_similarity_update(
 
 
 def pairwise_cosine_similarity(
-    x: Tensor, y: Optional[Tensor] = None, reduction: Optional[str] = None, zero_diagonal: Optional[bool] = None
+    x: Tensor,
+    y: Optional[Tensor] = None,
+    reduction: Literal["mean", "sum", "none", None] = None,
+    zero_diagonal: Optional[bool] = None,
 ) -> Tensor:
     r"""
     Calculates pairwise cosine similarity:

torchmetrics/functional/pairwise/euclidean.py

@@ -14,6 +14,7 @@
 from typing import Optional
 
 from torch import Tensor
+from typing_extensions import Literal
 
 from torchmetrics.functional.pairwise.helpers import _check_input, _reduce_distance_matrix
 
@@ -38,7 +39,10 @@ def _pairwise_euclidean_distance_update(
 
 
 def pairwise_euclidean_distance(
-    x: Tensor, y: Optional[Tensor] = None, reduction: Optional[str] = None, zero_diagonal: Optional[bool] = None
+    x: Tensor,
+    y: Optional[Tensor] = None,
+    reduction: Literal["mean", "sum", "none", None] = None,
+    zero_diagonal: Optional[bool] = None,
 ) -> Tensor:
     r"""
     Calculates pairwise euclidean distances:

torchmetrics/functional/pairwise/linear.py

@@ -14,6 +14,7 @@
 from typing import Optional
 
 from torch import Tensor
+from typing_extensions import Literal
 
 from torchmetrics.functional.pairwise.helpers import _check_input, _reduce_distance_matrix
 
@@ -37,7 +38,10 @@ def _pairwise_linear_similarity_update(
 
 
 def pairwise_linear_similarity(
-    x: Tensor, y: Optional[Tensor] = None, reduction: Optional[str] = None, zero_diagonal: Optional[bool] = None
+    x: Tensor,
+    y: Optional[Tensor] = None,
+    reduction: Literal["mean", "sum", "none", None] = None,
+    zero_diagonal: Optional[bool] = None,
 ) -> Tensor:
     r"""
     Calculates pairwise linear similarity:

torchmetrics/functional/pairwise/manhattan.py

@@ -14,6 +14,7 @@
 from typing import Optional
 
 from torch import Tensor
+from typing_extensions import Literal
 
 from torchmetrics.functional.pairwise.helpers import _check_input, _reduce_distance_matrix
 
@@ -37,7 +38,10 @@ def _pairwise_manhattan_distance_update(
 
 
 def pairwise_manhattan_distance(
-    x: Tensor, y: Optional[Tensor] = None, reduction: Optional[str] = None, zero_diagonal: Optional[bool] = None
+    x: Tensor,
+    y: Optional[Tensor] = None,
+    reduction: Literal["mean", "sum", "none", None] = None,
+    zero_diagonal: Optional[bool] = None,
 ) -> Tensor:
     r"""
     Calculates pairwise manhattan distance: