From 8304029ab34f01030e48d079d765043a8b2c8e07 Mon Sep 17 00:00:00 2001 From: Gertjan van Zwieten Date: Fri, 28 Jul 2023 14:52:45 +0200 Subject: [PATCH] support off-diagonal trace This patch reimplements the dispatched numpy.trace function to include support for the off-diagonal case. An important side effect of this patch is that the function no longer relies on the function._takediag helper function, which will be removed in a later commit. --- nutils/function.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/nutils/function.py b/nutils/function.py index 6ca383680..3f0dc8e55 100644 --- a/nutils/function.py +++ b/nutils/function.py @@ -4353,9 +4353,7 @@ def ravel(arg: Array): @implements(numpy.trace) def trace(arg: Array, offset: int = 0, axis1: int = 0, axis2: int = 1) -> Array: - if offset != 0: - raise NotImplementedError('traces over offset diagonal are not yet supported') - return numpy.sum(_takediag(arg, axis1, axis2), -1) + return numpy.sum(numpy.diagonal(arg, offset, axis1, axis2), -1) @implements(numpy.transpose) def transpose(array: Array, axes: Optional[Sequence[int]] = None) -> Array: