geneflowdecoder 0.1.1__tar.gz → 0.1.3__tar.gz

{geneflowdecoder-0.1.1/src/geneflowdecoder.egg-info → geneflowdecoder-0.1.3}/PKG-INFO

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

{geneflowdecoder-0.1.1 → geneflowdecoder-0.1.3}/pyproject.toml

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

@@ -388,15 +388,17 @@ 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)
+    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
-        return self.vel_proj(h)
+        v = self.vel_proj(h)                                   # [B,G,1] 或 [B,G]
+        v = v.reshape(B, G)                                    # ← 强制对齐，根治错位
+        return v
 
     # ----------------------------------------------------------------
     # objective
@@ -497,9 +499,9 @@ class GeneFlowDecoder(nn.Module):
         # 把 .X 或 obsm 留在 cuda 的情况），强制关闭，避免
         #   RuntimeError: cannot pin 'torch.cuda.FloatTensor'
         # 注：_as_tensor 对 cuda 输入原样返回，因此这里 X/z 仍在 GPU。
-        assert X.device.type == "cpu" or z.device.type == "cpu" or not pin_memory
         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,

{geneflowdecoder-0.1.1 → geneflowdecoder-0.1.3}/src/geneflowdecoder/_version.py

@@ -1,2 +1,2 @@
 # -*- coding: utf-8 -*-
-__version__ = "0.1.1"
+__version__ = "0.1.3"

{geneflowdecoder-0.1.1 → geneflowdecoder-0.1.3/src/geneflowdecoder.egg-info}/PKG-INFO

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