[FEATURE REQUEST] Adopt Array API standard to support torch, jax, etc.

[btjanaka]

Description

The Python Array API standard was introduced in 2021. It provides a common interface for array libraries. Currently, pyribs is exclusively dependent on NumPy. However, I believe we could support using a wide variety of libraries like PyTorch, CuPy, JAX, etc. by switching to the Array API, particularly by using the array-api-compat package https://github.com/data-apis/array-api-compat

This seems to be becoming quite popular, with adoption from libraries like scipy (scipy/scipy#18867) and Gymnasium (Farama-Foundation/Gymnasium#1388) (albeit still experimental or at least work-in-progress).

As a starting point, I think it would be great to integrate this and get pyribs working with NumPy and PyTorch on a single component, e.g., a GridArchive. PyTorch seems like a good choice since we use it in many tutorials.

TODO

• Initial setup (Start Array API infrastructure #573)
• ArrayStore (Support Array API in ArrayStore #571, Add CuPy support to ArrayStore via Array API #579)
• Tidy up ArrayStore dtypes (Use regular dtypes rather than scalar types in ArrayStore #577, Treat archive dtypes as regular dtypes rather than scalar dtypes #584)
• Deal with pickling -- right now, you can't pickle archives because the xp object is a module that cannot be pickled: Allow pickling classes that have the _xp array module attribute #644
• GridArchive (Support array API in GridArchive #578)
• GaussianEmitter
• Scheduler
• sphere example (end-to-end evaluation)
  • Benchmark performance on GPU
  • Update to support torch
• EvolutionStrategyEmitter

---

• Try setting up GPU CI: https://cirun.io -> It turns out that cirun is only for running CI on your existing cloud. This way, there is an alternative to running on the cloud from GitHub Actions. Since we do not have an existing cloud, it would not be possible to use cirun.
• Add Array API strict for testing

[btjanaka]

Having played around with this for a bit now, I think we should leave things as is in pyribs and not try to make everything use the array API to provide multiple backends. There are a few reasons for this:

1. Adding this feature is not trivial. In particular, while the array API does provide unity among array libraries, there are a number of operations we use in pyribs that are specific to one library. For example, ravel_multi_index and unravel_index (used in GridArchive) are in numpy and cupy but not in torch. Similarly, the aggregate method that we use from numpy groupies is only available for torch in a separate extension library that has non-trivial installation instructions. All this means that to implement the various archives with array API, we need to do a bit more work than "just" replacing np with xp.
2. A more complex library is also probably less accessible. We want beginners to be able to read the code in pyribs and just modify it.
3. Adding array API support increases maintenance burden, as we now have to test things on many more backends.
4. The current code is already quite fast; again, in the big problems, it appears that evaluation is still the bottleneck.

All this being said, there are still certainly use cases where it makes sense to have, for instance, an archive that places its data on the GPU. From that point of view, I think it would be great to add new objects that support this behavior, like TorchGridArchive. I could also see Scheduler supporting array API, as it doesn't have any super complex operations. To potentially support any of this development, I am keeping ArrayStore as is, with array API support. I have also cleaned up remaining WIP stuff in #645.