Add complex number support to equal (#528)

Author: kgryte

spec/API_specification/array_api/array_object.py

@@ -369,9 +369,31 @@ def __dlpack_device__(self: array, /) -> Tuple[Enum, int]:
         """
 
     def __eq__(self: array, other: Union[int, float, bool, array], /) -> array:
-        """
+        r"""
         Computes the truth value of ``self_i == other_i`` for each element of an array instance with the respective element of the array ``other``.
 
+        **Special Cases**
+
+        Let ``self`` equal ``x1`` and ``other`` equal ``x2``.
+
+        For real-valued floating-point operands,
+
+        - If ``x1_i`` is ``NaN`` or ``x2_i`` is ``NaN``, the result is ``False``.
+        - If ``x1_i`` is ``+infinity`` and ``x2_i`` is ``+infinity``, the result is ``True``.
+        - If ``x1_i`` is ``-infinity`` and ``x2_i`` is ``-infinity``, the result is ``True``.
+        - If ``x1_i`` is ``-0`` and ``x2_i`` is either ``+0`` or ``-0``, the result is ``True``.
+        - If ``x1_i`` is ``+0`` and ``x2_i`` is either ``+0`` or ``-0``, the result is ``True``.
+        - If ``x1_i`` is a finite number, ``x2_i`` is a finite number, and ``x1_i`` equals ``x2_i``, the result is ``True``.
+        - In the remaining cases, the result is ``False``.
+
+        For complex floating-point operands, let ``a = real(x1_i)``, ``b = imag(x1_i)``, ``c = real(x2_i)``, ``d = imag(x2_i)``, and
+
+        - If ``a``, ``b``, ``c``, or ``d`` is ``NaN``, the result is ``False``.
+        - In the remaining cases, the result is the logical AND of the equality comparison between the real values ``a`` and ``c`` (real components) and between the real values ``b`` and ``d`` (imaginary components), as described above for real-valued floating-point operands (i.e., ``a == c AND b == d``).
+
+        .. note::
+           For discussion of complex number equality, see :ref:`complex-numbers`.
+
         Parameters
         ----------
         self: array

spec/API_specification/array_api/elementwise_functions.py

@@ -812,9 +812,29 @@ def divide(x1: array, x2: array, /) -> array:
     """
 
 def equal(x1: array, x2: array, /) -> array:
-    """
+    r"""
     Computes the truth value of ``x1_i == x2_i`` for each element ``x1_i`` of the input array ``x1`` with the respective element ``x2_i`` of the input array ``x2``.
 
+    **Special Cases**
+
+    For real-valued floating-point operands,
+
+    - If ``x1_i`` is ``NaN`` or ``x2_i`` is ``NaN``, the result is ``False``.
+    - If ``x1_i`` is ``+infinity`` and ``x2_i`` is ``+infinity``, the result is ``True``.
+    - If ``x1_i`` is ``-infinity`` and ``x2_i`` is ``-infinity``, the result is ``True``.
+    - If ``x1_i`` is ``-0`` and ``x2_i`` is either ``+0`` or ``-0``, the result is ``True``.
+    - If ``x1_i`` is ``+0`` and ``x2_i`` is either ``+0`` or ``-0``, the result is ``True``.
+    - If ``x1_i`` is a finite number, ``x2_i`` is a finite number, and ``x1_i`` equals ``x2_i``, the result is ``True``.
+    - In the remaining cases, the result is ``False``.
+
+    For complex floating-point operands, let ``a = real(x1_i)``, ``b = imag(x1_i)``, ``c = real(x2_i)``, ``d = imag(x2_i)``, and
+
+    - If ``a``, ``b``, ``c``, or ``d`` is ``NaN``, the result is ``False``.
+    - In the remaining cases, the result is the logical AND of the equality comparison between the real values ``a`` and ``c`` (real components) and between the real values ``b`` and ``d`` (imaginary components), as described above for real-valued floating-point operands (i.e., ``a == c AND b == d``).
+
+    .. note::
+       For discussion of complex number equality, see :ref:`complex-numbers`.
+
     Parameters
     ----------
     x1: array

spec/design_topics/complex_numbers.rst

@@ -0,0 +1,10 @@
+.. _complex-numbers:
+
+Complex Numbers
+===============
+
+Mathematically, equality comparison between complex numbers depends on the choice of topology. For example, the complex plane has a continuum of infinities; however, when the complex plane is projected onto the surface of a sphere (a stereographic projection commonly referred to as the *Riemann sphere*), infinities coalesce into a single *point at infinity*, thus modeling the extended complex plane. For the former, the value :math:`\infty + 3j` is distinct from (i.e., does not equal) :math:`\infty + 4j`, while, for the latter, :math:`\infty + 3j` does equal :math:`\infty + 4j`.
+
+Modeling complex numbers as a Riemann sphere conveys certain mathematical niceties (e.g., well-behaved division by zero and preservation of the identity :math:`\frac{1}{\frac{1}{z}} = z`); however, translating the model to IEEE 754 floating-point operations can lead to some unexpected results. For example, according to IEEE 754, :math:`+\infty` and :math:`-\infty` are distinct values; hence, for equality comparison, if :math:`x = +\infty` and :math:`y = -\infty`, then :math:`x \neq y`. In contrast, if we convert :math:`x` and :math:`y` to their complex number equivalents :math:`x = +\infty + 0j` and :math:`y = -\infty + 0j` and then interpret within the context of the extended complex plane, we arrive at the opposite result; namely, :math:`x = y`.
+
+In short, given the constraints of floating-point arithmetic and the subtleties of signed zeros, infinities, NaNs, and their interaction, crafting a specification which always yields intuitive results and satisfies all use cases involving complex numbers is not possible. Instead, this specification attempts to follow precedent (e.g., C99, Python, Julia, NumPy, and elsewhere), while also minimizing surprise. The result is an imperfect balance in which certain APIs may appear to embrace the one-infinity model found in C/C++ for algebraic operations involving complex numbers (e.g., considering :math:`\infty + \operatorname{NaN}\ j` to be infinite, irrespective of the imaginary component's value, including NaN), while other APIs may rely on the complex plane with its multiplicity of infinities (e.g., in transcendental functions). Accordingly, consumers of this specification should expect that certain results involving complex numbers for one operation may not be wholly consistent with results involving complex numbers for another operation.

spec/design_topics/index.rst

@@ -11,6 +11,7 @@ Design topics & constraints
    device_support
    static_typing
    accuracy
+   complex_numbers
    branch_cuts
    complex_number_ordering
    C_API