v0.4: Use mambapy.pscan (proven, grad-safe parallel scan) — no more torch.associative_scan

liquidflow/model.py

@@ -1,12 +1,8 @@
 """
-LiquidFlow: A Novel Liquid-SSM Flow Matching Image Generator
-v0.3.0 — PARALLEL SSM scan via torch.associative_scan (O(log L) not O(L))
-
-Key changes from v0.2:
-- SSM uses torch.associative_scan for O(log L) parallel scan (no Python for-loop)
-- Fallback: Blelloch tree-scan in pure PyTorch for older PyTorch versions
-- patch_size=8 for 128×128 → L=256 tokens (not 1024)
-- patch_size=4 for 32/64 → fine at small sizes
+LiquidFlow v0.4 — Parallel SSM scan via mambapy.pscan (proven, grad-safe)
+
+Uses the battle-tested pscan from alxndrTL/mamba.py — O(log L) parallel,
+full autograd support, no custom CUDA kernels, works on any GPU.
 """
 
 import math
@@ -15,210 +11,35 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch.utils.checkpoint import checkpoint
 
-# ---- Parallel Scan Infrastructure ----
-
-HAS_NATIVE_SCAN = False
+# ---- Parallel Scan: mambapy (primary) or sequential fallback ----
 try:
-    from torch._higher_order_ops.associative_scan import associative_scan as _native_scan
-    HAS_NATIVE_SCAN = True
+    from mambapy.pscan import pscan as _pscan
+    HAS_PSCAN = True
 except ImportError:
-    pass
-
-
-def _ssm_combine(left, right):
-    """Associative operator for SSM: (a,b) ⊕ (a',b') = (a'*a, a'*b + b')"""
-    return (left[0] * right[0], right[0] * left[1] + right[1])
-
-
-def parallel_scan_native(A, X, dim=1):
-    """Use PyTorch built-in associative_scan (≥2.4). O(log L) parallel depth."""
-    mode = 'pointwise' if A.is_cuda else 'generic'
-    _, h_all = _native_scan(_ssm_combine, (A, X), dim=dim, combine_mode=mode)
-    return h_all
+    HAS_PSCAN = False
 
-
-def parallel_scan_blelloch(A, X):
-    """
-    Blelloch tree-scan fallback for older PyTorch.
-    Pure tensor ops, O(L log L) work, O(log L) depth.
-    A, X: (B, L, D) — L must be power of 2.
-    Returns: H (B, L, D) — all prefix scan results.
-    """
-    B, L, D = A.shape
-    assert L & (L - 1) == 0, f"L must be power of 2, got {L}"
-
-    Aa = A.clone()
-    Xa = X.clone()
-    num_steps = int(math.log2(L))
-
-    # Up-sweep (reduce): merge pairs → quads → ...
-    for k in range(num_steps):
-        s = 2 ** (k + 1)
-        half = s // 2
-        # right = op(left, right) for all pairs in parallel
-        Xa[:, s - 1::s] = Aa[:, s - 1::s] * Xa[:, half - 1::s] + Xa[:, s - 1::s]
-        Aa[:, s - 1::s] = Aa[:, s - 1::s] * Aa[:, half - 1::s]
-
-    # Clear last element (it has the total reduction, not needed for scan)
-    Xa[:, -1] = 0
-    Aa[:, -1] = 0
-
-    # Down-sweep: distribute prefix sums back
-    for k in range(num_steps - 1, -1, -1):
-        s = 2 ** (k + 1)
-        half = s // 2
-        # Save left child
-        tmp_a = Aa[:, half - 1::s].clone()
-        tmp_x = Xa[:, half - 1::s].clone()
-        # Left child ← parent
-        Aa[:, half - 1::s] = Aa[:, s - 1::s]
-        Xa[:, half - 1::s] = Xa[:, s - 1::s]
-        # Right child ← op(parent, saved left)
-        Xa[:, s - 1::s] = Aa[:, s - 1::s] * tmp_x + Xa[:, s - 1::s]  # WRONG — needs old right
-        Aa[:, s - 1::s] = Aa[:, s - 1::s] * tmp_a
-
-    # The Blelloch scan gives exclusive prefix sums. Convert to inclusive:
-    # h_t = A_t * prefix_{t-1} + X_t
-    # For inclusive: shift and apply one more step
-    # Actually, let's use the simpler Hillis-Steele approach which gives inclusive directly:
-    pass  # Blelloch is tricky to get right — use Hillis-Steele instead
-
-
-def parallel_scan_hillis_steele(A, X):
+def parallel_scan(A, X):
     """
-    Hillis-Steele inclusive parallel scan. Simpler than Blelloch.
-    O(L log L) work, O(log L) depth. All tensor operations.
-    A, X: (B, L, D). Returns H: (B, L, D) = all hidden states.
+    Parallel prefix scan for SSM: h_t = A_t * h_{t-1} + X_t
+    A, X: (B, L, ED, N). Returns H: (B, L, ED, N).
+    Uses mambapy.pscan (O(log L)), falls back to sequential if unavailable.
     """
