reuse _safe_matmul (#1014)

* additional `_safe_matmul`
* chlog

(cherry picked from commit 21384e5)

CHANGELOG.md

@@ -22,10 +22,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fixed multi device aggregation in `PearsonCorrCoef` ([#998](https://github.com/PyTorchLightning/metrics/pull/998))
 
 
-- Fixed MAP metric when using custom list of thresholds ([#995](https://github.com/PyTorchLightning/metrics/issues/995))
+- Fixed compatibility with future Pytorch 1.12 in `safe_matmul` ([#1011](https://github.com/PyTorchLightning/metrics/pull/1011), [#1014](https://github.com/PyTorchLightning/metrics/pull/1014))
 
 
-- Fixed compatibility with future Pytorch 1.12 in `pairwise_cosine_similarity` ([#1011](https://github.com/PyTorchLightning/metrics/pull/1011))
+- Fixed MAP metric when using custom list of thresholds ([#995](https://github.com/PyTorchLightning/metrics/issues/995))
 
 
 ## [0.8.1] - 2022-04-27

torchmetrics/functional/pairwise/cosine.py

@@ -18,16 +18,7 @@
 from typing_extensions import Literal
 
 from torchmetrics.functional.pairwise.helpers import _check_input, _reduce_distance_matrix
-
-
-def _safe_matmul(x: Tensor, y: Tensor) -> Tensor:
-    """Safe calculation of matrix multiplication.
-
-    If input is float16, will cast to float32 for computation and back again.
-    """
-    if x.dtype == torch.float16 or y.dtype == torch.float16:
-        return (x.float() @ y.T.float()).half()
-    return x @ y.T
+from torchmetrics.utilities.compute import _safe_matmul
 
 
 def _pairwise_cosine_similarity_update(

torchmetrics/functional/pairwise/linear.py

@@ -17,6 +17,7 @@
 from typing_extensions import Literal
 
 from torchmetrics.functional.pairwise.helpers import _check_input, _reduce_distance_matrix
+from torchmetrics.utilities.compute import _safe_matmul
 
 
 def _pairwise_linear_similarity_update(
@@ -31,7 +32,7 @@ def _pairwise_linear_similarity_update(
     """
     x, y, zero_diagonal = _check_input(x, y, zero_diagonal)
 
-    distance = x @ y.T
+    distance = _safe_matmul(x, y)
     if zero_diagonal:
         distance.fill_diagonal_(0)
     return distance

torchmetrics/utilities/compute.py

@@ -0,0 +1,25 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+from torch import Tensor
+
+
+def _safe_matmul(x: Tensor, y: Tensor) -> Tensor:
+    """Safe calculation of matrix multiplication.
+
+    If input is float16, will cast to float32 for computation and back again.
+    """
+    if x.dtype == torch.float16 or y.dtype == torch.float16:
+        return (x.float() @ y.T.float()).half()
+    return x @ y.T