geneflowdecoder 0.1.4__tar.gz → 0.1.6__tar.gz

{geneflowdecoder-0.1.4/src/geneflowdecoder.egg-info → geneflowdecoder-0.1.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: geneflowdecoder
-Version: 0.1.4
+Version: 0.1.6
 Summary: Gene Transformer + Conditional Rectified Flow decoder from cell-state z to log1p expression X.
 Author: GeneFlowDecoder Contributors
 License: MIT

{geneflowdecoder-0.1.4 → geneflowdecoder-0.1.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "geneflowdecoder"
-version = "0.1.4"
+version = "0.1.6"
 description = "Gene Transformer + Conditional Rectified Flow decoder from cell-state z to log1p expression X."
 readme = "README.md"
 requires-python = ">=3.9"

geneflowdecoder-0.1.6/src/geneflowdecoder/_core.py

@@ -0,0 +1,549 @@
+# -*- coding: utf-8 -*-
+"""
+GeneFlowDecoder 0.1.2
+z (cell state, *frozen*) --> X (log-normalized expression, log1p space)
+
+Public API (stable):
+    model = GeneFlowDecoder(n_genes=G, z_dim=D, ...)
+    model.fit(X, z, X_val=None, z_val=None, ...)
+    model(z)                 # sample X from z  (flow, stochastic)
+    model.predict(z)         # deterministic X  (NB mean, log1p space)
+    model.save(path) / GeneFlowDecoder.load(path, **init_kwargs)
+
+Requires: torch >= 2.1
+"""
+
+from __future__ import annotations
+
+import math
+import warnings
+from dataclasses import dataclass, field
+from typing import Any, Optional, Tuple, Union
+
+import numpy as np
+
+# ----------------------------------------------------------------------
+# 可选的 torch 导入：缺 torch 时 import 不崩，构造模型时才明确报错
+# ----------------------------------------------------------------------
+try:
+    import torch
+    import torch.nn as nn
+    import torch.nn.functional as F
+    from torch import Tensor
+
+    _HAS_TORCH = True
+except Exception:  # pragma: no cover
+    _HAS_TORCH = False
+    torch = None  # type: ignore[assignment]
+    nn = None  # type: ignore[assignment]
+    F = None  # type: ignore[assignment]
+
+    class Tensor:  # noqa: N801 -- placeholder type, never instantiated
+        """Stand-in so ``Array`` remains valid at annotation-check time."""
+
+        def __init_subclass__(cls, **_: Any) -> None:
+            pass
+
+    class _Noop:
+        def __getattr__(self, _name: str) -> "_Noop":
+            return self
+
+        def __call__(self, *__args: Any, **__kw: Any) -> None:
+            return None
+
+    nn = _Noop()  # type: ignore[assignment]
+    F = _Noop()  # type: ignore[assignment]
+
+
+Array = Union[np.ndarray, "Tensor"]
+
+
+def _as_tensor(x: Array, device: "torch.device") -> "Tensor":
+    """numpy -> torch; 已在目标设备的 tensor 原样返回；cuda 输入不再拉回 cpu。"""
+    if isinstance(x, np.ndarray):
+        x = torch.from_numpy(x)
+    t: Tensor = x
+    if t.device != device:
+        t = t.to(device, non_blocking=True)
+    return t
+
+
+def _require_torch() -> None:
+    if not _HAS_TORCH:
+        raise RuntimeError(
+            "GeneFlowDecoder requires PyTorch (>= 2.1). "
+            "Install it with: pip install 'torch>=2.1'"
+        )
+
+
+# ============================================================
+# 条件化模块
+# ============================================================
+
+class _AdaLNModulation(nn.Module):
+    def __init__(self, d: int, cond_dim: int):
+        super().__init__()
+        self.norm = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)
+        self.proj = nn.Sequential(
+            nn.Linear(cond_dim, d), nn.SiLU(), nn.Linear(d, 3 * d)
+        )
+        nn.init.zeros_(self.proj[-1].weight)
+        nn.init.zeros_(self.proj[-1].bias)
+
+    def forward(self, x: Tensor, z: Tensor) -> Tensor:
+        g, b, gate = self.proj(z).chunk(3, dim=-1)
+        g = g.unsqueeze(1)
+        b = b.unsqueeze(1)
+        gate = gate.unsqueeze(1).sigmoid()
+        return gate * (self.norm(x) * (1 + g) + b)
+
+
+class _CondBlock(nn.Module):
+    def __init__(self, d: int, cond_dim: int, n_heads: int, mlp_ratio: float = 4.0):
+        super().__init__()
+        self.attn_norm = _AdaLNModulation(d, cond_dim)
+        self.attn = nn.MultiheadAttention(d, n_heads, batch_first=True)
+        self.mlp_norm = _AdaLNModulation(d, cond_dim)
+        hd = int(d * mlp_ratio)
+        self.mlp = nn.Sequential(nn.Linear(d, hd), nn.GELU(), nn.Linear(hd, d))
+
+    def forward(self, x: Tensor, z: Tensor) -> Tensor:
+        h = self.attn_norm(x, z)
+        h, _ = self.attn(h, h, h, need_weights=False)
+        x = x + h
+        return x + self.mlp(self.mlp_norm(x, z))
+
+
+# ============================================================
+# Gene Transformer backbone
+# 输出恒为 [B, G, d]；chunked 分支显式保证拼接后形状
+# ============================================================
+
+class _GeneTransformer(nn.Module):
+    def __init__(
+        self,
+        n_genes: int,
+        d: int,
+        depth: int,
+        n_heads: int,
+        cond_dim: int,
+        chunk_size: int,
+        ckpt: bool,
+    ):
+        super().__init__()
+        self.d = d
+        self.n_genes = n_genes
+        self.chunk_size = max(1, chunk_size)
+        self._ckpt = ckpt
+
+        self.gene_embed = nn.Embedding(n_genes, d)
+        nn.init.normal_(self.gene_embed.weight, std=0.02)
+        self.pos_embed = nn.Parameter(torch.zeros(1, n_genes, d))
+        nn.init.normal_(self.pos_embed, std=0.02)
+
+        self.blocks = nn.ModuleList([_CondBlock(d, cond_dim, n_heads) for _ in range(depth)])
+        self.head_norm = nn.LayerNorm(d, eps=1e-6)
+
+    def _run(self, x: Tensor, z: Tensor) -> Tensor:
+        for blk in self.blocks:
+            if self._ckpt and self.training:
+                x = torch.utils.checkpoint.checkpoint(blk, x, z, use_reentrant=False)
+            else:
+                x = blk(x, z)
+        return self.head_norm(x)
+
+    def forward(self, gene_ids: Tensor, z: Tensor) -> Tensor:
+        B, G = gene_ids.shape[0], gene_ids.shape[1]
+        if G <= self.chunk_size:
+            h = self.gene_embed(gene_ids) + self.pos_embed[:, :G, :]
+            return self._run(h, z)
+
+        chunks: list[Tensor] = []
+        for s in range(0, G, self.chunk_size):
+            e = min(s + self.chunk_size, G)
+            gids = gene_ids[:, s:e]
+            h = self.gene_embed(gids) + self.pos_embed[:, s:e, :]
+            chunks.append(self._run(h, z))
+        out = torch.cat(chunks, dim=1)
+        # 防御性断言：backbone 永远交出 [B, G, d]
+        if out.shape != (B, G, self.d):  # pragma: no cover - 仅开发期暴露
+            raise RuntimeError(
+                f"backbone shape invariant broken: got {tuple(out.shape)}, "
+                f"expected ({B},{G},{self.d})"
+            )
+        return out
+
+
+# ============================================================
+# NB 输出头
+# ============================================================
+
+class _NBHead(nn.Module):
+    def __init__(self, d: int, n_genes: int):
+        super().__init__()
+        self.mu_proj = nn.Linear(d, n_genes)
+        self.theta = nn.Parameter(torch.zeros(n_genes))
+        nn.init.zeros_(self.mu_proj.weight)
+        nn.init.zeros_(self.mu_proj.bias)
+
+    def forward(self, h: Tensor) -> Tuple[Tensor, Tensor]:
+        # h: [B, G, d] -> log_mu: [B, G]
+        return self.mu_proj(h).squeeze(-1), F.softplus(self.theta) + 1e-3
+
+    @staticmethod
+    def nll(
+        log_mu: Tensor,
+        theta: Tensor,
+        x: Tensor,
+        mask: Optional[Tensor] = None,
+    ) -> Tensor:
+        mu = log_mu.exp().clamp(min=1e-6)
+        th = theta.unsqueeze(0).expand_as(x)
+        ll = (
+            torch.lgamma(x + th)
+            - torch.lgamma(th)
+            - torch.lgamma(x + 1)
+            + th * (th.log() - (th + mu).log())
+            + x * (mu.log() - (th + mu).log())
+        )
+        if mask is not None:
+            ll = ll * mask
+            return -ll.sum() / mask.sum().clamp(min=1.0)
+        return -ll.mean()
+
+
+@dataclass
+class FlowConfig:
+    """训练目标权重；均为非负标量。"""
+
+    sigma_min: float = 1e-4
+    nb_weight: float = 1.0
+    corr_weight: float = 0.1
+
+
+# ============================================================
+# 主类
+# ============================================================
+
+class GeneFlowDecoder(nn.Module):
+    """
+    z -> X (log1p expression).
+
+    ``z`` 默认视为**固定表征**（freeze_z=True）：不参与反传、不被优化器更新，
+    fit() 开头还会对其做一次 detach + clone 快照，训练中途外部 encoder
+    若继续更新也不会污染本 decoder 的输入。
+
+    Examples
+    --------
+    >>> model = GeneFlowDecoder(n_genes=2048, z_dim=256, d=512, depth=8)
+    >>> model.fit(X_train, z_train, epochs=200, batch_size=64)
+    >>> X_hat = model(z_test)          # 采样 (flow, stochastic)
+    >>> X_rec = model.predict(z_test)  # 确定性还原: NB mean (log1p 空间)
+    """
+
+    def __init__(
+        self,
+        n_genes: int,
+        z_dim: int,
+        d: int = 768,
+        depth: int = 16,
+        n_heads: int = 12,
+        chunk_size: int = 2048,
+        gradient_checkpoint: bool = True,
+        cfg: Optional[FlowConfig] = None,
+        freeze_z: bool = True,
+        device: str = "auto",
+    ):
+        _require_torch()
+        super().__init__()
+
+        if device == "auto":
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = torch.device(device)
+        self.n_genes = int(n_genes)
+        self.z_dim = int(z_dim)
+        self.freeze_z = bool(freeze_z)
+        self._cfg = cfg if cfg is not None else FlowConfig()
+
+        self.backbone = _GeneTransformer(
+            self.n_genes, d, depth, n_heads, self.z_dim, chunk_size, gradient_checkpoint
+        )
+        self.nb_head = _NBHead(d, self.n_genes)
+
+        # ★ 关键：每个 (cell, gene) 位置预测一个标量速度，out_features 必须是 1
+        self.vel_proj = nn.Linear(d, 1)
+        self.time_mlp = nn.Sequential(nn.Linear(128, d), nn.SiLU(), nn.Linear(d, d))
+
+        self._gene_ids = nn.Parameter(torch.arange(self.n_genes), requires_grad=False)
+
+        self.to(self.device)
+
+    # --------------------------------------------------------
+    # 时间编码
+    # --------------------------------------------------------
+    def _sincos(self, t: Tensor) -> Tensor:
+        half = 64
+        freqs = torch.exp(
+            -math.log(10000.0) * torch.arange(half, device=t.device) / (half - 1)
+        )
+        a = t[:, None] * freqs[None, :]
+        e = torch.cat([a.sin(), a.cos()], dim=-1)
+        if e.shape[-1] < 128:
+            e = F.pad(e, (0, 128 - e.shape[-1]))
+        return e
+
+    # --------------------------------------------------------
+    # z 隔离：freeze_z 时 detach；fit 入口额外 clone 成快照
+    # --------------------------------------------------------
+    def _maybe_detach(self, z: Tensor, *, snapshot: bool = False) -> Tensor:
+        if not self.freeze_z:
+            return z
+        z = z.detach()
+        return z.clone() if snapshot else z
+
+    # --------------------------------------------------------
+    # 速度场 v(x_t, t, z) -> [B, G]
+    # --------------------------------------------------------
+    def _velocity(self, x_t: Tensor, t: Tensor, z: Tensor) -> Tensor:
+        z = self._maybe_detach(z)
+        B, G = x_t.shape
+        t_emb = self.time_mlp(self._sincos(t)).unsqueeze(1)  # [B,1,d]
+
+        # 明确 [B, G] 的 gene id 网格，避免广播维歧义
+        gids = self._gene_ids[:G].unsqueeze(0).expand(B, -1)
+        h = self.backbone(gids, z) + t_emb              # [B,G,d]
+
+        v = self.vel_proj(h).squeeze(-1)                # [B,G,1]->[B,G]
+
+        if v.ndim != 2 or v.shape != (B, G):            # pragma: no cover
+            raise RuntimeError(
+                f"_velocity shape bug: got {tuple(v.shape)}, expected ({B},{G}). "
+                f"vel_proj 必须是 nn.Linear(d, 1)，请重装 >=0.1.2 版本。"
+            )
+        return v
+
+    # --------------------------------------------------------
+    # 损失
+    # --------------------------------------------------------
+    def _loss(
+        self,
+        x: Tensor,
+        z: Tensor,
+        mask: Optional[Tensor] = None,
+    ) -> Tensor:
+        z = self._maybe_detach(z)
+        B = x.shape[0]
+        t = torch.rand(B, device=self.device).clamp(
+            self._cfg.sigma_min, 1.0 - self._cfg.sigma_min
+        )
+        eps = torch.randn_like(x)
+        x_t = x.lerp(eps, t.unsqueeze(1))
+
+        v_pred = self._velocity(x_t, t, z)      # [B,G]
+        v_target = eps - x                       # [B,G]，rectified-flow 目标
+        mse = F.mse_loss(v_pred, v_target, reduction="none")
+        if mask is not None:
+            mse = mse * mask
+            mse = mse.sum() / mask.sum().clamp(min=1.0)
+        else:
+            mse = mse.mean()
+
+        x0_hat = x_t.lerp(v_pred.detach(), t.unsqueeze(1))
+
+        nb = torch.tensor(0.0, device=self.device)
+        if self._cfg.nb_weight > 0:
+            gids = self._gene_ids[: x.shape[1]].unsqueeze(0).expand(B, -1)
+            log_mu, theta = self.nb_head(self.backbone(gids, z))
+            nb = self.nb_head.nll(log_mu, theta, x, mask)
+
+        corr = torch.tensor(0.0, device=self.device)
+        if self._cfg.corr_weight > 0 and mask is None:
+            corr = self._corr(x0_hat, x)
+
+        return mse + self._cfg.nb_weight * nb + self._cfg.corr_weight * corr
+
+    @staticmethod
+    def _corr(p: Tensor, q: Tensor) -> Tensor:
+        p = p - p.mean(1, keepdim=True)
+        q = q - q.mean(1, keepdim=True)
+        d = p.norm(dim=1) * q.norm(dim=1) + 1e-8
+        return (1 - (p * q).sum(dim=1) / d).mean()
+
+    # --------------------------------------------------------
+    # fit
+    # --------------------------------------------------------
+    @staticmethod
+    def _want_pin_memory(X: Tensor, z: Tensor) -> bool:
+        """pin_memory 只对 CPU tensor 有意义；cuda 输入直接关掉。"""
+
+        def on_gpu(a: Tensor) -> bool:
+            return isinstance(a, torch.Tensor) and a.device.type != "cpu"
+
+        return not (on_gpu(X) or on_gpu(z))
+
+    def fit(
+        self,
+        X: Array,
+        z: Array,
+        X_val: Optional[Array] = None,
+        z_val: Optional[Array] = None,
+        epochs: int = 200,
+        batch_size: int = 64,
+        lr: float = 2e-4,
+        weight_decay: float = 0.02,
+        patience: int = 20,
+        num_workers: int = 0,
+        verbose: int = 1,
+    ) -> "GeneFlowDecoder":
+        X = _as_tensor(X, self.device)
+        z = _as_tensor(z, self.device)
+        # ★ 快照：freeze_z 时拷出一份独立副本，训练期与上游 z 彻底解耦
+        z = self._maybe_detach(z, snapshot=True)
+
+        if X.ndim != 2 or z.ndim != 2:
+            raise ValueError(f"X/z 必须是 2D: got {X.ndim}D / {z.ndim}D")
+        if X.shape[0] != z.shape[0]:
+            raise ValueError(f"X 与 z 样本数不一致: {X.shape[0]} vs {z.shape[0]}")
+        if z.shape[1] != self.z_dim:
+            raise ValueError(f"z 特征维 {z.shape[1]} 与 z_dim={self.z_dim} 不符")
+        if X.shape[1] != self.n_genes:
+            raise ValueError(
+                f"X 基因数 {X.shape[1]} 与 n_genes={self.n_genes} 不符；"
+                f"请用实际 X.shape[1] 构造模型"
+            )
+
+        # ★ 先算 pin_memory，再断言（0.1.1 的顺序 bug 已修正）
+        pin_memory = self._want_pin_memory(X, z)
+        assert (
+            X.device.type == "cpu" or z.device.type == "cpu" or not pin_memory
+        ), "cuda 输入不应开启 pin_memory"
+
+        ds = torch.utils.data.TensorDataset(X, z)
+        dl = torch.utils.data.DataLoader(
+            ds,
+            batch_size=int(batch_size),
+            shuffle=True,
+            num_workers=int(num_workers),
+            pin_memory=pin_memory,
+        )
+
+        has_val = False
+        if X_val is not None and z_val is not None:
+            Xv = _as_tensor(X_val, self.device)
+            zv = self._maybe_detach(_as_tensor(z_val, self.device), snapshot=True)
+            if Xv.shape[1] != self.n_genes or zv.shape[1] != self.z_dim:
+                raise ValueError("验证集维度与训练集不一致")
+            has_val = True
+
+        params = [p for p in self.parameters() if p.requires_grad]
+        opt = torch.optim.AdamW(params, lr=lr, weight_decay=weight_decay)
+        sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=int(epochs))
+
+        best_val = float("inf")
+        best_state: Optional[dict[str, Tensor]] = None
+        no_improve = 0
+
+        for ep in range(1, int(epochs) + 1):
+            self.train()
+            tot = 0.0
+            n_seen = 0
+            for xb, zb in dl:
+                opt.zero_grad()
+                loss = self._loss(xb, zb)
+                loss.backward()
+                nn.utils.clip_grad_norm_(params, 1.0)
+                opt.step()
+                tot += loss.item() * xb.shape[0]
+                n_seen += xb.shape[0]
+            sched.step()
+
+            if verbose:
+                msg = f"[fit] epoch {ep:3d}/{epochs} | train={tot / max(n_seen,1):.4f}"
+                if has_val:
+                    self.eval()
+                    with torch.no_grad():
+                        vl = self._loss(Xv, zv).item()
+                    msg += f" | val={vl:.4f}"
+                    if vl < best_val - 1e-5:
+                        best_val = vl
+                        no_improve = 0
+                        best_state = {
+                            k: v.detach().cpu().clone()
+                            for k, v in self.state_dict().items()
+                        }
+                    else:
+                        no_improve += 1
+                    if no_improve >= int(patience):
+                        if verbose:
+                            print(f"[fit] early stop at epoch {ep}")
+                        break
+                print(msg)
+
+        if best_state is not None:
+            self.load_state_dict(best_state)
+            if verbose:
+                print(f"[fit] restored best val={best_val:.4f}")
+        return self
+
+    # --------------------------------------------------------
+    # 采样 / 确定性还原
+    # --------------------------------------------------------
+    @torch.no_grad()
+    def forward(
+        self,
+        z: Array,
+        steps: int = 4,
+        method: str = "midpoint",
+        deterministic: bool = False,
+    ) -> Tensor:
+        """
+        给定 z 采样 X（log1p 空间）。
+        deterministic=True 等价于 predict()；steps=4 通常够用。
+        """
+        z = self._maybe_detach(_as_tensor(z, self.device))
+        if deterministic:
+            return self.predict(z)
+
+        self.eval()
+        B = z.shape[0]
+        G = self.n_genes
+        x = torch.randn(B, G, device=self.device)
+        ts = torch.linspace(0.0, 1.0, int(steps) + 1, device=self.device)
+        for i in range(int(steps)):
+            t0, t1 = ts[i], ts[i + 1]
+            tmid = (t0 + t1) / 2
+            dt = float(t1 - t0)
+            v = self._velocity(x, torch.full((B,), t0, device=self.device), z)
+            if method == "midpoint":
+                xm = x + 0.5 * dt * v
+                v = self._velocity(xm, torch.full((B,), tmid, device=self.device), z)
+            x = x + dt * v
+        return x
+
+    @torch.no_grad()
+    def predict(self, z: Array) -> Tensor:
+        """确定性还原：NB 均值（log1p 空间）。伪计数用 (exp(log_mu)-1).clamp(min=0)。"""
+        z = self._maybe_detach(_as_tensor(z, self.device))
+        self.eval()
+        B = z.shape[0]
+        gids = self._gene_ids[:self.n_genes].unsqueeze(0).expand(B, -1)
+        log_mu, _ = self.nb_head(self.backbone(gids, z))
+        return log_mu
+
+    # --------------------------------------------------------
+    # 持久化
+    # --------------------------------------------------------
+    def save(self, path: str) -> None:
+        torch.save(
+            {"state": self.state_dict(), "cfg": self._cfg, "freeze_z": self.freeze_z},
+            path,
+        )
+
+    @classmethod
+    def load(cls, path: str, **init_kwargs: Any) -> "GeneFlowDecoder":
+        ckpt = torch.load(path, map_location="cpu")
+        init_kwargs.setdefault("freeze_z", bool(ckpt.get("freeze_z", True)))
+        model = cls(**init_kwargs)
+        model._cfg = ckpt.get("cfg", model._cfg)
+        model.load_state_dict(ckpt["state"])
+        return model

