Repository: jax-ml/jax
Issue category: Performance Issue

Issue:
jnp.select should use jax.lax.select_n
### Description

jnp.select is a function that's like jnp.where, but selects between n different arrays. jnp.where is a wrapper around jax.lax.select, which makes it more flexible in terms of input shapes and dtypes.

jax.lax.select_n (added in #9482 ) is a generalization of jax.lax.select to pick from n arrays rather than 2; it translates at lower time into a binary tree of jax.lax.select's, which is helpful for efficiency when there are many branches. jnp.select doesn't use this: it linearly chains jnp.where's across the different arrays. It would likely help performance if it did, particularly in the case where the condlist entries have lower ndim than the choicelist.

### System info (python version, jaxlib version, accelerator, etc.)

jax:    0.4.27
jaxlib: 0.4.27
numpy:  1.26.3
python: 3.11.8 (stable, redacted, redacted) [Clang google3-trunk (3b5e7c83a6e226d5bd7ed2e9b67449b64812074c)]
jax.devices (16 total, 16 local): [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0) TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1) ... TpuDevice(id=14, process_index=0, coords=(3,1,0), core_on_chip=0) TpuDevice(id=15, process_index=0, coords=(3,1,0), core_on_chip=1)]
process_count: 1
platform: uname_result(system='Linux', node='ebe6e776479e64ea-65d7388ae36.borgtask.google.com', release='5.10.0-smp-1101.33.0.0', version='#1 [v5.10.0-1101.33.0.0] SMP @1708585970', machine='x86_64')



Hints:

Patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 3767633deefc..f7686ec467d6 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -1162,12 +1162,12 @@ def select(
     raise ValueError(msg.format(len(condlist), len(choicelist)))
   if len(condlist) == 0:
     raise ValueError("condlist must be non-empty")
-  choices = util.promote_dtypes(default, *choicelist)
-  choicelist = choices[1:]
-  output = choices[0]
-  for cond, choice in zip(condlist[::-1], choicelist[::-1]):
-    output = where(cond, choice, output)
-  return output
+  # Put the default at front with condition False because
+  # argmax returns zero for an array of False values.
+  choicelist = util.promote_dtypes(default, *choicelist)
+  conditions = stack(broadcast_arrays(False, *condlist))
+  idx = argmax(conditions.astype(bool), axis=0)
+  return lax.select_n(*broadcast_arrays(idx, *choicelist))
 
 
 @util.implements(np.bincount, lax_description="""\


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index ddc599792e63..ee6c3601a1bb 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -4532,6 +4532,7 @@ def testWhereScalarPromotion(self):
     # maximal set of dtypes.
     dtypes=itertools.combinations_with_replacement(all_dtypes, 3),
   )
+  @jax.numpy_rank_promotion('allow')
   def testSelect(self, n, shapes, dtypes):
     dtypes = dtypes[:n+1]
     rng = jtu.rand_default(self.rng())