-    B, L, D = A.shape
-
-    # Pad to power of 2 if needed
-    orig_L = L
-    next_pow2 = 1 << (L - 1).bit_length()
-    if next_pow2 != L:
-        pad = next_pow2 - L
-        A = F.pad(A, (0, 0, 0, pad), value=1.0)  # pad A with 1 (identity for mult)
-        X = F.pad(X, (0, 0, 0, pad), value=0.0)  # pad X with 0 (identity for add)
-        L = next_pow2
-
-    h = X.clone()  # (B, L, D)
-    a = A.clone()
-
-    num_steps = int(math.log2(L))
-    for d in range(num_steps):
-        stride = 2 ** d
-        # h[i] = a_shifted[i] * h[i-stride] + h[i]  (for i >= stride)
-        h_shifted = F.pad(h[:, :-stride], (0, 0, stride, 0))     # shift right by stride
-        a_shifted = F.pad(a[:, :-stride], (0, 0, stride, 0), value=1.0)
-
-        h = a_shifted * h_shifted + h  # this is wrong for multi-step...
-
-    # Actually Hillis-Steele doesn't directly work for (a,b) pairs.
-    # Let me implement the correct parallel prefix approach.
-    return h[:, :orig_L]
-
-
-def parallel_scan_correct(A, X):
-    """
-    Work-efficient parallel prefix scan for SSM recurrence.
-    h_t = A_t * h_{t-1} + X_t
-
-    Uses up-sweep + down-sweep on the (A, X) pair.
-    A, X: (B, L, D). Returns H: (B, L, D).
-    """
-    B, L, D = A.shape
-
-    # Pad L to power of 2
-    orig_L = L
-    next_pow2 = 1 << (L - 1).bit_length()
-    if next_pow2 != L:
-        pad = next_pow2 - L
-        A = F.pad(A, (0, 0, 0, pad), value=1.0)
-        X = F.pad(X, (0, 0, 0, pad), value=0.0)
-        L = next_pow2
-
-    # Work on clones
-    a = A.clone()
-    x = X.clone()
-
-    # Store intermediate values for down-sweep
-    a_levels = []
-    x_levels = []
-
-    # UP-SWEEP: reduce pairs
-    num_levels = int(math.log2(L))
-    for level in range(num_levels):
-        # Current length
-        cur_len = a.shape[1]
-        a_even = a[:, 0::2]  # left children
-        a_odd = a[:, 1::2]   # right children
-        x_even = x[:, 0::2]
-        x_odd = x[:, 1::2]
-
-        # Save for down-sweep
-        a_levels.append((a_even.clone(), a_odd.clone()))
-        x_levels.append((x_even.clone(), x_odd.clone()))
-
-        # Merge: right = right ⊕ left  →  (a_r*a_l, a_r*x_l + x_r)
-        a = a_odd * a_even
-        x = a_odd * x_even + x_odd
-
-    # After up-sweep, a and x have length 1 containing the full reduction.
-    # We need the inclusive prefix scan, not just the total.
-
-    # DOWN-SWEEP: propagate prefix sums
-    # Start with identity prefix (for position before the first element)
-    prefix_a = torch.ones(B, 1, D, device=A.device, dtype=A.dtype)
-    prefix_x = torch.zeros(B, 1, D, device=A.device, dtype=A.dtype)
-
-    for level in range(num_levels - 1, -1, -1):
-        a_even, a_odd = a_levels[level]
-        x_even, x_odd = x_levels[level]
-
-        # For each pair (even, odd) with prefix:
-        # Result for even = prefix ⊕ even
-        # Result for odd  = (prefix ⊕ even) ⊕ odd
-
-        # prefix ⊕ even: (prefix_a * a_even, prefix_a * x_even + prefix_x)
-        new_a_even = prefix_a * a_even
-        new_x_even = prefix_a * x_even + prefix_x
-
-        # (prefix ⊕ even) ⊕ odd: (new_a_even * a_odd, a_odd * new_x_even + x_odd)
-        # Wait, the operator order matters. SSM recurrence: h_t = A_t * h_{t-1} + X_t
-        # So element t is (A_t, X_t), and the scan computes h_t = result_x of prefix up to t.
-        # The operator is: (a_l, x_l) ⊕ (a_r, x_r) = (a_r * a_l, a_r * x_l + x_r)
-        new_a_odd = a_odd * new_a_even
-        new_x_odd = a_odd * new_x_even + x_odd
-
-        # Interleave back: [even_0, odd_0, even_1, odd_1, ...]
-        out_a = torch.stack([new_a_even, new_a_odd], dim=2).view(B, -1, D)
-        out_x = torch.stack([new_x_even, new_x_odd], dim=2).view(B, -1, D)
-
-        prefix_a = out_a
-        prefix_x = out_x
-
-    return prefix_x[:, :orig_L]
-
-
-def parallel_ssm_scan(A, X):
-    """
-    Top-level SSM parallel scan dispatcher.
-    A: (B, L, D) — discretized diagonal A (decay) per timestep
-    X: (B, L, D) — B_bar * u (input contribution) per timestep
-    Returns: H (B, L, D) — all hidden states h_1..h_L
-    """
-    if HAS_NATIVE_SCAN:
-        return parallel_scan_native(A, X, dim=1)
+    if HAS_PSCAN:
+        # pscan modifies X in-place, so clone
+        return _pscan(A, X.clone())
     else:
