refactor: implement nplike registry (1 of 2)

[agoose77]

TL;DR

This PR replaces a hard-coded loop-based search with a faster per-type lookup for each nplike.

TL

Context

We have an nplike_of L3 function in Awkward so that we can determine the most suitable array library among a set of complex input objects. This mechanism needs to extend beyond identifying the array library for the array type to layouts, builders, and potentially other objects.

Our existing mechanism for finding the appropriate nplike for a given object is flawed, due to the following points:

1. Our L1 API should not expose backend or nplike
2. Our L2 APIs cannot all be made homogeneous, e.g. _ext.ArrayBuilder should not have a backend (it has no concept of such a thing)
3. Each call to nplike_for performs a non-cached lookup of several attributes. For array types, it loops over each nplike
4. The set of nplikes is hard-coded

Extending nplike_for to include ArrayBuilder would require more customisation of this logic. At the same time, PRs like #2305 would benefit from the search for an nplike being as fast as possible.

Details

In general, we can do better by generalising the search mechanism. This PR makes such a change, by introducing a type-cache that allows us to short-circuit the existing for loop (unless there is a cache miss):

1. There is now a register_nplike function that lets the caller register an nplike. This replaces the need to explicitly enumerate the nplikes with the need to explicitly import all nplikes. Therefore, this is mainly an aesthetic change.
2. nplike_of is now only defined for individual array objects operated upon by nplikes, i.e. not ak.index.Index.
3. nplike_of is memoized by the array type.

Changes

• ak.index.Index(ak.index.Index(...)) is no longer supported
• nplike_of(obj) fails for objects that are not np.ndarray, or cp.ndarray, ..., or TypeTracerArray.

[codecov]

Codecov Report

Merging #2389 (8d7618a) into main (265f632) will increase coverage by 0.04%.
The diff coverage is 98.05%.

Additional details and impacted files

Impacted Files	Coverage Δ
src/awkward/_nplikes/numpylike.py	73.81% <75.00%> (+0.01%)	⬆️
src/awkward/_nplikes/dispatch.py	96.55% <96.55%> (ø)
src/awkward/_backends.py	88.75% <100.00%> (+4.90%)	⬆️
src/awkward/_kernels.py	68.93% <100.00%> (ø)
src/awkward/_layout.py	83.50% <100.00%> (+0.17%)	⬆️
src/awkward/_nplikes/__init__.py	77.77% <100.00%> (-12.97%)	⬇️
src/awkward/_nplikes/array_module.py	91.17% <100.00%> (+0.19%)	⬆️
src/awkward/_nplikes/cupy.py	42.46% <100.00%> (+1.62%)	⬆️
src/awkward/_nplikes/jax.py	84.61% <100.00%> (+0.83%)	⬆️
src/awkward/_nplikes/numpy.py	64.00% <100.00%> (+1.50%)	⬆️
... and 7 more

[jpivarski]

I like it! Just one question, though:

• nplike_of(obj) fails for objects that are not np.ndarray, or cp.ndarray, ..., or TypeTracerArray

Don't we need this to deal with array-like function arguments? The tests would fail if what I'm thinking is true, and that hasn't happened, so perhaps the array-like → array conversion (to_layout) has already happened by the time nplike_of gets called?

[agoose77]

I like it! Just one question, though:

• nplike_of(obj) fails for objects that are not np.ndarray, or cp.ndarray, ..., or TypeTracerArray

Don't we need this to deal with array-like function arguments? The tests would fail if what I'm thinking is true, and that hasn't happened, so perhaps the array-like → array conversion (to_layout) has already happened by the time nplike_of gets called?

Yes! The main point of this PR is to make nplike_of less smart. It should only be used when we know that we have array-like objects (or, at least, we use it knowing that we don't have content objects). If we do have contents, then we should be using backend_of, which gives us both index_nplike and nplike.

This PR is an alternative solution to an idea I had of giving Index objects a Backend instead of nplike, such that it would be possible to mix Index and Content objects in backend_of. However, it doesn't make sense to talk about such things; index objects don't need the concept of a backend. It would only serve to disambiguate mixing Index and Content, and the easier solution to that is not to; Index objects come from a Content object, so use the Content in backend_of.