Return type of jnp.where

[bitsandscraps]

Even though the document states that jnp.where returns

an [emphasis added] array with elements from x where condition is True, and elements from y elsewhere

the return type of jnp.where is set to be Array | tuple[Array, ...] in the code.
When does jnp.where return a tuple of arrays?

[jakevdp]

Hi, thanks for the question! jax.numpy.where follows the API of numpy.where, and the output type depends on the number of arguments. With one argument, jnp.where is equivalent to jnp.nonzero, and returns a tuple of indices:

>>> x = jnp.array([[True, False],
...                [False, True]])
>>> jnp.where(x)
(Array([0, 1], dtype=int32), Array([0, 1], dtype=int32))

With three arguments, it returns the type in the documentation excerpt you quoted:

>>> jnp.where(x, 100, 200)
Array([[100, 200],
       [200, 100]], dtype=int32)

I think this behavior is covered in the Note at the top of the documentation you linked to.

[bitsandscraps]

Thanks!