Merge pull request #24 from dwierichs/main

Allow additional keyword arguments in `tensorflow_diag`
jcmgray · May 14, 2024 · b367fa8 · b367fa8
2 parents 96b42b7 + cfa94de
commit b367fa8
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/autoray/autoray.py b/autoray/autoray.py
@@ -1871,12 +1871,12 @@ def numpy_like(ary, indices_or_sections, axis=0, **kwargs):
     return numpy_like
 
 
-def tensorflow_diag(x):
+def tensorflow_diag(x, **kwargs):
     nd = ndim(x)
     if nd == 2:
-        return do("linalg.diag_part", x)
+        return do("linalg.diag_part", x, **kwargs)
     elif nd == 1:
-        return do("linalg.diag", x)
+        return do("linalg.diag", x, **kwargs)
     else:
         raise ValueError("Input must be 1- or 2-d.")