Add a jit decorator to comparison operators.
Add a jit decorator to comparison operators.

We had previously omitted jit because of https://github.com/google/jax/issues/6713, but concrete remat no longer exists.

PiperOrigin-RevId: 474309905

google/jax,google/jax,google/jax,google/jax