Implement Binary Sparse Block Codes VSA model (#146)

* WIP segmented sparse vectors

* [github-action] formatting fixes

* Add documentation and rename

* [github-action] formatting fixes

* Update comments

* Rename to Sparse Block Codes

* Add tests and change segment size to block size

* [github-action] formatting fixes

* Add missing block_size argument

* [github-action] formatting fixes

* Update docs

* [github-action] formatting fixes

* Update naming to Binary SBC

* [github-action] formatting fixes

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: mikeheddes

docs/torchhd.rst

@@ -85,6 +85,7 @@ VSA Models
     MAPTensor
     HRRTensor
     FHRRTensor
+    BSBCTensor
     VTBTensor
 
 

torchhd/__init__.py

@@ -34,6 +34,7 @@
 from torchhd.tensors.map import MAPTensor
 from torchhd.tensors.hrr import HRRTensor
 from torchhd.tensors.fhrr import FHRRTensor
+from torchhd.tensors.bsbc import BSBCTensor
 from torchhd.tensors.vtb import VTBTensor
 
 from torchhd.functional import (
@@ -85,6 +86,7 @@
     "MAPTensor",
     "HRRTensor",
     "FHRRTensor",
+    "BSBCTensor",
     "VTBTensor",
     "functional",
     "embeddings",

torchhd/embeddings.py

@@ -108,6 +108,7 @@ def __init__(
         sparse: bool = False,
         device=None,
         dtype=None,
+        **kwargs,
     ) -> None:
         factory_kwargs = {"device": device, "dtype": dtype}
         # Have to call Module init explicitly in order not to use the Embedding init
@@ -116,6 +117,7 @@ def __init__(
         self.num_embeddings = num_embeddings
         self.embedding_dim = embedding_dim
         self.vsa = vsa
+        self.vsa_kwargs = kwargs
 
         if padding_idx is not None:
             if padding_idx > 0:
@@ -135,7 +137,7 @@ def __init__(
         self.sparse = sparse
 
         embeddings = functional.empty(
-            num_embeddings, embedding_dim, self.vsa, **factory_kwargs
+            num_embeddings, embedding_dim, self.vsa, **factory_kwargs, **self.vsa_kwargs
         )
         # Have to provide requires grad at the creation of the parameters to
         # prevent errors when instantiating a non-float embedding
@@ -148,7 +150,11 @@ def reset_parameters(self) -> None:
 
         with torch.no_grad():
             embeddings = functional.empty(
-                self.num_embeddings, self.embedding_dim, self.vsa, **factory_kwargs
+                self.num_embeddings,
+                self.embedding_dim,
+                self.vsa,
+                **factory_kwargs,
+                **self.vsa_kwargs,
             )
             self.weight.copy_(embeddings)
 
@@ -214,6 +220,7 @@ def __init__(
         sparse: bool = False,
         device=None,
         dtype=None,
+        **kwargs,
     ) -> None:
         factory_kwargs = {"device": device, "dtype": dtype}
         # Have to call Module init explicitly in order not to use the Embedding init
@@ -222,6 +229,7 @@ def __init__(
         self.num_embeddings = num_embeddings
         self.embedding_dim = embedding_dim
         self.vsa = vsa
+        self.vsa_kwargs = kwargs
 
         if padding_idx is not None:
             if padding_idx > 0:
@@ -241,7 +249,7 @@ def __init__(
         self.sparse = sparse
 
         embeddings = functional.identity(
-            num_embeddings, embedding_dim, self.vsa, **factory_kwargs
+            num_embeddings, embedding_dim, self.vsa, **factory_kwargs, **self.vsa_kwargs
         )
         # Have to provide requires grad at the creation of the parameters to
         # prevent errors when instantiating a non-float embedding
@@ -254,7 +262,11 @@ def reset_parameters(self) -> None:
 
         with torch.no_grad():
             embeddings = functional.identity(
-                self.num_embeddings, self.embedding_dim, self.vsa, **factory_kwargs
+                self.num_embeddings,
+                self.embedding_dim,
+                self.vsa,
+                **factory_kwargs,
+                **self.vsa_kwargs,
             )
             self.weight.copy_(embeddings)
 
@@ -266,7 +278,7 @@ def _fill_padding_idx_with_empty(self) -> None:
         if self.padding_idx is not None:
             with torch.no_grad():
                 empty = functional.empty(
-                    1, self.embedding_dim, self.vsa, **factory_kwargs
+                    1, self.embedding_dim, self.vsa, **factory_kwargs, **self.vsa_kwargs
                 )
                 self.weight[self.padding_idx].copy_(empty.squeeze(0))
 
@@ -332,6 +344,7 @@ def __init__(
         sparse: bool = False,
         device=None,
         dtype=None,
+        **kwargs,
     ) -> None:
         factory_kwargs = {"device": device, "dtype": dtype}
         # Have to call Module init explicitly in order not to use the Embedding init
@@ -340,6 +353,7 @@ def __init__(
         self.num_embeddings = num_embeddings
         self.embedding_dim = embedding_dim
         self.vsa = vsa
+        self.vsa_kwargs = kwargs
 
         if padding_idx is not None:
             if padding_idx > 0:
@@ -359,7 +373,7 @@ def __init__(
         self.sparse = sparse
 
         embeddings = functional.random(
-            num_embeddings, embedding_dim, self.vsa, **factory_kwargs
+            num_embeddings, embedding_dim, self.vsa, **factory_kwargs, **self.vsa_kwargs
         )
         # Have to provide requires grad at the creation of the parameters to
         # prevent errors when instantiating a non-float embedding
@@ -372,7 +386,11 @@ def reset_parameters(self) -> None:
 
         with torch.no_grad():
             embeddings = functional.random(
-                self.num_embeddings, self.embedding_dim, self.vsa, **factory_kwargs
+                self.num_embeddings,
+                self.embedding_dim,
+                self.vsa,
+                **factory_kwargs,
+                **self.vsa_kwargs,
             )
             self.weight.copy_(embeddings)
 
@@ -384,7 +402,7 @@ def _fill_padding_idx_with_empty(self) -> None:
         if self.padding_idx is not None:
             with torch.no_grad():
                 empty = functional.empty(
-                    1, self.embedding_dim, self.vsa, **factory_kwargs
+                    1, self.embedding_dim, self.vsa, **factory_kwargs, **self.vsa_kwargs
                 )
                 self.weight[self.padding_idx].copy_(empty.squeeze(0))
 
@@ -469,6 +487,7 @@ def __init__(
         sparse: bool = False,
         device=None,
         dtype=None,
+        **kwargs,
     ) -> None:
         factory_kwargs = {"device": device, "dtype": dtype}
         # Have to call Module init explicitly in order not to use the Embedding init
@@ -477,6 +496,7 @@ def __init__(
         self.num_embeddings = num_embeddings
         self.embedding_dim = embedding_dim
         self.vsa = vsa
+        self.vsa_kwargs = kwargs
         self.low = low
         self.high = high
         self.randomness = randomness
@@ -493,6 +513,7 @@ def __init__(
             self.vsa,
             randomness=randomness,
             **factory_kwargs,
+            **self.vsa_kwargs,
         )
         # Have to provide requires grad at the creation of the parameters to
         # prevent errors when instantiating a non-float embedding
@@ -508,6 +529,7 @@ def reset_parameters(self) -> None:
                 self.vsa,
                 randomness=self.randomness,
                 **factory_kwargs,
+                **self.vsa_kwargs,
             )
             self.weight.copy_(embeddings)
 
@@ -592,6 +614,7 @@ def __init__(
         sparse: bool = False,
         device=None,
         dtype=None,
+        **kwargs,
     ) -> None:
         factory_kwargs = {"device": device, "dtype": dtype}
         # Have to call Module init explicitly in order not to use the Embedding init
@@ -600,6 +623,7 @@ def __init__(
         self.num_embeddings = num_embeddings
         self.embedding_dim = embedding_dim
         self.vsa = vsa
+        self.vsa_kwargs = kwargs
         self.low = low
         self.high = high
 
@@ -610,7 +634,7 @@ def __init__(
         self.sparse = sparse
 
         embeddings = functional.thermometer(
-            num_embeddings, embedding_dim, self.vsa, **factory_kwargs
+            num_embeddings, embedding_dim, self.vsa, **factory_kwargs, **self.vsa_kwargs
         )
         # Have to provide requires grad at the creation of the parameters to
         # prevent errors when instantiating a non-float embedding
@@ -621,7 +645,11 @@ def reset_parameters(self) -> None:
 
         with torch.no_grad():
             embeddings = functional.thermometer(
-                self.num_embeddings, self.embedding_dim, self.vsa, **factory_kwargs
+                self.num_embeddings,
+                self.embedding_dim,
+                self.vsa,
+                **factory_kwargs,
+                **self.vsa_kwargs,
             )
             self.weight.copy_(embeddings)
 
@@ -704,6 +732,7 @@ def __init__(
         sparse: bool = False,
         device=None,
         dtype=None,
+        **kwargs,
     ) -> None:
         factory_kwargs = {"device": device, "dtype": dtype}
         # Have to call Module init explicitly in order not to use the Embedding init
@@ -712,6 +741,7 @@ def __init__(
         self.num_embeddings = num_embeddings
         self.embedding_dim = embedding_dim
         self.vsa = vsa
+        self.vsa_kwargs = kwargs
         self.phase = phase
         self.period = period
         self.randomness = randomness
@@ -728,6 +758,7 @@ def __init__(
             self.vsa,
             randomness=randomness,
             **factory_kwargs,
+            **self.vsa_kwargs,
         )
         # Have to provide requires grad at the creation of the parameters to
         # prevent errors when instantiating a non-float embedding
@@ -743,6 +774,7 @@ def reset_parameters(self) -> None:
                 self.vsa,
                 randomness=self.randomness,
                 **factory_kwargs,
+                **self.vsa_kwargs,
             )
             self.weight.copy_(embeddings)
 
@@ -945,6 +977,7 @@ def __init__(
         device=None,
         dtype=None,
         requires_grad: bool = False,
+        **kwargs,
     ):
         factory_kwargs = {
             "device": device,
@@ -954,10 +987,16 @@ def __init__(
         super(Density, self).__init__()
 
         # A set of random vectors used as unique IDs for features of the dataset.
-        self.key = Random(in_features, out_features, vsa, **factory_kwargs)
+        self.key = Random(in_features, out_features, vsa, **factory_kwargs, **kwargs)
         # Thermometer encoding used for transforming input data.
         self.density_encoding = Thermometer(
-            out_features + 1, out_features, vsa, low=low, high=high, **factory_kwargs
+            out_features + 1,
+            out_features,
+            vsa,
+            low=low,
+            high=high,
+            **factory_kwargs,
+            **kwargs,
         )
 
     def reset_parameters(self) -> None:

torchhd/functional.py

@@ -32,6 +32,7 @@
 from torchhd.tensors.map import MAPTensor
 from torchhd.tensors.hrr import HRRTensor
 from torchhd.tensors.fhrr import FHRRTensor
+from torchhd.tensors.bsbc import BSBCTensor
 from torchhd.tensors.vtb import VTBTensor
 from torchhd.types import VSAOptions
 
@@ -83,6 +84,8 @@ def get_vsa_tensor_class(vsa: VSAOptions) -> Type[VSATensor]:
         return HRRTensor
     elif vsa == "FHRR":
         return FHRRTensor
+    elif vsa == "BSBC":
+        return BSBCTensor
     elif vsa == "VTB":
         return VTBTensor
 
@@ -351,7 +354,10 @@ def level(
         dimensions,
         dtype=span_hv.dtype,
         device=span_hv.device,
-    )
+    ).as_subclass(vsa_tensor)
+
+    if vsa == "BSBC":
+        hv.block_size = span_hv.block_size
 
     for i in range(num_vectors):
         span_idx = int(i // levels_per_span)
@@ -372,7 +378,7 @@ def level(
             hv[i] = torch.where(threshold_v[span_idx] < t, span_start_hv, span_end_hv)
 
     hv.requires_grad = requires_grad
-    return hv.as_subclass(vsa_tensor)
+    return hv
 
 
 def thermometer(
@@ -461,7 +467,7 @@ def thermometer(
             device=rand_hv.device,
         )
     else:
-        raise ValueError(f"{vsa_tensor} HD/VSA model is not defined.")
+        raise ValueError(f"{vsa_tensor} HD/VSA model is not (yet) supported.")
 
     # Create hypervectors using the obtained step
     for i in range(1, num_vectors):
@@ -575,7 +581,10 @@ def circular(
         dimensions,
         dtype=span_hv.dtype,
         device=span_hv.device,
-    )
+    ).as_subclass(vsa_tensor)
+
+    if vsa == "BSBC":
+        hv.block_size = span_hv.block_size
 
     mutation_history = deque()
 
@@ -618,7 +627,7 @@ def circular(
             hv[i // 2] = mutation_hv
 
     hv.requires_grad = requires_grad
-    return hv.as_subclass(vsa_tensor)
+    return hv
 
 
 def bind(input: VSATensor, other: VSATensor) -> VSATensor: