Add streamlined _ensure_arraylike utility

[jakevdp]

Over time we've landed on the following pattern for user-facing APIs in jax.numpy, jax.scipy, jax.random, etc.:

import jax.numpy as jnp
from jax._src.numpy.util import _check_arraylike
from jax.typing import ArrayLike
from jax import Array

def func(x: ArrayLike, ...) -> Array:
  _check_arraylike("func", x)  # this is essentially a runtime check of isinstance(x, ArrayLike)
  x = jnp.asarray(x)  # this actually converts x to Array so that, e.g. it can be passed to lax functions
  # do some computation with x
  return out

There are a couple improvements we could make here:

1. We essentially never want to call _check_arraylike without following with jnp.asarray on the inputs. Whe should combine these into a single utility! This would prevent failures like the one reported in jnp.sort and jnp.vdot raise error under jax.disable_jit() #14906.
2. jnp.asarray is a pretty heavy function; if we know that an input passes _check_arraylike, we can convert to an array much more efficiently.
3. This unified utility might also handle the stackable checks that are currently in numpy, via an extra keyword like allow_stackable=True. Also perhaps incorporate check_no_float0s via allow_float0s=False.
4. _check_arraylike lives in jax._src.numpy.util, but it is far more generally applicable. It should probably live in jax._util, or perhaps even be a public utility so that we can recommend this pattern for downstream APIs (e.g. here: https://jax.readthedocs.io/en/latest/jax.typing.html#jax-typing-best-practices)

[soraros]

@jakevdp Thanks for investigating this!

Wow, I didn't know the existence of this (jax.readthedocs.io/en/latest/jax.typing.html#jax-typing-best-practices). It's such a nature place where all the array API related (type checking, dtype checking, dtype promotion, staticness, etc.) tools can go. Do you think we could open an umbrella issue to track the inception of such a set of (end user facing) tools?

[jakevdp]

That doc is only a week or two old. #12049 is a relevant tracking issue