-        return parallel_scan_correct(A, X)
+        # Sequential fallback
+        B, L, ED, N = A.shape
+        h = torch.zeros(B, ED, N, device=A.device, dtype=A.dtype)
+        ys = []
+        for i in range(L):
+            h = A[:, i] * h + X[:, i]
+            ys.append(h)
+        return torch.stack(ys, dim=1)
 
 
-# ============================================================
-# 1. LIQUID TIME-CONSTANT CELL
-# ============================================================
-
 class LiquidCfCCell(nn.Module):
-    """CfC: gate=σ(-f_τ), out = gate*h + (1-gate)*f_x. Bounded by sigmoid."""
+    """CfC cell: sigmoid gating guarantees bounded dynamics."""
     def __init__(self, input_dim, hidden_dim):
         super().__init__()
         self.backbone = nn.Linear(input_dim, hidden_dim)
@@ -231,16 +52,8 @@ class LiquidCfCCell(nn.Module):
         return gate * h + (1.0 - gate) * f_x
 
 
-# ============================================================
-# 2. SELECTIVE SSM — PARALLEL SCAN
-# ============================================================
-
 class SelectiveSSM(nn.Module):
-    """
-    Selective SSM with PARALLEL scan. No Python for-loops over L.
-    Uses torch.associative_scan on GPU, tree-scan fallback on CPU.
-    Training speed: O(L log L) parallel vs O(L) sequential.
-    """
+    """Selective SSM with parallel scan via mambapy.pscan."""
     def __init__(self, d_model, d_state=16, d_conv=4, expand=2):
         super().__init__()
         self.d_model = d_model
@@ -267,83 +80,48 @@ class SelectiveSSM(nn.Module):
         B, L, _ = x.shape
         xz = self.in_proj(x)
         x_inner, z = xz.chunk(2, dim=-1)
-
         x_conv = F.silu(self.conv1d(x_inner.transpose(1, 2))[:, :, :L].transpose(1, 2))
 
         x_ssm = self.x_proj(x_conv)
-        B_sel = x_ssm[:, :, :self.d_state]                    # (B, L, N)
-        C_sel = x_ssm[:, :, self.d_state:2*self.d_state]      # (B, L, N)
-        dt = F.softplus(self.dt_proj(x_ssm[:, :, -1:]))       # (B, L, d_inner)
+        B_sel = x_ssm[:, :, :self.d_state]
+        C_sel = x_ssm[:, :, self.d_state:2*self.d_state]
+        dt = F.softplus(self.dt_proj(x_ssm[:, :, -1:]))
 
         A = -torch.exp(self.A_log)  # (d_inner, N)
 
