geneflowdecoder 0.1.0__tar.gz → 0.1.2__tar.gz

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

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

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

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

@@ -28,31 +28,121 @@ try:
     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
-    torch = None  # type: ignore
-    nn = object   # type: ignore
-    F = object    # type: ignore
-
-    def _no_grad(fn):  # identity decorator when torch is absent
-        return fn
-
+    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
 
-try:
-    from torch import Tensor as _TorchTensor  # noqa: F401
-    Tensor = _TorchTensor  # type: ignore
-except Exception:
-    pass
+Array = Union[np.ndarray, Tensor]  # type: ignore[valid-type]
 
-Array = Union[np.ndarray, torch.Tensor]
 
+def _as_tensor(x: Array, device: "torch.device") -> "torch.Tensor":
+    """把 numpy / CPU / GPU 输入规范化成 torch 张量。
 
-def _as_tensor(x: Array, device: torch.device) -> torch.Tensor:
+    关键：输入**已经**在 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):
-        x = torch.from_numpy(x)
+        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
 # ----------------------------------------------------------------------
@@ -244,6 +334,11 @@ class GeneFlowDecoder(nn.Module):
         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)
@@ -398,13 +493,20 @@ class GeneFlowDecoder(nn.Module):
         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=True,
+            pin_memory=pin_memory,
         )
 
         has_val = False

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

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

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

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