Mixed-precision reconstruction, fused-kernel dispatch, and CUDA graph capture for tomography INR

src/quantem/core/config.py

@@ -34,6 +34,15 @@
     if "cuda" in str(e):
         NUM_DEVICES = 0
     _defaults["has_cupy"] = False
+try:
+    import quantem.cuda  # type: ignore  # noqa: F401
+
+    _defaults["has_quantem_cuda"] = True
+except ModuleNotFoundError:
+    _defaults["has_quantem_cuda"] = False
+except Exception:
+    # installed but unloadable (e.g. libcudart missing at runtime)
+    _defaults["has_quantem_cuda"] = False
 
 
 defaults: list[Mapping] = [_defaults]

src/quantem/core/ml/activation_functions.py

@@ -1,10 +1,66 @@
-from typing import Callable
+import os
+from typing import Any, Callable
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
 
+class _TruncExpReference(torch.autograd.Function):
+    """Exponential with the torch-ngp clamped backward used by tomography."""
+
+    @staticmethod
+    def forward(ctx: Any, values: torch.Tensor, offset: float) -> torch.Tensor:
+        ctx.save_for_backward(values)
+        ctx.offset = offset
+        return torch.exp(values - offset)
+
+    @staticmethod
+    def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor, None]:
+        (values,) = ctx.saved_tensors
+        shifted = values - ctx.offset
+        return grad_output * torch.exp(shifted.clamp(max=15)), None
+
+
+def trunc_exp(values: torch.Tensor, offset: float = 0.0) -> torch.Tensor:
+    """Apply trunc-exp, using the fused CUDA density tail when available."""
+    use_fused = (
+        os.environ.get("QUANTEM_DENSITY_TAIL_FUSED", "1") != "0"
+        and values.is_cuda
+        and values.dtype in (torch.float32, torch.bfloat16)
+        and values.ndim in (1, 2)
+        and values.numel() > 0
+    )
+    if use_fused:
+        try:
+            import quantem.cuda.core.ml as cuda_ml
+        except (ImportError, OSError, RuntimeError):
+            pass
+        else:
+            fused = getattr(cuda_ml, "density_tail", None)
+            if fused is not None:
+                return fused(values, float(offset))
+    return _TruncExpReference.apply(values, float(offset))
+
+
+# Object constraints consume this explicit capability instead of guessing from
+# a callable name or implementation detail.
+trunc_exp.quantem_guarantees_nonnegative = True  # type: ignore[attr-defined]
+
+
+class TruncExpActivation(nn.Module):
+    """Configurable trunc-exp activation with a non-negativity capability."""
+
+    quantem_guarantees_nonnegative = True
+
+    def __init__(self, offset: float = 0.0) -> None:
+        super().__init__()
+        self.offset = float(offset)
+
+    def forward(self, values: torch.Tensor) -> torch.Tensor:
+        return trunc_exp(values, self.offset)
+
+
 class ModReLU(nn.Module):
     """Modulated ReLU activation for complex-valued inputs.
 

src/quantem/core/ml/inr.py

@@ -111,18 +111,20 @@ def _build(self) -> None:
         self.net = nn.Sequential(*net_list)
 
         if self.winner_initialization:
-            if type(self.winner_initialization) is int:
-                rng = torch.Generator()
-                rng.manual_seed(self.winner_initialization)
-            else:
-                rng = torch.Generator()
-                rng.manual_seed(42)
+            seed = self.winner_initialization if type(self.winner_initialization) is int else 42
+            rng = torch.Generator()
+            rng.manual_seed(seed)
+            # torch.randn_like ignores generators, so the noise must come from
+            # torch.randn with the seeded generator -- otherwise the "winner" seed
+            # silently has no effect and the perturbation is not reproducible.
             with torch.no_grad():
-                self.net[0].linear.weight += (  # type: ignore[reportAttributeAccessIssue]
-                    torch.randn_like(self.net[0].linear.weight) * 5 / self.first_omega_0  # type:ignore
-                )
-                self.net[1].linear.weight += (  # type: ignore[reportAttributeAccessIssue]
-                    torch.randn_like(self.net[1].linear.weight) * 0.1 / self.hidden_omega_0  # type:ignore
+                w0 = self.net[0].linear.weight  # type: ignore[reportAttributeAccessIssue]
+                w0 += torch.randn(w0.shape, generator=rng, dtype=w0.dtype) * 5 / self.first_omega_0
+                w1 = self.net[1].linear.weight  # type: ignore[reportAttributeAccessIssue]
+                w1 += (
+                    torch.randn(w1.shape, generator=rng, dtype=w1.dtype)
+                    * 0.1
+                    / self.hidden_omega_0
                 )
 
     def forward(self, coords: torch.Tensor) -> torch.Tensor: