added missing diag rules and fix to unary rule

cola/linalg/diag_trace.py

@@ -40,6 +40,12 @@ def diag(A: Identity, k=0, **kwargs):
     else:
         return A.xnp.zeros(A.shape[0] - k, A.dtype)
 
+@dispatch
+def diag(A: Diagonal, k=0, **kwargs):
+    if k == 0:
+        return A.v
+    else:
+        return A.xnp.zeros(A.shape[0] - k, A.dtype)
 
 @dispatch
 def diag(A: Sum, k=0, **kwargs):

cola/linalg/logdet.py

@@ -73,7 +73,8 @@ def slogdet(A: LinearOperator, **kwargs) -> Array:
         return slogdet(cola.decompositions.cholesky_decomposed(A), **kws)
     elif method in ('iterative', 'approx') or (method == 'auto' and
                                                (np.prod(A.shape) > 1e6 and kws['tol'] >= 3e-2)):
-        return 1., stochastic_lanczos_quad(A, A.xnp.log, **kws)
+        one = A.xnp.array(1., dtype=A.dtype, device=A.device)
+        return one, stochastic_lanczos_quad(A, A.xnp.log, **kws)
     else:
         raise ValueError(f"Unknown method {method} or CoLA didn't fit any selection criteria")
 

cola/linalg/unary.py

@@ -115,12 +115,13 @@ def apply_unary(f: Callable, A: Diagonal, **kwargs):
 
 @dispatch
 def apply_unary(f: Callable, A: BlockDiag, **kwargs):
-    fAs = [apply_unary(f, a, **kwargs) for a in A.blocks]
+    fAs = [apply_unary(f, a, **kwargs) for a in A.Ms]
     return BlockDiag(*fAs, multiplicities=A.multiplicities)
 
 @dispatch
 def apply_unary(f: Callable, A: Identity, **kwargs):
-    return f(1.)*A
+    one = A.xnp.array(1., dtype=A.dtype, device=A.device)
+    return f(one)*A
 
 @dispatch
 def apply_unary(f: Callable, A: ScalarMul, **kwargs):