Repository: google/jax
Issue: Combining where and pad gives garbage data

Reproducer, running on CPU:

```python
import jax
import jax.numpy as jnp
import haiku as hk

key = hk.PRNGSequence(17)


def create_network():
    action_space_1 = 32
    action_space_2 = 15

    input_shapes = [
        1,
        6,
        action_space_2,
        2, 
    ]

    def network(inputs):
        mask_choice = inputs[..., 0, None]
        mask_idx = slice(sum(input_shapes[:2]), sum(input_shapes[:3]))

        out_value = hk.Linear(1)(inputs)

        out_action_1 = hk.Linear(action_space_1)(inputs)

        out_action_2 = hk.Linear(action_space_2)(inputs)
        out_action_2 = jnp.where(
            inputs[..., mask_idx] == 1,
            -jnp.inf,
            out_action_2,
        )

        def pad_action(logit, num_pad):
            return jnp.pad(logit, [(0, 0), (0, num_pad)], constant_values=-jnp.inf)

        if action_space_1 < action_space_2:
            out_action_1 = pad_action(out_action_1, action_space_2 - action_space_1)

        elif action_space_1 > action_space_2:
            out_action_2 = pad_action(out_action_2, action_space_1 - action_space_2)

        out_action = jnp.where(mask_choice == 0, out_action_2, out_action_1)

        return out_action, jnp.squeeze(out_value, axis=-1)

    forward = hk.without_apply_rng(hk.transform(network))
    return forward, input_shapes


net, input_shape = create_network()
network = jax.jit(net.apply)

weights = net.init(
    next(key), jnp.ones((1, sum(input_shape)), dtype=jnp.float32)
)

print(network(weights, jnp.zeros((1, sum(input_shape)), dtype=jnp.float32)))
```

Expected output:

```
(DeviceArray([[  0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,
                0.,   0.,   0.,   0.,   0., -inf, -inf, -inf, -inf, -inf,
              -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,
              -inf, -inf]], dtype=float32), DeviceArray([0.], dtype=float32))
```

Actual output:

```
(DeviceArray([[ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
               0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
               0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
               0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
               0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
                        -inf,  1.4012985e-45,  0.0000000e+00,
               1.2611686e-44,  0.0000000e+00, -3.0224281e-27,
               3.0944874e-41,  1.4012985e-45,  0.0000000e+00,
               1.2611686e-44,  0.0000000e+00,  0.0000000e+00,
               0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
               0.0000000e+00,  0.0000000e+00]], dtype=float32), DeviceArray([0.], dtype=float32))
```

Note how the `-inf` values are replaced by bit garbage.

(Sorry for the lengthy script, I couldn't really condense it any further. This uses `haiku`, but I don't think this is a `haiku` issue.)



Changed files:
tests/lax_test.py, jax/test_util.py

diff --git a/jax/test_util.py b/jax/test_util.py
index 3871ed474..19afbe8bf 100644
--- a/jax/test_util.py
+++ b/jax/test_util.py
@@ -873,7 +873,9 @@ class JaxTestCase(parameterized.TestCase):
     """Assert that x and y arrays are exactly equal."""
     if check_dtypes:
       self.assertDtypesMatch(x, y)
-    np.testing.assert_array_equal(x, y, err_msg=err_msg)
+    # Work around https://github.com/numpy/numpy/issues/18992
+    with np.errstate(over='ignore'):
+      np.testing.assert_array_equal(x, y, err_msg=err_msg)
 
   def assertArraysAllClose(self, x, y, *, check_dtypes=True, atol=None,
                            rtol=None, err_msg=''):
diff --git a/tests/lax_test.py b/tests/lax_test.py
index 028d6004d..9ca26b04e 100644
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -2406,6 +2406,21 @@ class LaxTest(jtu.JaxTestCase):
     self.assertEqual(lax.igamma(1., np.inf), 1.)
     self.assertEqual(lax.igammac(1., np.inf), 0.)
 
+  @unittest.skipIf(jax.lib.version < (0, 1, 66),
+                   "Test fails on jaxlib 0.1.65 or earlier.")
+  def testRegressionIssue5728(self):
+    # The computation in this test gave garbage data on CPU due to an LLVM bug.
+    @jax.jit
+    def f(inputs):
+      out_action_2 = lax.slice_in_dim(inputs, 0, 15, axis=-1)
+      mask = lax.slice_in_dim(inputs, 7, 22, axis=-1)
+      out_action_2 = lax.select(lax.eq(mask, np.float32(0)),
+                                lax.broadcast(np.float32(42), (1, 15)),
+                                out_action_2)
+      return lax.pad(out_action_2, np.float32(42), [(0, 0, 0), (0, 15, 0)])
+    self.assertArraysEqual(np.full((1, 30), np.float32(42)),
+                           f(np.zeros((1, 24), dtype=np.float32)))
+
 
 class LazyConstantTest(jtu.JaxTestCase):
   def _Check(self, make_const, expected):