Skip to content

Commit

Permalink
ENH: add create_diagonal
Browse files Browse the repository at this point in the history
Co-authored-by: Jake Bowhay <[email protected]>
  • Loading branch information
lucascolley and j-bowhay committed Nov 14, 2024
1 parent b83ba61 commit 57e26a8
Showing 1 changed file with 8 additions and 0 deletions.
8 changes: 8 additions & 0 deletions src/array_api_extra/_funcs.py
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,14 @@ def cov(m: Array, /, *, xp: ModuleType) -> Array:
return xp.squeeze(c, axis=axes)


def create_diagonal(x: Array, /, *, offset: int = 0, xp: ModuleType) -> Array:
n = x.shape[0] + abs(offset)
diag = xp.zeros(n**2, dtype=x.dtype)
i = offset if offset >= 0 else abs(offset) * n
diag[i : min(n * (n - offset), diag.shape[0]) : n + 1] = x
return xp.reshape(diag, (n, n))


def _mean(
x: Array,
/,
Expand Down

0 comments on commit 57e26a8

Please sign in to comment.