Fix cosine similarity for batch processing

Author: mikeheddes

torchhd/tensors/fhrr.py

@@ -366,8 +366,8 @@ def cosine_similarity(self, others: "FHRRTensor", *, eps=1e-08) -> Tensor:
         others_dot = torch.sum(torch.real(others * torch.conj(others)), dim=-1)
         others_mag = torch.sqrt(others_dot)
 
-        if self.dim() > 1:
-            magnitude = self_mag.unsqueeze(-1) * others_mag.unsqueeze(0)
+        if self.dim() >= 2:
+            magnitude = self_mag.unsqueeze(-1) * others_mag.unsqueeze(-2)
         else:
             magnitude = self_mag * others_mag
 

torchhd/tensors/hrr.py

@@ -353,8 +353,8 @@ def cosine_similarity(self, others: "HRRTensor", *, eps=1e-08) -> Tensor:
         others_dot = torch.sum(others * others, dim=-1)
         others_mag = torch.sqrt(others_dot)
 
-        if self.dim() > 1:
-            magnitude = self_mag.unsqueeze(-1) * others_mag.unsqueeze(0)
+        if self.dim() >= 2:
+            magnitude = self_mag.unsqueeze(-1) * others_mag.unsqueeze(-2)
         else:
             magnitude = self_mag * others_mag
 

torchhd/tensors/map.py

@@ -335,7 +335,7 @@ def cosine_similarity(self, others: "MAPTensor", *, eps=1e-08) -> Tensor:
         others_mag = torch.sqrt(others_dot)
 
         if self.dim() >= 2:
-            magnitude = self_mag.unsqueeze(-1) * others_mag.unsqueeze(0)
+            magnitude = self_mag.unsqueeze(-1) * others_mag.unsqueeze(-2)
         else:
             magnitude = self_mag * others_mag