Allow batch dimensions in dot and cos similarity (#119)

Author: mikeheddes

torchhd/__init__.py

@@ -34,6 +34,7 @@
     dot_similarity,
     hamming_similarity,
     multiset,
+    multibundle,
     multibind,
     bundle_sequence,
     bind_sequence,
@@ -83,6 +84,7 @@
     "dot_similarity",
     "hamming_similarity",
     "multiset",
+    "multibundle",
     "multibind",
     "bundle_sequence",
     "bind_sequence",

torchhd/bsc.py

@@ -406,14 +406,17 @@ def permute(self, shifts: int = 1) -> "BSC":
     def dot_similarity(self, others: "BSC") -> Tensor:
         """Inner product with other hypervectors."""
         dtype = torch.get_default_dtype()
+        device = self.device
 
-        min_one = torch.tensor(-1.0, dtype=dtype)
-        plus_one = torch.tensor(1.0, dtype=dtype)
+        min_one = torch.tensor(-1.0, dtype=dtype, device=device)
+        plus_one = torch.tensor(1.0, dtype=dtype, device=device)
 
         self_as_bipolar = torch.where(self.bool(), min_one, plus_one)
         others_as_bipolar = torch.where(others.bool(), min_one, plus_one)
 
-        return F.linear(self_as_bipolar, others_as_bipolar)
+        if others.dim() >= 2:
+            others_as_bipolar = others_as_bipolar.mT
+        return torch.matmul(self_as_bipolar, others_as_bipolar)
 
     def cos_similarity(self, others: "BSC") -> Tensor:
         """Cosine similarity with other hypervectors."""

torchhd/fhrr.py

@@ -354,7 +354,9 @@ def permute(self, shifts: int = 1) -> "FHRR":
 
     def dot_similarity(self, others: "FHRR") -> Tensor:
         """Inner product with other hypervectors"""
-        return F.linear(self, others.conj()).real
+        if others.dim() >= 2:
+            others = others.mT
+        return torch.matmul(self, others.conj()).real
 
     def cos_similarity(self, others: "FHRR", *, eps=1e-08) -> Tensor:
         """Cosine similarity with other hypervectors"""

torchhd/hrr.py

@@ -339,7 +339,9 @@ def permute(self, shifts: int = 1) -> "HRR":
 
     def dot_similarity(self, others: "HRR") -> Tensor:
         """Inner product with other hypervectors"""
-        return F.linear(self, others)
+        if others.dim() >= 2:
+            others = others.mT
+        return torch.matmul(self, others)
 
     def cos_similarity(self, others: "HRR", *, eps=1e-08) -> Tensor:
         """Cosine similarity with other hypervectors"""

torchhd/map.py

@@ -320,7 +320,9 @@ def clipping(self, kappa) -> "MAP":
     def dot_similarity(self, others: "MAP") -> Tensor:
         """Inner product with other hypervectors"""
         dtype = torch.get_default_dtype()
-        return F.linear(self.to(dtype), others.to(dtype))
+        if others.dim() >= 2:
+            others = others.mT
+        return torch.matmul(self.to(dtype), others.to(dtype))
 
     def cos_similarity(self, others: "MAP", *, eps=1e-08) -> Tensor:
         """Cosine similarity with other hypervectors"""
@@ -332,10 +334,10 @@ def cos_similarity(self, others: "MAP", *, eps=1e-08) -> Tensor:
         others_dot = torch.sum(others * others, dim=-1, dtype=dtype)
         others_mag = others_dot.sqrt()
 
-        if self.dim() > 1:
+        if self.dim() >= 2:
             magnitude = self_mag.unsqueeze(-1) * others_mag.unsqueeze(0)
         else:
             magnitude = self_mag * others_mag
 
         magnitude = magnitude.clamp(min=eps)
-        return F.linear(self.to(dtype), others.to(dtype)) / magnitude
+        return self.dot_similarity(others) / magnitude