Add None as reduction option in CosineSimilarity

CHANGELOG.md

@@ -20,6 +20,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added Symmetric Mean Absolute Percentage error (SMAPE) ([#375](https://github.com/PyTorchLightning/metrics/issues/375))
 
 
+- Added `None` as reduction option in `CosineSimilarity` metric ([#400](https://github.com/PyTorchLightning/metrics/pull/400))
+
+
 - Added Permutation Invariant Training metric (PIT) ([#294](https://github.com/PyTorchLightning/metrics/issues/294))
 
 

torchmetrics/functional/regression/cosine_similarity.py

@@ -61,7 +61,12 @@ def _cosine_similarity_compute(preds: Tensor, target: Tensor, reduction: str = '
  preds_norm = preds.norm(dim=-1)
  target_norm = target.norm(dim=-1)
  similarity = dot_product / (preds_norm * target_norm)
- reduction_mapping = {"sum": torch.sum, "mean": torch.mean, "none": lambda x: x}
+ reduction_mapping = {
+ "sum": torch.sum,
+ "mean": torch.mean,
+ "none": lambda x: x,
+ None: lambda x: x,
+ }
  return reduction_mapping[reduction](similarity)
 
 

torchmetrics/regression/cosine_similarity.py

@@ -77,10 +77,13 @@ def __init__(
  process_group=process_group,
  dist_sync_fn=dist_sync_fn
  )
+ allowed_reduction = ('sum', 'mean', 'none', None)
+ if reduction not in allowed_reduction:
+ raise ValueError(f"Expected argument `reduction` to be one of {allowed_reduction}" f" but got {reduction}")
+ self.reduction = reduction
 
  self.add_state("preds", [], dist_reduce_fx="cat")
  self.add_state("target", [], dist_reduce_fx="cat")
- self.reduction = reduction
 
  def update(self, preds: Tensor, target: Tensor) -> None: # type: ignore
  """