mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #1091 from fehiepsi/tril

Browse Source 
expose tril_indices, triu_indices similar to diag_indices
This commit is contained in:
Matthew Johnson 2019-08-01 20:58:22 -07:00 committed by GitHub
commit 1f3b4ae97e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
jax/numpy/lax_numpy.py
View File

@ -1615,6 +1615,8 @@ def trace(a, offset=0, axis1=0, axis2=1, dtype=None, out=None):
diag_indices = onp.diag_indices
tril_indices = onp.tril_indices
triu_indices = onp.triu_indices
@_wraps(onp.diagonal)