Add variational and spatial token autoencoders

bluemath_tk/deeplearning/layers.py

@@ -736,9 +736,12 @@ def _extend_for_multilayer(param, num_layers):
 
 class LinearSelfAttention(nn.Module):
     """
-    Softmax-free, Performer-style linear attention on the time axis.
+    Softmax-free linear attention with an ELU+1 feature map.
 
-    Provides O(B * L * D * H) scaling, good for large sequence lengths.
+    Attention contractions scale linearly with sequence length. Including
+    dense projections, total compute is approximately ``O(B * L * D**2)``;
+    accumulated key-value state across all heads is approximately
+    ``O(B * D**2 / H)``.
 
     Parameters
     ----------
@@ -750,7 +753,22 @@ class LinearSelfAttention(nn.Module):
 
     def __init__(self, d_model: int, num_heads: int = 4):
         super().__init__()
-        assert d_model % num_heads == 0
+        for name, value in (
+            ("d_model", d_model),
+            ("num_heads", num_heads),
+        ):
+            if (
+                not isinstance(value, int)
+                or isinstance(value, bool)
+                or value < 1
+            ):
+                raise ValueError(
+                    f"{name} must be a positive integer."
+                )
+        if d_model % num_heads != 0:
+            raise ValueError(
+                "d_model must be divisible by num_heads."
+            )
         self.d_model = d_model
         self.num_heads = num_heads
         self.d_head = d_model // num_heads
@@ -778,6 +796,11 @@ def forward(self, x):
             Output sequences, shape (B, L, D).
         """
         # x: (B, L, D)
+        if x.dim() != 3 or x.size(-1) != self.d_model:
+            raise ValueError(
+                "LinearSelfAttention expects input shape "
+                f"(B, L, {self.d_model}), got {tuple(x.shape)}."
+            )
         B, L, D = x.shape
         Q = self.Wq(x)
         K = self.Wk(x)