Repository: google/jax

Issue:
Improve jnp.where error message about not handling pytree
jax.numpy.where doesn't seem to handle pytree?

```python
import jax.numpy as jnp
jnp.where(True, {'a':1}, {'b':2})
```

gives an uninformative error message

```bash
>>> jnp.where(True, {'a':1}, {'b':2})
Traceback (most recent call last):
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/core.py", line 1248, in _len
    return self.shape[0]
IndexError: tuple index out of range

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/numpy/lax_numpy.py", line 2152, in where
    return _where(condition, x, y)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/traceback_util.py", line 165, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/api.py", line 429, in cache_miss
    out_flat = xla.xla_call(
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/core.py", line 1675, in bind
    return call_bind(self, fun, *args, **params)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/core.py", line 1687, in call_bind
    outs = top_trace.process_call(primitive, fun, tracers, params)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/core.py", line 596, in process_call
    return primitive.impl(f, *tracers, **params)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/dispatch.py", line 142, in _xla_call_impl
    compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/linear_util.py", line 272, in memoized_fun
    ans = call(fun, *args)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/dispatch.py", line 169, in _xla_callable_uncached
    return lower_xla_callable(fun, device, backend, name, donated_invars,
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/profiler.py", line 206, in wrapper
    return func(*args, **kwargs)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/dispatch.py", line 197, in lower_xla_callable
    jaxpr, out_avals, consts = pe.trace_to_jaxpr_final(
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/profiler.py", line 206, in wrapper
    return func(*args, **kwargs)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/interpreters/partial_eval.py", line 1643, in trace_to_jaxpr_final
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/interpreters/partial_eval.py", line 1614, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers_)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/numpy/lax_numpy.py", line 2129, in _where
    x, y = _promote_dtypes(x, y)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/numpy/lax_numpy.py", line 514, in _promote_dtypes
    to_dtype, weak_type = dtypes._lattice_result_type(*args)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/dtypes.py", line 362, in _lattice_result_type
    dtypes, weak_types = zip(*(_dtype_and_weaktype(arg) for arg in args))
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/dtypes.py", line 362, in <genexpr>
    dtypes, weak_types = zip(*(_dtype_and_weaktype(arg) for arg in args))
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/dtypes.py", line 251, in _dtype_and_weaktype
    return dtype(value), any(value is typ for typ in _weak_types) or is_weakly_typed(value)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/dtypes.py", line 355, in dtype
    dt = np.result_type(x)
  File "<__array_function__ internals>", line 5, in result_type
  File "/usr/local/google/home/zhangqiaorjc/.pyenv/versions/3.9.4/lib/python3.9/site-packages/numpy/core/_internal.py", line 61, in _usefields
    names, formats, offsets, titles = _makenames_list(adict, align)
  File "/usr/local/google/home/zhangqiaorjc/.pyenv/versions/3.9.4/lib/python3.9/site-packages/numpy/core/_internal.py", line 29, in _makenames_list
    n = len(obj)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/core.py", line 464, in __len__
    return self.aval._len(self)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/core.py", line 1250, in _len
    raise TypeError("len() of unsized object") from err  # same as numpy error
jax._src.traceback_util.UnfilteredStackTrace: TypeError: len() of unsized object

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/numpy/lax_numpy.py", line 2152, in where
    return _where(condition, x, y)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/numpy/lax_numpy.py", line 2129, in _where
    x, y = _promote_dtypes(x, y)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/numpy/lax_numpy.py", line 514, in _promote_dtypes
    to_dtype, weak_type = dtypes._lattice_result_type(*args)
  File "<__array_function__ internals>", line 5, in result_type
  File "/usr/local/google/home/zhangqiaorjc/.pyenv/versions/3.9.4/lib/python3.9/site-packages/numpy/core/_internal.py", line 61, in _usefields
    names, formats, offsets, titles = _makenames_list(adict, align)
  File "/usr/local/google/home/zhangqiaorjc/.pyenv/versions/3.9.4/lib/python3.9/site-packages/numpy/core/_internal.py", line 29, in _makenames_list
    n = len(obj)
TypeError: len() of unsized object
```


Discussion:
@mattjj this should be the repro of what I saw earlier today
It's not clear to me that `jnp.where` should be able to handle pytrees – after all, its API is modeled after `np.where`, which does not support non-array inputs, and in general no other `jax.numpy` function that I'm aware of is written to support pytree inputs. This looks to me like a specific case of the idea discussed in #1012 and #8504, so perhaps the best approach would be to do this with the package that came out of those discussions: https://github.com/google/tree-math/ 

We could certainly improve the error message, though, probably by calling `_check_arraylike` on the inputs to `jnp.where`.
Thanks Jake! Yeah, I ended up using jax.tree_map. The error message could perhaps be clearer about jnp.where not handling pytree as you suggested.

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
@@ -2147,6 +2147,7 @@ def where(condition, x=None, y=None, *, size=None, fill_value=None):
     _check_arraylike("where", condition)
     return nonzero(condition, size=size, fill_value=fill_value)
   else:
+    _check_arraylike("where", condition, x, y)
     if size is not None or fill_value is not None:
       raise ValueError("size and fill_value arguments cannot be used in three-term where function.")
     return _where(condition, x, y)


Test patch: