Skip to content

Commit

Permalink
DOC: create_diagonal: add docs
Browse files Browse the repository at this point in the history
  • Loading branch information
lucascolley committed Nov 14, 2024
1 parent 0a53e57 commit be3664d
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 1 deletion.
1 change: 1 addition & 0 deletions docs/api-reference.md
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
atleast_nd
cov
create_diagonal
expand_dims
kron
```
38 changes: 37 additions & 1 deletion src/array_api_extra/_funcs.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -141,6 +141,42 @@ 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([2, 4, 8])
>>> xpx.create_diagonal(x, xp=xp)
Array([[2, 0, 0],
[0, 4, 0],
[0, 0, 8]], dtype=array_api_strict.int64)
>>> xpx.create_diagonal(x, offset=-2, xp=xp)
Array([[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0],
[2, 0, 0, 0, 0],
[0, 4, 0, 0, 0],
[0, 0, 8, 0, 0]], dtype=array_api_strict.int64)
"""
if x.ndim != 1:
err_msg = "`x` must be 1-dimensional."
raise ValueError(err_msg)
Expand Down

0 comments on commit be3664d

Please sign in to comment.