[mosaic_gpu] Fixed FragmentedArray comparisons with literals

PiperOrigin-RevId: 698343858

Author: superbobry

tests/mosaic/gpu_test.py

@@ -1318,18 +1318,21 @@ def kernel(ctx, dst, _):
           operator.ne,
       ],
       dtype=[jnp.float32, jnp.int32, jnp.uint32],
+      rhs_is_literal=[False, True]
   )
-  def test_comparison(self, op, dtype, m=64, n=32):
+  def test_comparison(self, op, dtype, rhs_is_literal, m=64, n=32):
     def kernel(ctx, dst, _):
       iota = iota_tensor(m, n, dtype)
-      op(iota, iota + 1).store_untiled(dst)
+      rhs = 0 if rhs_is_literal else iota + 1
+      op(iota, rhs).store_untiled(dst)
 
     out_shape = jax.ShapeDtypeStruct((m, n), jnp.bool)
     result = mgpu.as_gpu_kernel(
         kernel, (1, 1, 1), (128, 1, 1), (), out_shape, ()
     )()
     iota = np.arange(m * n, dtype=dtype).reshape(m, n)
-    np.testing.assert_array_equal(result, op(iota, iota + 1))
+    rhs = rhs = 0 if rhs_is_literal else iota + 1
+    np.testing.assert_array_equal(result, op(iota, rhs))
 
   @parameterized.product(
       op=[operator.and_, operator.or_, operator.xor],