Repository: google/jax

Issue:
stablehlo failed to infer result types when using vmapped cond of KeyArray
Minimal repro:
```python
import jax
jax.config.update('jax_enable_custom_prng', True)

def func(x):
  return jax.lax.cond(x, jax.random.PRNGKey, jax.random.PRNGKey, 0)

x = jax.numpy.array([True, False])
print(jax.vmap(func)(x))
```
Result:
```pytb
Traceback (most recent call last):
  File "/Users/vanderplas/github/google/jax/jax2tf_failure.py", line 8, in <module>
    print(jax.vmap(func)(x))
          ^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax2tf_failure.py", line 5, in func
    return jax.lax.cond(x, jax.random.PRNGKey, jax.random.PRNGKey, 0)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/.local/share/virtualenvs/jax-LBbfM5ix/lib/python3.11/site-packages/jaxlib/mlir/dialects/_stablehlo_ops_gen.py", line 4402, in __init__
    results = _ods_ir.InferTypeOpInterface(SelectOp).inferReturnTypes(
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: Failed to infer result types
```


Discussion:

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -3372,6 +3372,12 @@ def _select_jvp(primals, tangents):
 
 def _select_hlo_lowering(ctx, which, *cases):
   which_aval = ctx.avals_in[0]
+  aval_out, = ctx.avals_out
+
+  if core.is_opaque_dtype(aval_out.dtype):
+    return [aval_out.dtype._rules.select_mlir(
+        ctx, ctx.avals_in, aval_out, which, *cases)]
+
   if which_aval.dtype == np.dtype(np.bool_):
     assert len(cases) <= 2
     if len(cases) == 1: return cases
diff --git a/jax/_src/prng.py b/jax/_src/prng.py
--- a/jax/_src/prng.py
+++ b/jax/_src/prng.py
@@ -578,6 +578,28 @@ def _comparison_mlir(direction, reduction_op, identity,
   ne_mlir = staticmethod(partial(_comparison_mlir, 'NE', hlo.OrOp,
                                  lax_internal._get_bitwise_or_identity))
 
+  @staticmethod
+  def select_mlir(ctx, avals_in, aval_out, which, *cases):
+    assert all(aval_case == aval_out for aval_case in avals_in[1:])
+    assert avals_in[0].ndim == aval_out.ndim
+    select_lower = lax_internal._select_hlo_lowering
+    key_shape = aval_out.dtype.impl.key_shape
+
+    aval_which = avals_in[0]
+    aval_which_bcast = core.ShapedArray(
+        (*aval_which.shape, *key_shape), aval_which.dtype)
+    aval_out_raw = core.ShapedArray(
+        (*aval_out.shape, *key_shape), np.dtype('uint32'))
+    aval_cases_raw = [aval_out_raw] * (len(avals_in) - 1)
+
+    bcast_dims = list(range(aval_which.ndim))
+    which_bcast = mlir.broadcast_in_dim(
+        ctx, which, aval_which_bcast, broadcast_dimensions=bcast_dims)
+
+    return mlir.delegate_lowering(ctx, select_lower, which_bcast, *cases,
+                                  avals_in=[aval_which_bcast, *aval_cases_raw],
+                                  avals_out=[aval_out_raw])[0]
+
   @staticmethod
   def device_put_sharded(vals, aval, sharding, devices):
     physical_aval = keys_aval_to_base_arr_aval(aval)


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -2925,6 +2925,24 @@ def gather_mlir(ctx, avals_in, aval_out, x, indices, *,
                                   avals_in=[aval_x_raw, aval_indices],
                                   avals_out=[aval_y_raw])[0]
 
+  @staticmethod
+  def select_mlir(ctx, avals_in, aval_out, which, *cases):
+    assert all(aval_case == aval_out for aval_case in avals_in[1:])
+    assert avals_in[0].ndim == aval_out.ndim
+    select_lower = lax_internal._select_hlo_lowering
+    aval_which = avals_in[0]
+    aval_which_bcast = core.ShapedArray(
+        (*aval_which.shape, 2), aval_which.dtype)
+    aval_out_raw = core.ShapedArray(
+        (*aval_out.shape, 2), np.dtype('uint32'))
+    aval_cases_raw = [aval_out_raw] * (len(avals_in) - 1)
+    bcast_dims = list(range(aval_which.ndim))
+    which_bcast = mlir.broadcast_in_dim(
+        ctx, which, aval_which_bcast, broadcast_dimensions=bcast_dims)
+    return mlir.delegate_lowering(ctx, select_lower, which_bcast, *cases,
+                                  avals_in=[aval_which_bcast, *aval_cases_raw],
+                                  avals_out=[aval_out_raw])[0]
+
 
 class FooTy:
   name = 'foo'
@@ -3197,6 +3215,13 @@ def test_gather(self):
     self.assertIsInstance(ys, FooArray)
     self.assertEqual(ys.shape, (3, 2, 1))
 
+  def test_select(self):
+    ks = jax.jit(lambda: make((3,)))()
+    cs = jnp.array([True, False, False])
+    ys = jax.jit(lax.select)(cs, ks, ks)
+    self.assertIsInstance(ys, FooArray)
+    self.assertEqual(ys.shape, (3,))
+
   def test_xla_reverse_bug(self):
     # Regression test for b/248295786
     # This was an XLA bug related to an incorrect optimization of reverse
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -1835,6 +1835,25 @@ def test_gather(self):
     self.assertIsInstance(ys, random.KeyArray)
     self.assertEqual(ys.shape, (3, 2, 1))
 
+  @skipIf(not config.jax_enable_custom_prng,
+          'requires config.jax_enable_custom_prng')
+  def test_select(self):
+    ks = self.make_keys(3, 2)
+    cs = jnp.array([True, False, False, True, False, True]).reshape(3, 2)
+    ys = jax.jit(lax.select)(cs, ks, ks)
+    self.assertIsInstance(ys, random.KeyArray)
+    self.assertEqual(ys.shape, (3, 2))
+
+  @skipIf(not config.jax_enable_custom_prng,
+          'requires config.jax_enable_custom_prng')
+  def test_select2(self):
+    # See https://github.com/google/jax/issues/15869
+    def f(x):
+      keys = lax.broadcast(jax.random.PRNGKey(0), x.shape)
+      return lax.select(x, keys, keys)
+    x = jnp.array([True, False, False])
+    f(x)  # doesn't crash
+
   def test_device_put(self):
     device = jax.devices()[0]
     keys = self.make_keys(4)