From 31b9a0c501b6bfb0b10851f801756114d7161a94 Mon Sep 17 00:00:00 2001
From: Athan Reines <kgryte@gmail.com>
Date: Thu, 1 Dec 2022 01:29:31 -0800
Subject: [PATCH] Add complex number support to `linalg.trace`
---
spec/API_specification/array_api/linalg.py | 9 ++++-----
1 file changed, 4 insertions(+), 5 deletions(-)
diff --git a/spec/API_specification/array_api/linalg.py b/spec/API_specification/array_api/linalg.py
index b3595e1fa..619037c0c 100644
--- a/spec/API_specification/array_api/linalg.py
+++ b/spec/API_specification/array_api/linalg.py
@@ -447,14 +447,12 @@ def trace(x: array, /, *, offset: int = 0, dtype: Optional[dtype] = None) -> arr
- If ``N`` is ``0``, the sum is ``0`` (i.e., the empty sum).
- For floating-point operands,
-
- - If ``x_i`` is ``NaN``, the sum is ``NaN`` (i.e., ``NaN`` values propagate).
+ For both real-valued and complex floating-point operands, special cases must be handled as if the operation is implemented by successive application of :func:`~array_api.add`.
Parameters
----------
x: array
- input array having shape ``(..., M, N)`` and whose innermost two dimensions form ``MxN`` matrices. Should have a real-valued data type.
+ input array having shape ``(..., M, N)`` and whose innermost two dimensions form ``MxN`` matrices. Should have a numeric data type.
offset: int
offset specifying the off-diagonal relative to the main diagonal.
@@ -466,8 +464,9 @@ def trace(x: array, /, *, offset: int = 0, dtype: Optional[dtype] = None) -> arr
dtype: Optional[dtype]
data type of the returned array. If ``None``,
- - if the default data type corresponding to the data type "kind" (integer or floating-point) of ``x`` has a smaller range of values than the data type of ``x`` (e.g., ``x`` has data type ``int64`` and the default data type is ``int32``, or ``x`` has data type ``uint64`` and the default data type is ``int64``), the returned array must have the same data type as ``x``.
+ - if the default data type corresponding to the data type "kind" (integer, real-valued floating-point, or complex floating-point) of ``x`` has a smaller range of values than the data type of ``x`` (e.g., ``x`` has data type ``int64`` and the default data type is ``int32``, or ``x`` has data type ``uint64`` and the default data type is ``int64``), the returned array must have the same data type as ``x``.
- if ``x`` has a real-valued floating-point data type, the returned array must have the default real-valued floating-point data type.
+ - if ``x`` has a complex floating-point data type, the returned array must have the default complex floating-point data type.
- if ``x`` has a signed integer data type (e.g., ``int16``), the returned array must have the default integer data type.
- if ``x`` has an unsigned integer data type (e.g., ``uint16``), the returned array must have an unsigned integer data type having the same number of bits as the default integer data type (e.g., if the default integer data type is ``int32``, the returned array must have a ``uint32`` data type).