add lazy diag

Author: jcmgray

autoray/lazy/__init__.py

@@ -10,6 +10,7 @@
     tensordot,
     einsum,
     trace,
+    diag,
     matmul,
     kron,
     clip,
@@ -69,6 +70,7 @@
     "einsum",
     "conj",
     "trace",
+    "diag",
     "matmul",
     "kron",
     "clip",

autoray/lazy/core.py

@@ -12,7 +12,6 @@
     get_dtype_name,
     register_function,
     astype,
-    complex_add_re_im,
 )
 
 
@@ -1119,6 +1118,24 @@ def trace(a):
     )
 
 
+@lazy_cache("diag")
+def diag(a):
+    a = ensure_lazy(a)
+
+    if a.ndim == 1:
+        new_shape = (a.size, a.size)
+    elif a.ndim == 2:
+        new_shape = (min(a.shape),)
+    else:
+        raise ValueError("Input must be 1- or 2-d.")
+
+    return a.to(
+        fn=get_lib_fn(a.backend, "diag"),
+        args=(a,),
+        shape=new_shape,
+    )
+
+
 @lazy_cache("matmul")
 def matmul(x1, x2):
     backend = find_common_backend(x1, x2)

tests/test_lazy.py

@@ -379,6 +379,20 @@ def test_indexing():
         assert_allclose(a[key], b[key].compute())
 
 
+@pytest.mark.parametrize('shape', [
+    (3,),
+    (2, 2),
+    (3, 4),
+    (4, 3),
+])
+def test_diag(shape):
+    a = do('random.uniform', size=shape, like='numpy')
+    b = lazy.array(a)
+    ad = do('diag', a)
+    bd = do('diag', b)
+    assert_allclose(ad, bd.compute())
+
+
 def test_einsum():
     a = do('random.uniform', size=(2, 3, 4, 5), like='numpy')
     b = do('random.uniform', size=(4, 5), like='numpy')