geneflowdecoder 0.1.2__tar.gz → 0.1.4__tar.gz

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

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

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

@@ -1,6 +1,6 @@
 [project]
 name = "geneflowdecoder"
-version = "0.1.2"
+version = "0.1.4"
 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.2 → geneflowdecoder-0.1.4}/src/geneflowdecoder/_core.py

@@ -388,15 +388,25 @@ class GeneFlowDecoder(nn.Module):
         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:
-        z = self._maybe_detach(z)
-        B = x_t.shape[0]
-        t_emb = self.time_mlp(self._sincos(t)).unsqueeze(1)
-        gids = self._gene_ids.unsqueeze(0).expand(B, -1)
-        h = self.backbone(gids, z) + t_emb
-        return self.vel_proj(h)
+    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

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

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

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

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