Upload fractus/nn/stats.py with huggingface_hub

fractus/nn/stats.py

@@ -0,0 +1,45 @@
+"""Numerical utilities for fractus.
+
+Ported from the original system (src/math/stats.rs) in pure PyTorch, differentiable.
+
+elu_plus_one : strictly positive feature map for linear attention.
+    φ(x, α) = x + 1              if x > 0
+            = α(e^x - 1) + 1     otherwise
+    With α=1 (default), φ is strictly positive (min e^x > 0 for x→-∞,
+    = 1 at x=0). This positivity guarantees that the denominator of causal
+    linear attention stays well-defined.
+
+stable_softmax : softmax with max subtraction (no overflow).
+"""
+
+import torch
+
+
+def elu_plus_one(x: torch.Tensor, alpha: float = 1.0) -> torch.Tensor:
+    """ELU+1 strictly positive feature map, differentiable.
+
+    Args:
+        x : tensor of arbitrary shape.
+        alpha : ELU coefficient (1.0 by default, as in the original).
+    Returns:
+        tensor of the same shape, strictly positive.
+    """
+    # We use the direct formula (differentiable via torch.where):
+    # positive branch: x + 1; negative branch: alpha * (exp(x) - 1) + 1.
+    pos = x + 1.0
+    neg = alpha * (torch.exp(x) - 1.0) + 1.0
+    return torch.where(x > 0, pos, neg)
+
+
+def stable_softmax(logits: torch.Tensor, dim: int = -1) -> torch.Tensor:
+    """Numerically stable softmax (max subtraction).
+
+    If the exponential sum is < 1e-10, returns the uniform 1/N
+    (limit behavior inherited from the original stats.rs:56-57).
+    """
+    max_logits, _ = logits.max(dim=dim, keepdim=True)
+    exp = torch.exp(logits - max_logits)
+    denom = exp.sum(dim=dim, keepdim=True)
+    # Limit behavior: uniform if denom ~ 0.
+    uniform = torch.full_like(exp, 1.0 / exp.shape[dim])
+    return torch.where(denom > 1e-10, exp / denom, uniform)