DOC: create_diagonal: add docs

Author: lucascolley

docs/api-reference.md

@@ -8,6 +8,7 @@
 
     atleast_nd
     cov
+    create_diagonal
     expand_dims
     kron
 ```

src/array_api_extra/_funcs.py

@@ -6,7 +6,7 @@
 if TYPE_CHECKING:
     from ._typing import Array, ModuleType
 
-__all__ = ["atleast_nd", "cov", "expand_dims", "kron"]
+__all__ = ["atleast_nd", "cov", "create_diagonal", "expand_dims", "kron"]
 
 
 def atleast_nd(x: Array, /, *, ndim: int, xp: ModuleType) -> Array:
@@ -141,6 +141,41 @@ def cov(m: Array, /, *, xp: ModuleType) -> Array:
 
 
 def create_diagonal(x: Array, /, *, offset: int = 0, xp: ModuleType) -> Array:
+    """
+    Construct a diagonal array.
+
+    Parameters
+    ----------
+    x : array
+        A 1-D array
+    offset : int, optional
+        Offset from the leading diagonal (default is ``0``).
+        Use positive ints for diagonals above the leading diagonal,
+        and negative ints for diagonals below the leading diagonal.
+
+    Returns
+    -------
+    res : array
+        A 2-D array with `x` on the diagonal (offset by `offset`).
+
+    Examples
+    --------
+    >>> import array_api_strict as xp
+    >>> import array_api_extra as xpx
+    >>> x = xp.asarray([0, 4, 8])
+
+    >>> xpx.create_diagonal(x, xp=xp)
+    Array([[-2.1,  0. ,  0. ],
+           [ 0. , -1. ,  0. ],
+           [ 0. ,  0. ,  4.3]], dtype=array_api_strict.float64)
+    >>> xpx.create_diagonal(x, offset=-2, xp=xp)
+    Array([[ 0. ,  0. ,  0. ,  0. ,  0. ],
+           [ 0. ,  0. ,  0. ,  0. ,  0. ],
+           [-2.1,  0. ,  0. ,  0. ,  0. ],
+           [ 0. , -1. ,  0. ,  0. ,  0. ],
+           [ 0. ,  0. ,  4.3,  0. ,  0. ]], dtype=array_api_strict.float64)
+
+    """
     if x.ndim != 1:
         err_msg = "`x` must be 1-dimensional."
         raise ValueError(err_msg)