Modernize APIs, improve numerical stability, bump to 0.5.0

- kronecker_product: replace manual implementation with torch.kron
- GELS: replace deprecated torch.cholesky with torch.linalg.cholesky
- Lookahead optimizer: fix deprecated add_(scalar, tensor) signature
- ls_cov: use torch.linalg.lstsq for params, torch.linalg.solve
  instead of explicit .inverse() for better numerical stability
- StandardScaler: precompute reciprocal to multiply instead of divide
- Bump version 0.4.3 → 0.5.0

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: LemonPi

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "arm_pytorch_utilities"
-version = "0.4.3"
+version = "0.5.0"
 description = "Utilities for working with pytorch"
 readme = "README.md" # Optional
 

src/arm_pytorch_utilities/linalg.py

@@ -42,20 +42,7 @@ def kronecker_product(t1, t2):
     Computes the Kronecker product between two tensors.
     See https://en.wikipedia.org/wiki/Kronecker_product
     """
-    t1_height, t1_width = t1.size()
-    t2_height, t2_width = t2.size()
-    out_height = t1_height * t2_height
-    out_width = t1_width * t2_width
-
-    tiled_t2 = t2.repeat(t1_height, t1_width)
-    expanded_t1 = (
-        t1.unsqueeze(2)
-            .unsqueeze(3)
-            .repeat(1, t2_height, t2_width, 1)
-            .view(out_height, out_width)
-    )
-
-    return expanded_t1 * tiled_t2
+    return torch.kron(t1, t2)
 
 
 def cov(x, rowvar=False, bias=False, ddof=None, aweights=None):
@@ -152,16 +139,16 @@ def forward(ctx, A, b):
         # A: (..., M, N)
         # b: (..., M, K)
         # https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/ops/linalg_ops.py#L267
-        u = torch.cholesky(torch.matmul(A.transpose(-1, -2), A), upper=True)
-        ret = torch.cholesky_solve(torch.matmul(A.transpose(-1, -2), b), u, upper=True)
-        ctx.save_for_backward(u, ret, A, b)
+        L = torch.linalg.cholesky(torch.matmul(A.transpose(-1, -2), A))
+        ret = torch.cholesky_solve(torch.matmul(A.transpose(-1, -2), b), L, upper=False)
+        ctx.save_for_backward(L, ret, A, b)
         return ret
 
     @staticmethod
     def backward(ctx, grad_output):
         # https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/ops/linalg_grad.py#L223
         chol, x, a, b = ctx.saved_tensors
-        z = torch.cholesky_solve(grad_output, chol, upper=True)
+        z = torch.cholesky_solve(grad_output, chol, upper=False)
         xzt = torch.matmul(x, z.transpose(-1, -2))
         zx_sym = xzt + xzt.transpose(-1, -2)
         grad_A = - torch.matmul(a, zx_sym) + torch.matmul(b, z.transpose(-1, -2))
@@ -197,11 +184,13 @@ def ls(X, Y, weights=None):
 def ls_cov(X, Y, weights=None, make_symmetric=True, sigreg=1e-4):
     X, Y = _apply_weights(X, Y, weights)
 
-    pinvXX = X.pinverse()
-    params = (pinvXX @ Y).t()
+    # Solve least squares via lstsq (more stable than pinverse)
+    result = torch.linalg.lstsq(X, Y)
+    params = result.solution.t()
 
     # estimate covariance according to: http://users.stat.umn.edu/~helwig/notes/mvlr-Notes.pdf (see up to slide 66)
     # hat/projection matrix - Yhat = H*Y
+    pinvXX = X.pinverse()
     H = X @ pinvXX
 
     N = X.shape[0]
@@ -231,8 +220,9 @@ def ls_cov(X, Y, weights=None, make_symmetric=True, sigreg=1e-4):
         XXXX = XXXX_sym
         error_covariance = error_covariance_sym
 
-    # TODO might be able to use cholesky decomp here since XXXX > 0
-    covariance = kronecker_product(error_covariance, XXXX.inverse())
+    # Use solve instead of explicit inverse: solve(A, I) = A^{-1}, more numerically stable
+    XXXX_inv = torch.linalg.solve(XXXX, torch.eye(XXXX.shape[0], dtype=XXXX.dtype, device=XXXX.device)).contiguous()
+    covariance = kronecker_product(error_covariance, XXXX_inv)
 
     return params, covariance
 

src/arm_pytorch_utilities/optim.py

@@ -106,12 +106,12 @@ def step(self, closure=None):
             for group in self.optimizer.param_groups:
                 for p in group['params']:
                     param_state = self.state[p]
-                    p.data.mul_(self.la_alpha).add_(1.0 - self.la_alpha, param_state['cached_params'])  # crucial line
+                    p.data.mul_(self.la_alpha).add_(param_state['cached_params'], alpha=1.0 - self.la_alpha)  # crucial line
                     param_state['cached_params'].copy_(p.data)
                     if self.pullback_momentum == "pullback":
                         internal_momentum = self.optimizer.state[p]["momentum_buffer"]
                         self.optimizer.state[p]["momentum_buffer"] = internal_momentum.mul_(self.la_alpha).add_(
-                            1.0 - self.la_alpha, param_state["cached_mom"])
+                            param_state["cached_mom"], alpha=1.0 - self.la_alpha)
                         param_state["cached_mom"] = self.optimizer.state[p]["momentum_buffer"]
                     elif self.pullback_momentum == "reset":
                         self.optimizer.state[p]["momentum_buffer"] = torch.zeros_like(p.data)

src/arm_pytorch_utilities/preprocess.py

@@ -311,13 +311,15 @@ def __init__(self, **kwargs):
         super().__init__(**kwargs)
         self._m = None
         self._s = None
+        self._inv_s = None
 
     def fit(self, X):
         self._m = X.mean(0, keepdim=True)
         self._s = X.std(0, unbiased=False, keepdim=True)
+        self._inv_s = 1.0 / self._s
 
     def transform(self, X):
-        return (X - self._m) / self._s
+        return (X - self._m) * self._inv_s
 
     def inverse_transform(self, X):
         return (X * self._s) + self._m