Update PTST.py

PTST.py

@@ -117,6 +117,32 @@ class MultiHeadSelfAttention(nn.Module):
         out = nn.Dropout(rate=self.dropout)(out, deterministic=deterministic)
         return out
 
+class ChannelSelfAttention(nn.Module):
+    d_model: int
+    tau: float = 1.0   # softmax temperature
+    col_norm: bool = True  # match your C^- convention
+    init_scale: float = 0.0
+
+    @nn.compact
+    def __call__(self, x, deterministic: bool):
+        """
+        x: (B, T, d_model)
+        returns: (B, T, d_model)
+        """
+        B, T, D = x.shape
+        assert D == self.d_model
+
+        # Channel affinity per sample: S = x^T x / T  => (B, D, D)
+        S = jnp.einsum("btd,bte->bde", x, x) / jnp.maximum(T, 1)
+
+        # Softmax to get mixing weights (row-stochastic by default)
+        C_row = nn.softmax(S / self.tau, axis=-1)  # rows sum to 1
+
+        # If you want column-normalized C^- for right-multiply, use transpose:
+        C = jnp.swapaxes(C_row, -1, -2) if self.col_norm else C_row  # (B,D,D)
+
+        scale = self.param("scale", lambda k, s: jnp.array(self.init_scale, jnp.float32), ())
+        return x + scale * (x @ C)  # residual; starts near identity
 
 class FeedForward(nn.Module):
     d_model: int
@@ -134,7 +160,7 @@ class FeedForward(nn.Module):
         return x
 
 
-class TransformerEncoderBlock(nn.Module):
+class TransformerEncoderBlockX(nn.Module):
     d_model: int
     n_heads: int
     dropout: float = 0.0
@@ -155,6 +181,51 @@ class TransformerEncoderBlock(nn.Module):
         x = x + h2
         return x
 
+class TransformerEncoderBlock(nn.Module):
+    d_model: int
+    n_heads: int
+    dropout: float = 0.0
+    mlp_ratio: float = 4.0
+
+    # new options
+    use_cam: bool = False
+    use_ffn: bool = True
+
+    # CAM options
+    cam_tau: float = 1.0
+    cam_col_norm: bool = True
+    cam_init_scale: float = 0.0
+
+    @nn.compact
+    def __call__(self, x, deterministic: bool):
+        # --- Token self-attention (always on, per your current design) ---
+        h = nn.LayerNorm(name="attn_ln")(x)
+        h = MultiHeadSelfAttention(
+            d_model=self.d_model, n_heads=self.n_heads, dropout=self.dropout
+        )(h, deterministic=deterministic)
+        x = x + h
+
+        # --- Channel self-attention (optional) ---
+        if self.use_cam:
+            hc = nn.LayerNorm(name="cam_ln")(x)
+            hc = ChannelSelfAttention(
+                d_model=self.d_model,
+                tau=self.cam_tau,
+                col_norm=self.cam_col_norm,
+                init_scale=self.cam_init_scale,
+                name="cam",
+            )(hc, deterministic=deterministic)
+            x = x + hc  # residual around CAM
+
+        # --- FFN (optional) ---
+        if self.use_ffn:
+            h2 = nn.LayerNorm(name="ffn_ln")(x)
+            h2 = FeedForward(
+                d_model=self.d_model, mlp_ratio=self.mlp_ratio, dropout=self.dropout
+            )(h2, deterministic=deterministic)
+            x = x + h2
+
+        return x
 
 class PTSTBackbone(nn.Module):
     """