Array objects of arbitrary rank are infeasible - require a reasonable range of ranks instead

[Zac-HD]

https://github.com/data-apis/array-api/blob/main/spec/API_specification/array_object.rst#17 says:

Furthermore, a conforming implementation of the array API standard must support array objects of arbitrary rank N (i.e., number of dimensions), where N is greater than or equal to zero.

Unfortunately this is infeasible: on any computer there's going to be a finite maximum dimensionality, and in practice it's usually much smaller - e.g. numpy.array_api is limited to 32-dimensional arrays. When it's impossible to comply with the standard, implementations will quite reasonably choose whatever noncompliant option makes sense for them.

I'd therefore propose that this should instead require implementations to support arrays with between zero and 32 dimensions inclusive (or another reasonable constant TBD), and permit implementations to support a larger number of dimensions if they wish.

[tylerjereddy]

From the pykokkos side, I asked @crtrott on Kokkos team, since we ultimately depend on the compiled C++ Kokkos bindings under the hood.

Kokkos right now supports rank 8

Mdspan supports arbitrary

In practice your compiler will hit recursion limits at between rank 100 and 200 or so

But mdspan is the multidimensional array which comes the closest to truly arbitrary as far as i know

[rgommers]

That seems like a reasonable change to make. I think the absolute minimum number of dimensions that are needed to support the features used in the standard (e.g., linear algebra routines with a batch dimension) is 3. In practice, >3 are used though, so it'd be good to set the number higher.

Is there any library that does not support 8 dimensions?

[leofang]

8 is way too small for some APIs like einsum. We've hit several issues with cuQuantum in which NumPy/CuPy only support up to 32 ndim. cuTENSOR has supported up to 64 (and will further relax in a future release) and PyTorch has no limitation.

[kgryte]

@Zac-HD What is motivating this request? Am I correct that this stems from a concern in Hypothesis?

Without further context, I am leery about imposing any explicit limits on array rank in the specification, as does not make sense to me to arbitrarily impose constraints based on present practice, when we don't have a crystal ball for the future.

While not constrained in Python, ndim is likely constrained in practice (e.g., either int32 or int64). In which case, I'd be fine saying that the number of dimensions must fit within an int32.

[kgryte]

@Zac-HD Is your proposal that the spec modify its language something to the effect of

Furthermore, a conforming implementation of the array API standard must support array objects of rank N (i.e., number of dimensions), where N is on the interval [0, 32] and where the upper bound is the minimum upper bound (i.e., a conforming implementation must support arrays having at least 32 dimensions and may support arrays having more than 32 dimensions).

[asmeurer]

8 is way too small for some APIs like einsum. We've hit several issues with cuQuantum in which NumPy/CuPy only support up to 32 ndim. cuTENSOR has supported up to 64 (and will further relax in a future release) and PyTorch has no limitation.

Presumably many of the dimensions in these cases are size 1? A rank $k$ array without trivial dimensions has at least $2^k$ entries. I'm assuming that's where NumPy's 32 limit comes from, since intp can be int32.

[Zac-HD]

I ran into this when thinking about how to generate validly-shaped arrays for Xarray (pydata/xarray#6908 (comment)), and obviously it does have implications for our Array-API array_shapes() strategy. My objection is that the current specification is impossible to implement, and if the spec is inconsistent with any implementation then the spec should probably change.

I'm certainly not proposing that the specification should cap the maximum array rank. It does specify that rank-zero arrays must be supported, and I think that's worth keeping. I'm also dubious about requiring all implementations to support some specific rank - in practice rank-N arrays are widely used up to at least rank-4, but the comments above that requiring rank-8 would be too little seem more domain-specific to me. To propose some specific language:

Furthermore, a conforming implementation of the array API standard must support array objects of rank N (i.e., number of dimensions), where N is greater than or equal to zero, and document their maximum supported rank N.

Depending on the library this might be a deliberate choice like Numpy's rank-32 limit or Kokkos rank-8, or it might be implementation-defined like "your compiler will likely crash somewhere in the hundreds of dimensions".

In summary: if literally nobody complies with the spec, the spec should change.

[leofang]

Presumably many of the dimensions in these cases are size 1?

Precisely. In many cases we encounter tensors with many axes of extent 1, so ndim can easily go beyond 32 while the whole array can still be stored on single GPU while making sense.

My objection is that the current specification is impossible to implement, and if the spec is inconsistent with any implementation then the spec should probably change.

As mentioned above, this is not true. PyTorch does not have this limitation, so you get at least 1 compliant library 🙂 (Update: CuPy can easily go beyond ndim=32 too)

[leofang]

Furthermore, a conforming implementation of the array API standard must support array objects of rank N (i.e., number of dimensions), where N is greater than or equal to zero, and document their maximum supported rank N.

I would think this is a safe statement.

[Zac-HD]

As mentioned above, this is not true. PyTorch does not have this limitation, so you get at least 1 compliant library 🙂 (Update: CuPy can easily go beyond ndim=32 too)

In [6]: torch.zeros((1,)*2000)
Out[6]: ---------------------------------------------------------------------------
RecursionError                            Traceback (most recent call last)
...
RecursionError: maximum recursion depth exceeded while calling a Python object

Still not arbitrarily large! (I'm impressed that it handles rank-1000 though 🤯)

[kgryte]

My sense is that we'd still want a minimum upper bound. Otherwise, users lack portability guarantees. E.g., one may be able to perform a batch operation on a rank 10 array in library X, but then encounter an error in library Y. Would be beneficial if spec compliant array libraries all had a minimum level of support.

[leofang]

Right I think Zac was suggesting the same thing too. Would also be nice to learn about the limitation in distributed array libraries. I just checked cuNumeric, they offer prebuilt binaries up to ndim=4, and building from source would support up to ndim=9. Ping @jakirkham for Dask.

[kgryte]

@leofang I'd want something a bit stronger than the following, however:

Furthermore, a conforming implementation of the array API standard must support array objects of rank N (i.e., number of dimensions), where N is greater than or equal to zero, and document their maximum supported rank N.

[tylerjereddy]

I believe the pykokkos-base templated C++ bindings built through CMake also experience a massive increase in compile time as you increase the max supported ranks from 4 to 7 or so. IIRC you can be talking about more than an hour vs. a minute and a half or so. Not sure if other libs experience that as well.