Add complex number support to sqrt

spec/API_specification/array_api/elementwise_functions.py

@@ -1482,28 +1482,55 @@ def square(x: array, /) -> array:
     """
 
 def sqrt(x: array, /) -> array:
-    """
-    Calculates the square root, having domain ``[0, +infinity]`` and codomain ``[0, +infinity]``, for each element ``x_i`` of the input array ``x``. After rounding, each result must be indistinguishable from the infinitely precise result (as required by IEEE 754).
+    r"""
+    Calculates the principal square root for each element ``x_i`` of the input array ``x``.
+
+    .. note::
+       After rounding, each result must be indistinguishable from the infinitely precise result (as required by IEEE 754).
 
     **Special cases**
 
-    For floating-point operands,
+    For real-valued floating-point operands,
 
     - If ``x_i`` is ``NaN``, the result is ``NaN``.
     - If ``x_i`` is less than ``0``, the result is ``NaN``.
     - If ``x_i`` is ``+0``, the result is ``+0``.
     - If ``x_i`` is ``-0``, the result is ``-0``.
     - If ``x_i`` is ``+infinity``, the result is ``+infinity``.
 
+    For complex floating-point operands, let ``a = real(x_i)``, ``b = imag(x_i)``, and
+
+    - If ``a`` is either ``+0`` or ``-0`` and ``b`` is ``+0``, the result is ``+0 + 0j``.
+    - If ``a`` is any value (including ``NaN``) and ``b`` is ``+infinity``, the result is ``+infinity + infinity j``.
+    - If ``a`` is a finite number and ``b`` is ``NaN``, the result is ``NaN + NaN j``.
+    - If ``a`` ``-infinity`` and ``b`` is a positive (i.e., greater than ``0``) finite number, the result is ``NaN + NaN j``.
+    - If ``a`` is ``+infinity`` and ``b`` is a positive (i.e., greater than ``0``) finite number, the result is ``+0 + infinity j``.
+    - If ``a`` is ``-infinity`` and ``b`` is ``NaN``, the result is ``NaN + infinity j`` (sign of the imaginary component is unspecified).
+    - If ``a`` is ``+infinity`` and ``b`` is ``NaN``, the result is ``+infinity + NaN j``.
+    - If ``a`` is ``NaN`` and ``b`` is any value, the result is ``NaN + NaN j``.
+    - If ``a`` is ``NaN`` and ``b`` is ``NaN``, the result is ``NaN + NaN j``.
+
+    .. note::
+       For complex floating-point operands, ``sqrt(conj(x))`` must equal ``conj(sqrt(x))``.
+
+    .. note::
+       By convention, the branch cut of the square root is the negative real axis :math:`(-\infty, 0)`.
+
+       The square root is a continuous function from above the branch cut, taking into account the sign of the imaginary component.
+
+       Accordingly, for complex arguments, the function returns the square root in the range of the right half-plane, including the imaginary axis (i.e., the plane defined by :math:`[0, +\infty)` along the real axis and :math:`(-\infty, +\infty)` along the imaginary axis).
+
+       *Note: branch cuts have provisional status* (see :ref:`branch-cuts`).
+
     Parameters
     ----------
     x: array
-        input array. Should have a real-valued floating-point data type.
+        input array. Should have a floating-point data type.
 
     Returns
     -------
     out: array
-        an array containing the square root of each element in ``x``. The returned array must have a real-valued floating-point data type determined by :ref:`type-promotion`.
+        an array containing the square root of each element in ``x``. The returned array must have a floating-point data type determined by :ref:`type-promotion`.
     """
 
 def subtract(x1: array, x2: array, /) -> array:

spec/design_topics/branch_cuts.rst

@@ -0,0 +1,17 @@
+.. _branch-cuts:
+
+Branch Cuts
+===========
+
+In the mathematical field of complex analysis, a **branch cut** is a curve in the complex plane across which an analytic multi-valued function is discontinuous. Branch cuts are often taken as lines or line segments, and the choice of any particular branch cut is a matter of convention.
+
+For example, consider the function :math:`z^2` which maps a complex number :math:`z` to a well-defined number :math:`z^2`. The function's inverse function :math:`\sqrt{z}` does not, however, map to a single value. For example, for :math:`z = 1`, :math:`\sqrt{1} = \pm 1`. While one can choose a unique principal value for this and similar functions (e.g., in this case, the principal square root is :math:`+1`), choices cannot be made continuous over the whole complex plane, as lines of discontinuity must occur. To handle discontinuities, one commonly adopts branch cuts, which are not, in general, unique. Instead, one chooses a branch cut as a matter of convention in order to give simple analytic properties.
+
+Branch cuts do not arise for single-valued trigonometric, hyperbolic, integer power, or exponential functions; however, branch cuts do arise for their multi-valued inverses.
+
+In contrast to real-valued floating-point numbers which have well-defined behavior as specified in IEEE 754, complex-valued floating-point numbers have no equivalent specification. Accordingly, this specification chooses to follow C99 conventions for special cases and branch cuts for those functions supporting complex numbers. For those functions which do not have C99 equivalents (e.g., linear algebra APIs), the specification relies on dominant conventions among existing array libraries.
+
+.. warning::
+   All branch cuts documented in this specification are considered **provisional**. While conforming implementations of the array API standard should adopt the branch cuts described in this standard, consumers of array API standard implementations should **not** assume that branch cuts are consistent between implementations. 
+
+   Provided no issues arise due to the choice of branch cut, the provisional status is likely to be removed in a future revision of this standard.

spec/design_topics/index.rst

@@ -11,5 +11,6 @@ Design topics & constraints
    device_support
    static_typing
    accuracy
+   branch_cuts
    C_API
    parallelism

spec/purpose_and_scope.md

@@ -434,6 +434,9 @@ a (usually fixed-size) multidimensional container of items of the same type and
 **axis**:
 an array dimension.
 
+**branch cut**:
+a curve in the complex plane across which a given complex function fails to be continuous.
+
 **broadcast**:
 automatic (implicit) expansion of array dimensions to be of equal sizes without copying array data for the purpose of making arrays with different shapes have compatible shapes for element-wise operations.