{geneflowdecoder-0.1.4 → geneflowdecoder-0.1.6}/src/geneflowdecoder/_version.py

@@ -1,2 +1,2 @@
 # -*- coding: utf-8 -*-
-__version__ = "0.1.4"
+__version__ = "0.1.6"

{geneflowdecoder-0.1.4 → geneflowdecoder-0.1.6/src/geneflowdecoder.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: geneflowdecoder
-Version: 0.1.4
+Version: 0.1.6
 Summary: Gene Transformer + Conditional Rectified Flow decoder from cell-state z to log1p expression X.
 Author: GeneFlowDecoder Contributors
 License: MIT

geneflowdecoder-0.1.4/src/geneflowdecoder/_core.py

@@ -1,679 +0,0 @@
-# -*- coding: utf-8 -*-
-"""
-GeneFlowDecoder
-===============
-
-z (fixed cell-state representation) --> X (log-normalized expression, log1p space)
-
-Public API
-----------
-    model = GeneFlowDecoder(n_genes=2048, z_dim=256)
-    model.fit(X, z)                      # train
-    X_hat = model(z)                     # sample (generative)
-    X_hat = model.predict(z)             # deterministic (NB mean)
-
-This module is framework-agnostic beyond ``torch`` and ``numpy``.
-"""
-
-from __future__ import annotations
-
-import math
-from dataclasses import dataclass, field
-from typing import Optional, Tuple, Union
-
-import numpy as np
-
-try:
-    import torch
-    import torch.nn as nn
-    import torch.nn.functional as F
-    _no_grad = torch.no_grad
-    _TORCH_AVAILABLE = True
-except Exception:  # pragma: no cover - soft dep for import-only environments
-    import types
-
-    torch = types.ModuleType("torch")
-    nn = types.ModuleType("torch.nn")
-    F = types.ModuleType("torch.nn.functional")
-    class _NoGrad:
-        """无 torch 时充当 ``@torch.no_grad()`` 与 ``with torch.no_grad():`` 的透明桩。"""
-
-        def __call__(self, fn=None):
-            return fn if fn is not None else self
-
-        def __enter__(self):
-            return self
-
-        def __exit__(self, *args):
-            pass
-
-    _no_grad = _NoGrad()
-    _TORCH_AVAILABLE = False
-
-    class _Device(str):
-        """``torch.device(device)`` 的桩：保留字符串语义与 ``.type`` 属性。"""
-
-        def __new__(cls, dev):
-            dev = str(dev)
-            if dev.startswith("cuda"):
-                dev = "cuda"
-            return super().__new__(cls, dev)
-
-        @property
-        def type(self):
-            return str(self)
-
-
-    # 让 ``class Foo(nn.Module)`` 与 ``Tensor`` 注解在无 torch 时仍能解析。
-    # 运行时真正调用前向/构造模型会立即报清晰的错误，而不是在 import 期崩。
-    class _ModuleStub:
-        def __init_subclass__(cls, **kwargs):
-            pass
-
-    nn.Module = _ModuleStub  # type: ignore[attr-defined]
-
-    # 补齐 GeneFlowDecoder 类体里会用到的 torch 属性，保持 import 期可解析
-    torch.Tensor = object  # type: ignore[attr-defined]
-    torch.no_grad = _no_grad  # 即 _NoGrad() 实例  # type: ignore[attr-defined]
-    torch.device = _Device  # type: ignore[attr-defined]
-    torch.cuda = types.ModuleType("torch.cuda")
-    torch.cuda.is_available = lambda: False
-    torch.utils = types.ModuleType("torch.utils")
-    torch.utils.data = types.ModuleType("torch.utils.data")
-    torch.utils.data.TensorDataset = object
-    torch.utils.data.DataLoader = object
-    Tensor = object  # placeholder for type annotations only
-else:
-    try:
-        from torch import Tensor as _TorchTensor  # noqa: F401
-        Tensor = _TorchTensor  # type: ignore
-    except Exception:
-        Tensor = object  # placeholder for type annotations only
-
-Array = Union[np.ndarray, Tensor]  # type: ignore[valid-type]
-
-
-def _as_tensor(x: Array, device: "torch.device") -> "torch.Tensor":
-    """把 numpy / CPU / GPU 输入规范化成 torch 张量。
-
-    关键：输入**已经**在 GPU 上时，保留其原设备，不再 ``.to(device)``。
-    否则 cuda 输入会被拷回 CPU，DataLoader 的 pin_memory 又会对它报错：
-        RuntimeError: cannot pin 'torch.cuda.FloatTensor'
-    """
-    if not _TORCH_AVAILABLE:
-        raise RuntimeError(
-            "GeneFlowDecoder requires PyTorch, but it could not be imported. "
-            "Install it with: pip install 'torch>=2.1'"
-        )
-    if isinstance(x, np.ndarray):
-        return torch.as_tensor(x, device=device)      # CPU numpy -> CPU tensor
-    if x.device.type != "cpu":
-        return x                                     # 已在 GPU：原样返回
-    return x.to(device, non_blocking=True)
-
-
-def _is_cuda_array(x: Array) -> bool:
-    """判断输入是否已经落在 GPU 上。
-
-    numpy 数组一定在 CPU；torch 张量则可能已在 cuda。DataLoader 的
-    ``pin_memory`` 只接受 CPU tensor，输入已经是 cuda 时必须降级为 False，
-    否则会在 ``pin_memory()`` 处抛出
-    ``RuntimeError: cannot pin 'torch.cuda.FloatTensor'``。
-
-    注意：无 torch 时 ``torch.Tensor`` 只是一个占位桩（object），此时 numpy
-    数组不该被当成 Tensor，故先用 ``not isinstance(x, (np.ndarray, ...))``
-    兜底，保证纯 numpy 输入永远返回 False（即开启 pin_memory）。
-    """
-    if isinstance(x, np.ndarray):
-        return False
-    try:
-        return isinstance(x, torch.Tensor) and x.device.type != "cpu"
-    except Exception:
-        return False
-
-
-def _want_pin_memory(X: Array, z: Array) -> bool:
-    """输入全在 CPU 时开启 pin_memory 加速；任一在 GPU 则关闭。
-
-    numpy 输入一律视为 CPU 数据，走加速路径；只有在 torch 可用且输入
-    device 为 cuda 时才关闭。
-    """
-    if isinstance(X, np.ndarray) and isinstance(z, np.ndarray):
-        return True
-    return not (_is_cuda_array(X) or _is_cuda_array(z))
-
-
-# ----------------------------------------------------------------------
-# 1. Conditional modulation & transformer blocks
-# ----------------------------------------------------------------------
-
-class _AdaLNModulation(nn.Module):
-    def __init__(self, d: int, cond_dim: int):
-        super().__init__()
-        self.norm = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)
-        self.proj = nn.Sequential(
-            nn.Linear(cond_dim, d), nn.SiLU(), nn.Linear(d, 3 * d)
-        )
-        nn.init.zeros_(self.proj[-1].weight)
-        nn.init.zeros_(self.proj[-1].bias)
-
-    def forward(self, x: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
-        g, b, gate = self.proj(z).chunk(3, dim=-1)
-        g, b, gate = g.unsqueeze(1), b.unsqueeze(1), gate.unsqueeze(1).sigmoid()
-        return gate * (self.norm(x) * (1 + g) + b)
-
-
-class _CondBlock(nn.Module):
-    def __init__(self, d: int, cond_dim: int, n_heads: int, mlp_ratio: float = 4.0):
-        super().__init__()
-        self.attn_norm = _AdaLNModulation(d, cond_dim)
-        self.attn = nn.MultiheadAttention(d, n_heads, batch_first=True)
-        self.mlp_norm = _AdaLNModulation(d, cond_dim)
-        h = int(d * mlp_ratio)
-        self.mlp = nn.Sequential(nn.Linear(d, h), nn.GELU(), nn.Linear(h, d))
-
-    def forward(self, x: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
-        h = self.attn_norm(x, z)
-        h, _ = self.attn(h, h, h, need_weights=False)
-        x = x + h
-        return x + self.mlp(self.mlp_norm(x, z))
-
-
-class _GeneTransformer(nn.Module):
-    def __init__(
-        self,
-        n_genes: int,
-        d: int,
-        depth: int,
-        n_heads: int,
-        cond_dim: int,
-        chunk_size: int,
-        ckpt: bool,
-    ):
-        super().__init__()
-        self.d = d
-        self.n_genes = n_genes
-        self.chunk_size = chunk_size
-        self.gene_embed = nn.Embedding(n_genes, d)
-        nn.init.normal_(self.gene_embed.weight, std=0.02)
-        self.pos_embed = nn.Parameter(torch.zeros(1, n_genes, d))
-        nn.init.normal_(self.pos_embed, std=0.02)
-
-        self.blocks = nn.ModuleList(
-            [_CondBlock(d, cond_dim, n_heads) for _ in range(depth)]
-        )
-        self.head_norm = nn.LayerNorm(d, eps=1e-6)
-        self._ckpt = ckpt
-
-    def _run(self, x: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
-        for blk in self.blocks:
-            if self._ckpt and self.training:
-                x = torch.utils.checkpoint.checkpoint(
-                    blk, x, z, use_reentrant=False
-                )
-            else:
-                x = blk(x, z)
-        return self.head_norm(x)
-
-    def forward(self, gene_ids: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
-        B, G = gene_ids.shape[0], gene_ids.shape[-1]
-        if G <= self.chunk_size:
-            h = self.gene_embed(gene_ids) + self.pos_embed[:, :G, :]
-            return self._run(h, z)
-
-        outs: list[torch.Tensor] = []
-        s = 0
-        while s < G:
-            e = min(s + self.chunk_size, G)
-            h = self.gene_embed(gene_ids[..., s:e]) + self.pos_embed[:, s:e, :]
-            outs.append(self._run(h, z))
-            s = e
-        return torch.cat(outs, dim=1)
-
-
-# ----------------------------------------------------------------------
-# 2. Negative-binomial output head
-# ----------------------------------------------------------------------
-
-class _NBHead(nn.Module):
-    def __init__(self, d: int, n_genes: int):
-        super().__init__()
-        self.mu_proj = nn.Linear(d, n_genes)
-        self.theta = nn.Parameter(torch.zeros(n_genes))
-        nn.init.zeros_(self.mu_proj.weight)
-        nn.init.zeros_(self.mu_proj.bias)
-
-    def forward(self, h: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-        return self.mu_proj(h), F.softplus(self.theta) + 1e-3
-
-    @staticmethod
-    def nll(
-        log_mu: torch.Tensor,
-        theta: torch.Tensor,
-        x: torch.Tensor,
-        mask: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        mu = log_mu.exp().clamp(min=1e-6)
-        th = theta.unsqueeze(0).expand_as(x)
-        ll = (
-            torch.lgamma(x + th)
-            - torch.lgamma(th)
-            - torch.lgamma(x + 1)
-            + th * (th.log() - (th + mu).log())
-            + x * (mu.log() - (th + mu).log())
-        )
-        if mask is not None:
-            ll = ll * mask
-            return -ll.sum() / mask.sum().clamp(min=1.0)
-        return -ll.mean()
-
-
-@dataclass
-class FlowConfig:
-    """Hyper-parameters of the conditional flow matching objective."""
-
-    sigma_min: float = 1e-4
-    nb_weight: float = 1.0
-    corr_weight: float = 0.1
-
-
-# ----------------------------------------------------------------------
-# 3. GeneFlowDecoder
-# ----------------------------------------------------------------------
-
-class GeneFlowDecoder(nn.Module):
-    """
-    Conditional Rectified-Flow decoder from a **fixed** cell-state vector *z*
-    to a log1p expression profile *X*.
-
-    Parameters
-    ----------
-    n_genes : int
-        Number of genes (output dimension).
-    z_dim : int
-        Dimensionality of the cell-state representation.
-    d, depth, n_heads : int
-        Transformer width, number of layers, number of attention heads.
-    chunk_size : int
-        Chunk size for the gene axis (enables >10k gene panels on one GPU).
-    gradient_checkpoint : bool
-        Enable activation checkpointing inside the transformer.
-    cfg : FlowConfig | None
-        Flow / loss hyper-parameters.
-    freeze_z : bool, default ``True``
-        When ``True`` (recommended), *z* is detached everywhere: it is treated
-        as a fixed, non-trainable conditioning signal and is never updated.
-    device : ``"auto"`` | ``str`` | ``torch.device``
-        Where to place the model. ``"auto"`` picks CUDA if available.
-
-    Examples
-    --------
-    >>> import numpy as np
-    >>> from geneflowdecoder import GeneFlowDecoder
-    >>> G, D = 2048, 128
-    >>> model = GeneFlowDecoder(n_genes=G, z_dim=D, d=512, depth=8)
-    >>> X = np.random.randn(1024, G).astype("float32")
-    >>> z = np.random.randn(1024, D).astype("float32")
-    >>> model.fit(X, z, epochs=10, batch_size=64)
-    >>> X_hat = model(z[:4])           # sampled, log1p space
-    >>> X_rec  = model.predict(z[:4])  # deterministic NB mean
-    """
-
-    def __init__(
-        self,
-        n_genes: int,
-        z_dim: int,
-        *,
-        d: int = 768,
-        depth: int = 16,
-        n_heads: int = 12,
-        chunk_size: int = 2048,
-        gradient_checkpoint: bool = True,
-        cfg: Optional[FlowConfig] = None,
-        freeze_z: bool = True,
-        device: str = "auto",
-    ):
-        super().__init__()
-        if not _TORCH_AVAILABLE:
-            raise RuntimeError(
-                "GeneFlowDecoder requires PyTorch, but it could not be imported. "
-                "Install it with: pip install 'torch>=2.1'"
-            )
-        if device == "auto":
-            device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.device = torch.device(device)
-        self.n_genes = int(n_genes)
-        self.z_dim = int(z_dim)
-        self.freeze_z = bool(freeze_z)
-        self._cfg = cfg or FlowConfig()
-
-        self.backbone = _GeneTransformer(
-            n_genes, d, depth, n_heads, z_dim, chunk_size, gradient_checkpoint
-        )
-        self.nb_head = _NBHead(d, n_genes)
-        self.vel_proj = nn.Linear(d, n_genes)
-        self.time_mlp = nn.Sequential(
-            nn.Linear(128, d), nn.SiLU(), nn.Linear(d, d)
-        )
-        self._gene_ids = nn.Parameter(
-            torch.arange(self.n_genes), requires_grad=False
-        )
-        self.to(self.device)
-
-    # ----------------------------------------------------------------
-    # internals
-    # ----------------------------------------------------------------
-    def _sincos(self, t: torch.Tensor) -> torch.Tensor:
-        half = 64
-        freqs = torch.exp(
-            -math.log(10000.0) * torch.arange(half, device=t.device) / (half - 1)
-        )
-        a = t[:, None] * freqs[None, :]
-        e = torch.cat([a.sin(), a.cos()], dim=-1)
-        if e.shape[-1] < 128:
-            e = F.pad(e, (0, 128 - e.shape[-1]))
-        return e
-
-    def _maybe_detach(
-        self, z: torch.Tensor, *, snapshot: bool = False
-    ) -> torch.Tensor:
-        """冻结 z 时截断计算图；``snapshot=True`` 再额外拷一份独立副本。
-
-        snapshot 只在 fit 入口用：把当轮 z 固化下来，之后外部 encoder
-        若继续 backward / 原地更新，训练集 z 的数值也不会跟着变。
-        单次推理（forward / predict）不开 snapshot，省一次显存拷贝。
-        """
-        if not self.freeze_z:
-            return z
-        z = z.detach()
-        return z.clone() if snapshot else z
-
-    def _velocity(self, x_t: torch.Tensor, t: torch.Tensor,
-                z: torch.Tensor) -> torch.Tensor:
-        if self.freeze_z:
-            z = z.detach()
-        B, G = x_t.shape
-        t_emb = self.time_mlp(self._sincos(t)).unsqueeze(1)        # [B,1,d]
-        gids = self._gene_ids[:G].unsqueeze(0).expand(B, -1)       # [B,G]
-        h = self.backbone(gids, z) + t_emb                         # [B,G,d]
-
-        # ★ 关键：只对特征维 d 做线性，结果天然是 [B, G]
-        v = self.vel_proj(h.squeeze(-2) if h.ndim == 2 else h)     # [B,G,d]→[B,G]
-        if v.ndim == 3:
-            v = v.squeeze(-1)                                      # [B,G,1]→[B,G]
-        if v.ndim != 2 or v.shape != (B, G):
-            raise RuntimeError(
-                f"_velocity shape bug: got {tuple(v.shape)}, expected ({B},{G}). "
-                f"检查 vel_proj 是否为 nn.Linear(d, 1) 而非 nn.Linear(d, n_genes)"
-            )
-        return v
-
-    # ----------------------------------------------------------------
-    # objective
-    # ----------------------------------------------------------------
-    def _loss(
-        self,
-        x: torch.Tensor,
-        z: torch.Tensor,
-        mask: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        z = self._maybe_detach(z)
-        B = x.shape[0]
-        t = torch.rand(B, device=self.device).clamp(
-            self._cfg.sigma_min, 1.0 - self._cfg.sigma_min
-        )
-        eps = torch.randn_like(x)
-        x_t = x.lerp(eps, t.unsqueeze(1))
-
-        v_pred = self._velocity(x_t, t, z)
-        v_target = eps - x  # rectified-flow target
-        mse = F.mse_loss(v_pred, v_target, reduction="none")
-        if mask is not None:
-            mse = mse * mask
-            mse = mse.sum() / mask.sum().clamp(min=1.0)
-        else:
-            mse = mse.mean()
-
-        x0_hat = x_t.lerp(v_pred.detach(), t.unsqueeze(1))
-
-        nb = torch.tensor(0.0, device=self.device)
-        if self._cfg.nb_weight > 0:
-            gids = self._gene_ids.unsqueeze(0).expand(B, -1)
-            log_mu, theta = self.nb_head(self.backbone(gids, z))
-            nb = self.nb_head.nll(log_mu, theta, x, mask)
-
-        corr = torch.tensor(0.0, device=self.device)
-        if self._cfg.corr_weight > 0 and mask is None:
-            corr = self._corr(x0_hat, x)
-
-        return (
-            mse
-            + self._cfg.nb_weight * nb
-            + self._cfg.corr_weight * corr
-        )
-
-    @staticmethod
-    def _corr(p: torch.Tensor, q: torch.Tensor) -> torch.Tensor:
-        p = p - p.mean(1, keepdim=True)
-        q = q - q.mean(1, keepdim=True)
-        denom = p.norm(dim=1) * q.norm(dim=1) + 1e-8
-        return (1 - (p * q).sum(dim=1) / denom).mean()
-
-    # ----------------------------------------------------------------
-    # fit
-    # ----------------------------------------------------------------
-    def fit(
-        self,
-        X: Array,
-        z: Array,
-        *,
-        X_val: Optional[Array] = None,
-        z_val: Optional[Array] = None,
-        epochs: int = 200,
-        batch_size: int = 64,
-        lr: float = 2e-4,
-        weight_decay: float = 0.02,
-        patience: int = 20,
-        num_workers: int = 0,
-        verbose: int = 1,
-    ) -> "GeneFlowDecoder":
-        """
-        Train the decoder.
-
-        Parameters
-        ----------
-        X : array, shape ``(N, n_genes)``
-            Log1p-normalized expression matrix.
-        z : array, shape ``(N, z_dim)``
-            Cell-state vectors. **Never updated** when ``freeze_z=True``.
-        X_val, z_val : arrays, optional
-            Validation split. If given, early stopping is enabled and the
-            best weights (lowest validation loss) are restored at the end.
-        epochs, batch_size, lr, weight_decay, patience
-            Standard optimisation knobs.
-        num_workers : int
-            Workers for the training DataLoader.
-        verbose : int
-            0 = silent, 1 = epoch logs.
-
-        Returns
-        -------
-        self
-        """
-        X = _as_tensor(X, self.device)
-        z = self._maybe_detach(_as_tensor(z, self.device), snapshot=True)
-
-        # pin_memory 要求 CPU tensor；若 X/z 已在 GPU（常见于 scVI / mudata
-        # 把 .X 或 obsm 留在 cuda 的情况），强制关闭，避免
-        #   RuntimeError: cannot pin 'torch.cuda.FloatTensor'
-        # 注：_as_tensor 对 cuda 输入原样返回，因此这里 X/z 仍在 GPU。
-        pin_memory = _want_pin_memory(X, z)
-        assert X.device.type == "cpu" or z.device.type == "cpu" or not pin_memory
-        
-        ds = torch.utils.data.TensorDataset(X, z)
-        dl = torch.utils.data.DataLoader(
-            ds,
-            batch_size=int(batch_size),
-            shuffle=True,
-            num_workers=int(num_workers),
-            pin_memory=pin_memory,
-        )
-
-        has_val = False
-        if X_val is not None and z_val is not None:
-            Xv = _as_tensor(X_val, self.device)
-            zv = self._maybe_detach(_as_tensor(z_val, self.device), snapshot=True)
-            has_val = True
-
-        # Only decoder-owned parameters are optimised.
-        params = [p for p in self.parameters() if p.requires_grad]
-        opt = torch.optim.AdamW(params, lr=lr, weight_decay=weight_decay)
-        sched = torch.optim.lr_scheduler.CosineAnnealingLR(
-            opt, T_max=int(epochs)
-        )
-
-        best_val = float("inf")
-        best_state: Optional[dict[str, torch.Tensor]] = None
-        no_improve = 0
-
-        for ep in range(1, int(epochs) + 1):
-            self.train()
-            tot = 0.0
-            for xb, zb in dl:
-                opt.zero_grad()
-                loss = self._loss(xb, zb)
-                loss.backward()
-                nn.utils.clip_grad_norm_(params, 1.0)
-                opt.step()
-                tot += loss.item() * xb.shape[0]
-            sched.step()
-
-            if verbose:
-                msg = (
-                    f"[fit] epoch {ep:3d}/{epochs} "
-                    f"| train={tot / len(ds):.4f}"
-                )
-                if has_val:
-                    self.eval()
-                    with torch.no_grad():
-                        vl = self._loss(Xv, zv).item()
-                    msg += f" | val={vl:.4f}"
-                    if vl < best_val - 1e-5:
-                        best_val, no_improve = vl, 0
-                        best_state = {
-                            k: v.detach().cpu().clone()
-                            for k, v in self.state_dict().items()
-                        }
-                    else:
-                        no_improve += 1
-                    if no_improve >= int(patience):
-                        if verbose:
-                            print(f"[fit] early stop at epoch {ep}")
-                        break
-                print(msg)
-
-        if best_state is not None:
-            self.load_state_dict(best_state)
-            if verbose:
-                print(f"[fit] restored best val={best_val:.4f}")
-        return self
-
-    # ----------------------------------------------------------------
-    # inference
-    # ----------------------------------------------------------------
-    @torch.no_grad()
-    def forward(
-        self,
-        z: Array,
-        *,
-        steps: int = 4,
-        method: str = "midpoint",
-        deterministic: bool = False,
-    ) -> torch.Tensor:
-        """
-        Sample an expression profile from *z* (log1p space).
-
-        Parameters
-        ----------
-        z : array, shape ``(B, z_dim)``
-        steps : int
-            Number of flow-integration steps (4 is usually enough).
-        method : ``"midpoint"`` | ``"euler"``
-        deterministic : bool
-            If ``True``, returns the deterministic NB mean
-            (equivalent to ``self.predict(z)``).
-        """
-        z = self._maybe_detach(_as_tensor(z, self.device))
-        if deterministic:
-            return self.predict(z)
-
-        self.eval()
-        B = z.shape[0]
-        x = torch.randn(B, self.n_genes, device=self.device)
-        ts = torch.linspace(0.0, 1.0, int(steps) + 1, device=self.device)
-        for i in range(int(steps)):
-            t0, t1 = ts[i], ts[i + 1]
-            tmid = (t0 + t1) / 2
-            dt = (t1 - t0).item()
-            v = self._velocity(x, torch.full((B,), t0, device=self.device), z)
-            if method == "midpoint":
-                xm = x + 0.5 * dt * v
-                v = self._velocity(
-                    xm, torch.full((B,), tmid, device=self.device), z
-                )
-            x = x + dt * v
-        return x
-
-    @torch.no_grad()
-    def predict(self, z: Array) -> torch.Tensor:
-        """
-        Deterministic reconstruction: the NB mean in log1p space.
-
-        Convert to pseudo-counts with ``(log_mu.exp() - 1).clamp(min=0)``.
-        """
-        z = self._maybe_detach(_as_tensor(z, self.device))
-        self.eval()
-        B = z.shape[0]
-        gids = self._gene_ids.unsqueeze(0).expand(B, -1)
-        log_mu, _ = self.nb_head(self.backbone(gids, z))
-        return log_mu
-
-    # ----------------------------------------------------------------
-    # persistence
-    # ----------------------------------------------------------------
-    def save(self, path: str) -> None:
-        """Save model weights + config to ``path``."""
-        torch.save(
-            {
-                "state": self.state_dict(),
-                "cfg": self._cfg,
-                "freeze_z": self.freeze_z,
-                "n_genes": self.n_genes,
-                "z_dim": self.z_dim,
-            },
-            path,
-        )
-
-    @classmethod
-    def load(
-        cls,
-        path: str,
-        *,
-        device: str = "auto",
-        **init_kwargs,
-    ) -> "GeneFlowDecoder":
-        """
-        Load a saved model.
-
-        ``n_genes`` / ``z_dim`` / ``freeze_z`` are read from the checkpoint
-        unless explicitly overridden in ``init_kwargs``.
-        """
-        ckpt = torch.load(path, map_location="cpu")
-        init_kwargs.setdefault("n_genes", ckpt.get("n_genes"))
-        init_kwargs.setdefault("z_dim", ckpt.get("z_dim"))
-        init_kwargs.setdefault("freeze_z", ckpt.get("freeze_z", True))
-        init_kwargs["device"] = device
-        model = cls(**init_kwargs)
-        model._cfg = ckpt.get("cfg", model._cfg)
-        model.load_state_dict(ckpt["state"])
-        return model