-        y = self._parallel_ssm(x_conv, dt, A, B_sel, C_sel)
-
-        y = y + x_conv * self.D.unsqueeze(0).unsqueeze(0)
-        return self.out_proj(y * F.silu(z))
-
-    def _parallel_ssm(self, x, dt, A, B, C):
-        """
-        Parallel selective scan. No Python for-loop.
-        x:  (B, L, d_inner)
-        dt: (B, L, d_inner)
-        A:  (d_inner, N) — negative
-        B:  (B, L, N)
-        C:  (B, L, N)
-        Returns: y (B, L, d_inner)
-        """
-        Bs, L, d_inner = x.shape
-        N = A.shape[1]
-
-        # Discretize: A_bar = exp(dt * A)  — per (batch, pos, channel, state)
-        # dt: (B, L, d_inner) → (B, L, d_inner, 1)
-        # A:  (d_inner, N)   → (1, 1, d_inner, N)
+        # Discretize
         A_bar = torch.exp(dt.unsqueeze(-1) * A.unsqueeze(0).unsqueeze(0))  # (B, L, d_inner, N)
+        BX = dt.unsqueeze(-1) * B_sel.unsqueeze(2) * x_conv.unsqueeze(-1) # (B, L, d_inner, N)
 
-        # B_bar * x: dt * B * x  → (B, L, d_inner, N)
-        BX = dt.unsqueeze(-1) * B.unsqueeze(2) * x.unsqueeze(-1)  # (B, L, d_inner, N)
-
-        # Flatten (d_inner, N) → D for the scan
-        D = d_inner * N
-        A_flat = A_bar.reshape(Bs, L, D)  # (B, L, D)
-        BX_flat = BX.reshape(Bs, L, D)    # (B, L, D)
-
-        # PARALLEL SCAN: h_t = A_t * h_{t-1} + BX_t
-        h_flat = parallel_ssm_scan(A_flat, BX_flat)  # (B, L, D)
+        # Pad L to power of 2 for pscan
+        orig_L = L
+        next_pow2 = 1 << (L - 1).bit_length()
+        if next_pow2 != L:
+            pad = next_pow2 - L
+            A_bar = F.pad(A_bar, (0,0, 0,0, 0,pad), value=1.0)
+            BX = F.pad(BX, (0,0, 0,0, 0,pad), value=0.0)
 
-        # Unflatten and apply C
-        h = h_flat.reshape(Bs, L, d_inner, N)  # (B, L, d_inner, N)
+        # PARALLEL SCAN — O(log L), full grad support
+        h_all = parallel_scan(A_bar, BX)  # (B, L_padded, d_inner, N)
+        h_all = h_all[:, :orig_L]
 
-        # y_t = sum_n(C_t_n * h_t_n)  →  (B, L, d_inner)
-        y = (h * C.unsqueeze(2)).sum(-1)
-
-        return y
+        # Output: y_t = C_t · h_t
+        y = (h_all * C_sel.unsqueeze(2)).sum(-1)  # (B, L, d_inner)
 
+        y = y + x_conv * self.D.unsqueeze(0).unsqueeze(0)
+        return self.out_proj(y * F.silu(z))
 
-# ============================================================
-# 3. ZIGZAG SCAN PATTERNS
-# ============================================================
 
 def create_scan_patterns(H, W):
-    total = H * W
-    idx = torch.arange(total)
-    grid = idx.view(H, W)
-    patterns = [
-        idx.clone(),
-        idx.flip(0),
-        grid.t().contiguous().view(-1),
-        torch.cat([grid[i].flip(0) if i % 2 else grid[i] for i in range(H)]),
-    ]
+    total = H * W; idx = torch.arange(total); grid = idx.view(H, W)
+    patterns = [idx.clone(), idx.flip(0), grid.t().contiguous().view(-1),
+                torch.cat([grid[i].flip(0) if i % 2 else grid[i] for i in range(H)])]
     inv = []
     for p in patterns:
         i = torch.zeros_like(p); i[p] = torch.arange(total); inv.append(i)
     return patterns, inv
 
 
-# ============================================================
-# 4. LIQUID-SSM BLOCK
-# ============================================================
-
 class LiquidSSMBlock(nn.Module):
     def __init__(self, d_model, d_state=16, d_conv=4, expand=2, dropout=0.0):
         super().__init__()
@@ -356,35 +134,27 @@ class LiquidSSMBlock(nn.Module):
             nn.Linear(d_model, d_model * 4), nn.GELU(), nn.Dropout(dropout),
             nn.Linear(d_model * 4, d_model), nn.Dropout(dropout))
         self.mix_alpha = nn.Parameter(torch.tensor(0.5))
-
     def _ssm_fwd(self, x): return self.ssm(self.norm1(x))
     def _liq_fwd(self, x): return self.liquid(self.norm2(x))
-
     def forward(self, x, scan_idx=None, unscan_idx=None):
         xs = x[:, scan_idx] if scan_idx is not None else x
         if self.training and x.requires_grad:
             so = checkpoint(self._ssm_fwd, xs, use_reentrant=False)
             lo = checkpoint(self._liq_fwd, x, use_reentrant=False)
         else:
-            so = self._ssm_fwd(xs)
-            lo = self._liq_fwd(x)
+            so = self._ssm_fwd(xs); lo = self._liq_fwd(x)
         if unscan_idx is not None: so = so[:, unscan_idx]
         a = torch.sigmoid(self.mix_alpha)
         x = x + a * so + (1 - a) * lo
         return x + self.ff(self.norm3(x))
 
 
-# ============================================================
-# 5. EMBEDDINGS
-# ============================================================
-
 class SinusoidalPosEmb(nn.Module):
     def __init__(self, dim): super().__init__(); self.dim = dim
     def forward(self, t):
         h = self.dim // 2; e = math.log(10000)/(h-1)
         e = torch.exp(torch.arange(h, device=t.device)*-e)
-        e = t.unsqueeze(-1)*e.unsqueeze(0)
-        return torch.cat([e.sin(), e.cos()], -1)
+        return torch.cat([(t.unsqueeze(-1)*e.unsqueeze(0)).sin(), (t.unsqueeze(-1)*e.unsqueeze(0)).cos()], -1)
 
 class AdaptiveLayerNorm(nn.Module):
     def __init__(self, d, c):
@@ -395,10 +165,6 @@ class AdaptiveLayerNorm(nn.Module):
         return self.norm(x) * (1+s.unsqueeze(1)) + b.unsqueeze(1)
 
 
-# ============================================================
-# 6. LIQUIDFLOW VELOCITY NETWORK
-# ============================================================
-
 class LiquidFlowNet(nn.Module):
     def __init__(self, img_size=128, patch_size=8, in_channels=3, d_model=256,
                  depth=8, d_state=16, d_conv=4, expand=2, dropout=0.0, num_classes=0):
@@ -406,11 +172,10 @@ class LiquidFlowNet(nn.Module):
         self.img_size = img_size; self.patch_size = patch_size
         self.in_channels = in_channels; self.d_model = d_model
         self.depth = depth; self.num_classes = num_classes
-        self.nph = img_size // patch_size; self.npw = img_size // patch_size
-        self.num_patches = self.nph * self.npw
+        self.num_patches_h = img_size // patch_size
+        self.num_patches_w = img_size // patch_size
+        self.num_patches = self.num_patches_h * self.num_patches_w
         self.patch_dim = in_channels * patch_size * patch_size
-        # Alias for backward compat
-        self.num_patches_h = self.nph; self.num_patches_w = self.npw
 
         self.patch_embed = nn.Sequential(nn.Linear(self.patch_dim, d_model), nn.LayerNorm(d_model))
         self.pos_embed = nn.Parameter(torch.randn(1, self.num_patches, d_model) * 0.02)
