From 4788fc46b8ece269f7fee985e5891cf17e356fec Mon Sep 17 00:00:00 2001 From: Athan Reines Date: Mon, 5 Dec 2022 11:02:18 -0800 Subject: [PATCH] Add complex number support to `square` --- .../array_api/elementwise_functions.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/spec/API_specification/array_api/elementwise_functions.py b/spec/API_specification/array_api/elementwise_functions.py index c68076aa0..0f9986baa 100644 --- a/spec/API_specification/array_api/elementwise_functions.py +++ b/spec/API_specification/array_api/elementwise_functions.py @@ -1713,13 +1713,22 @@ def sinh(x: array, /) -> array: """ def square(x: array, /) -> array: - """ - Squares (``x_i * x_i``) each element ``x_i`` of the input array ``x``. + r""" + Squares each element ``x_i`` of the input array ``x``. + + The square of a number ``x_i`` is defined as + + .. math:: + x_i^2 = x_i \cdot x_i + + **Special cases** + + For floating-point operands, special cases must be handled as if the operation is implemented as ``x * x`` (see :func:`~array_api.multiply`). Parameters ---------- x: array - input array. Should have a real-valued data type. + input array. Should have a numeric data type. Returns -------