Array API compatible wrapper for :py:func:`np.trace `. See its docstring for more information.
(x: Array, /, *, offset: int = 0, dtype: Optional[Dtype] = None)
| 376 | |
| 377 | # Note: trace is the numpy top-level namespace, not np.linalg |
| 378 | def trace(x: Array, /, *, offset: int = 0, dtype: Optional[Dtype] = None) -> Array: |
| 379 | """ |
| 380 | Array API compatible wrapper for :py:func:`np.trace <numpy.trace>`. |
| 381 | |
| 382 | See its docstring for more information. |
| 383 | """ |
| 384 | if x.dtype not in _numeric_dtypes: |
| 385 | raise TypeError('Only numeric dtypes are allowed in trace') |
| 386 | |
| 387 | # Note: trace() works the same as sum() and prod() (see |
| 388 | # _statistical_functions.py) |
| 389 | if dtype is None: |
| 390 | if x.dtype == float32: |
| 391 | dtype = float64 |
| 392 | elif x.dtype == complex64: |
| 393 | dtype = complex128 |
| 394 | # Note: trace always operates on the last two axes, whereas np.trace |
| 395 | # operates on the first two axes by default |
| 396 | return Array._new(np.asarray(np.trace(x._array, offset=offset, axis1=-2, axis2=-1, dtype=dtype))) |
| 397 | |
| 398 | # Note: vecdot is not in NumPy |
| 399 | def vecdot(x1: Array, x2: Array, /, *, axis: int = -1) -> Array: |