jax

5ab78d31 - [JAX] Use WEAKORDER comparison type for floating-point sort comparators

[JAX] Use WEAKORDER comparison type for floating-point sort comparators Update `lt_to_p`, `le_to_p`, and `eq_to_p` to lower floating-point comparisons to `stablehlo.compare` with `comparison_type = WEAKORDER` when not in forward-compatibility mode and `jaxlib` (>= 514) and `libtpu` (>= 0.0.52) support it, and remove explicit zero/NaN canonicalization from `_operands_to_keys`. When lowering for forward-compatible export or older `jaxlib`/`libtpu` versions, fall back to canonicalizing zeros and NaNs via `_canonicalize_float_for_sort` and emitting `TOTALORDER`. PiperOrigin-RevId: 1008722740