q and dtype

Author: evanatyourservice

kron_torch/kron.py

@@ -376,8 +376,10 @@ def _solve_triangular_right(X, A):
     orig_dtype = X.dtype
     X = X.to(dtype=torch.float32, non_blocking=True)
     A = A.to(dtype=torch.float32, non_blocking=True)
-    return torch.linalg.solve_triangular(A, X.reshape(-1, q.size(0)), upper=True, left=False).reshape_as(X)
-    
+    out = torch.linalg.solve_triangular(
+        A, X.reshape(-1, X.size(-1)), upper=True, left=False
+    ).reshape_as(X)
+    return out.to(dtype=orig_dtype, non_blocking=True)
 
 @torch.compile(fullgraph=True, dynamic=False)
 def _calc_A_and_conjB(exprA, G, Q, V):

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "kron-torch"
-version = "0.2.8"
+version = "0.2.9"
 description = "An implementation of PSGD Kron optimizer in PyTorch."
 readme = { file = "README.md", content-type = "text/markdown" }
 license = { file = "LICENSE" }