@@ -420,15 +185,13 @@ class LiquidFlowNet(nn.Module):
         self.blocks = nn.ModuleList([LiquidSSMBlock(d_model, d_state, d_conv, expand, dropout) for _ in range(depth)])
         self.adaln = nn.ModuleList([AdaptiveLayerNorm(d_model, d_model) for _ in range(depth)])
         self.skips = nn.ModuleList([nn.Linear(d_model*2, d_model) for _ in range(depth//2)])
-
         self.final_norm = nn.LayerNorm(d_model)
         self.final_proj = nn.Linear(d_model, self.patch_dim)
 
-        pats, ipats = create_scan_patterns(self.nph, self.npw)
+        pats, ipats = create_scan_patterns(self.num_patches_h, self.num_patches_w)
         for i,(p,ip) in enumerate(zip(pats, ipats)):
             self.register_buffer(f'scan_{i}', p); self.register_buffer(f'unscan_{i}', ip)
         self.n_scans = len(pats)
-
         self.pre_conv = nn.Conv2d(d_model, d_model, 3, padding=1, groups=d_model)
         self.post_conv = nn.Conv2d(d_model, d_model, 3, padding=1, groups=d_model)
         self._init_weights()
@@ -445,32 +208,28 @@ class LiquidFlowNet(nn.Module):
 
     def patchify(self, x):
         B,C,H,W = x.shape; p = self.patch_size
-        return x.unfold(2,p,p).unfold(3,p,p).contiguous().view(B,C,self.nph,self.npw,p*p).permute(0,2,3,1,4).contiguous().view(B,self.num_patches,self.patch_dim)
+        return x.unfold(2,p,p).unfold(3,p,p).contiguous().view(B,C,self.num_patches_h,self.num_patches_w,p*p).permute(0,2,3,1,4).contiguous().view(B,self.num_patches,self.patch_dim)
 
     def unpatchify(self, x):
         B=x.shape[0]; p=self.patch_size
-        return x.view(B,self.nph,self.npw,self.in_channels,p,p).permute(0,3,1,4,2,5).contiguous().view(B,self.in_channels,self.nph*p,self.npw*p)
+        return x.view(B,self.num_patches_h,self.num_patches_w,self.in_channels,p,p).permute(0,3,1,4,2,5).contiguous().view(B,self.in_channels,self.num_patches_h*p,self.num_patches_w*p)
 
     def forward(self, x, t, class_label=None):
         B = x.shape[0]
         tok = self.patch_embed(self.patchify(x)) + self.pos_embed
-        h = tok.view(B,self.nph,self.npw,self.d_model).permute(0,3,1,2)
+        h = tok.view(B,self.num_patches_h,self.num_patches_w,self.d_model).permute(0,3,1,2)
         tok = self.pre_conv(h).permute(0,2,3,1).contiguous().view(B,self.num_patches,self.d_model)
-
         te = self.time_embed(t)
         if self.class_embed is not None and class_label is not None: te = te + self.class_embed(class_label)
-
         sk = []
         for i,(blk,aln) in enumerate(zip(self.blocks, self.adaln)):
-            tok = aln(tok, te)
-            si = i % self.n_scans
+            tok = aln(tok, te); si = i % self.n_scans
             if i < self.depth//2: sk.append(tok)
             tok = blk(tok, getattr(self,f'scan_{si}'), getattr(self,f'unscan_{si}'))
             if i >= self.depth//2:
                 j = self.depth-1-i
                 if j < len(sk): tok = self.skips[j](torch.cat([tok, sk[j]], -1))
-
-        h = tok.view(B,self.nph,self.npw,self.d_model).permute(0,3,1,2)
+        h = tok.view(B,self.num_patches_h,self.num_patches_w,self.d_model).permute(0,3,1,2)
         tok = self.post_conv(h).permute(0,2,3,1).contiguous().view(B,self.num_patches,self.d_model)
         return self.unpatchify(self.final_proj(self.final_norm(tok)))
 
@@ -478,24 +237,16 @@ class LiquidFlowNet(nn.Module):
         return sum(p.numel() for p in self.parameters() if p.requires_grad)
 
 
-# ============================================================
-# 7. MODEL CONFIGURATIONS
-# ============================================================
-
 def liquidflow_tiny(img_size=128, num_classes=0):
-    """~4M params — Colab free tier"""
     ps = 4 if img_size <= 64 else 8
     return LiquidFlowNet(img_size=img_size, patch_size=ps, d_model=192, depth=6, d_state=8, expand=2, num_classes=num_classes)
 
 def liquidflow_small(img_size=128, num_classes=0):
-    """~10M params — production 128×128"""
     ps = 4 if img_size <= 64 else 8
     return LiquidFlowNet(img_size=img_size, patch_size=ps, d_model=256, depth=8, d_state=16, expand=2, num_classes=num_classes)
 
 def liquidflow_base(img_size=256, num_classes=0):
-    """~25M params — 256×256"""
     return LiquidFlowNet(img_size=img_size, patch_size=8, d_model=384, depth=10, d_state=16, expand=2, num_classes=num_classes)
 
 def liquidflow_512(img_size=512, num_classes=0):
-    """~25M params — 512×512"""
     return LiquidFlowNet(img_size=img_size, patch_size=16, d_model=384, depth=10, d_state=16, expand=2, num_classes=num_classes)