Repository: google/jax

Issue:
`jnp.sort` and `jnp.vdot` raise error under `jax.disable_jit()`
### Description

`jnp.sort`, `jnp.vdot` and a few others in `jax.numpy` will raise under `jax.disable_jit()`. IMO, the fact that exactly *these* functions go wrong is not important, nor is the fact that `jax.disable_jit()` is involved.

Taking a look at their implementation, on can spot a pattern:

```python
@partial(jit, static_argnums=(...))
def f(a: ArrayLike):
  util._check_arraylike('f', a)
  bar = foo(a.ravel())  # << not good
  ...
```

Without jit, the scalar case will get caught by the tests, lead us to the fix naturally. These functions do array conversion implicitly on the jit boundary hides the fact that array methods might not be available.

Fixing them is easy enough, but I want to ask/discuss something more general here:
1. Maybe we should insert these explicit conversions after `_check_arraylike` everywhere, in the form of `asarray` or explicit `reshape` or `ravel`.
2. If the pattern in 1. is desirable, is it always safe to use the simpler and faster variant `util._asarray`? I think if we can, we should.
3. I see similar problems (that JAX maintainers have rather canonical solutions in JAX itself) popping up almost every time someone tries to contribute to `jax.scipy` etc. Can we have a contributor guide documenting these issues? (`array` conversion, `dtype` handling, type hinting, `jnp` vs `lax`, standard ways to deal with static arguments: `concrete_or_error`, etc.)
4. What is the best place to ask these broader questions that are also half bug reports? Should the Github discussions be preferred?

@jakevdp I bore you the most in many related issues before, so I'm gonna tag you here.


Discussion:
Agreed we should always call `_check_arraylike` followed by casting the value using `asarray` in order to address this problem, though it's a common-enough pattern we might want to replace it with a more streamlined utility.
> we might want to replace it with a more streamlined utility

Maybe something like
```python
a, b = _ensure_arraylike('f', a, b)
```
which uses something a bit lighter than `jnp.array` internally. This might also helps with typing. Not sure if MyPy supports this kind of type refinement though, as I see many `xxx_arr` in `lax_numpy.py`.

That said, I also see instances like
```python
util._stackable(a) or util._check_arraylike("transpose", a)
```
and
```python
core.is_special_dim_size(repeats) or util._check_arraylike("repeat", repeats)
```
So whatever the new thing is, it better work well will these as well.

In any case, I'm all for a more complete `numpy_util` which I could use in my own work. Also looking forward to less _internal_ tools for `dtype` checking, type promotion, static argument, etc.
Yeah, that's roughly what I've been thinking. Though for typing, if you want to use array attributes you have to rename the variables, e.g.
```python
a_arr, b_arr = util._arraylike(a, b)
```
It's one of my biggest annoyances with the current generation of static type checkers, actually
> It's one of my biggest annoyances with the current generation of static type checkers, actually

Yeah. Pyright seems to work without renaming though, so there is hope.
Fix for this particular issue here: #14908

I'll probably tackle the simplification of `check_arraylike` in a later PR – it's a bigger job.
@jakevdp Thanks!

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -3057,7 +3057,7 @@ def vdot(a, b, *, precision=None):
   util._check_arraylike("vdot", a, b)
   if issubdtype(_dtype(a), complexfloating):
     a = ufuncs.conj(a)
-  return dot(a.ravel(), b.ravel(), precision=precision)
+  return dot(ravel(a), ravel(b), precision=precision)
 
 
 @util._wraps(np.tensordot, lax_description=_PRECISION_DOC)
@@ -3487,9 +3487,9 @@ def sort(a, axis: Optional[int] = -1, kind='quicksort', order=None):
     raise ValueError("'order' argument to sort is not supported.")
 
   if axis is None:
-    return lax.sort(a.ravel(), dimension=0)
+    return lax.sort(ravel(a), dimension=0)
   else:
-    return lax.sort(a, dimension=_canonicalize_axis(axis, ndim(a)))
+    return lax.sort(asarray(a), dimension=_canonicalize_axis(axis, ndim(a)))
 
 @util._wraps(np.sort_complex)
 @jit


Test patch: