RFC: add cbrt

[steff456]

This RFC requests to include a new API in the array API specification for the purpose of computing the cube root.

Overview

Based on array comparison data, the API is available in the majority of libraries in the PyData ecosystem.

The Array API Specification currently includes pow and sqrt, but does not include the IEEE 754 function cbrt.

• pow:

  array-api/src/array_api_stubs/_2022_12/elementwise_functions.py

  Line 1786 in 3d91878

  def pow(x1: array, x2: array, /) -> array:

• sqrt:

  array-api/src/array_api_stubs/_2022_12/elementwise_functions.py

  Line 2139 in 3d91878

  def sqrt(x: array, /) -> array:

While the cube root could be implemented in terms of pow, this is not desirable as the rational 1/3 is not typically equal to 1.0/3.0 due to limited floating-point precision. Cube root implementations are generally more accurate than the equivalent operation via pow.

Prior art

• NumPy: https://numpy.org/doc/stable/reference/generated/numpy.cbrt.html
• PyTorch: (not currently implemented)
• MXNet: https://mxnet.apache.org/versions/master/api/python/docs/api/np/generated/mxnet.np.cbrt.html
• TensorFlow: https://www.tensorflow.org/api_docs/python/tf/experimental/numpy/cbrt
• C99: https://en.cppreference.com/w/c/numeric/math/cbrt

Proposal:

def cbrt(x: array, /) -> array

cc @kgryte

[asmeurer]

It looks like the NumPy cbrt doesn't support complex arguments.

Presumably for a floating-point negative argument cbrt should return a real negative number, but for a complex real negative argument it should return the principle cube root (if we add complex support).

[kgryte]

Complex cube root is tricky (see JuliaLang/julia#36534). My initial sense is that we'd want to leave unspecified if complex dtypes are supported.

[asmeurer]

Even without implementing the complex version there's still the fact that cbrt(-1) returning -1 is a non-principle cube root (differing from power):

>>> np.power(-1, 1/3)
<stdin>:1: RuntimeWarning: invalid value encountered in power
nan
>>> np.cbrt(-1)
-1.0

[kgryte]

Agreed. For negative reals, cbrt should return a negative real number. And thanks for highlighting a key limitation of power.

[asmeurer]

I don't really see it as a limitation. There are very good reasons for power to return principle roots, even if that means returning nan instead of a real value for noncomplex floating-point inputs.

Of course, sometimes you do want a real cube root. And it seems that that's the C standard for what cbrt does, so we should definitely do that.

I agree that omitting an extension to complex numbers is probably best because that would require choosing between either returning different values for negative reals, or using a non-principle branch cut, and both seem problematic. Maybe worth bringing up in a consortium meeting though.

[kgryte]

Linking to a (closed) feature request on PyTorch: pytorch/pytorch#25766