Make checkify error code assignment local and deterministic.

`checkify` works in two stages:
1. Tracing: It traces the input function into a jaxpr, where `checkify.check()` calls emit `check_p` primitive equations.
2. Discharging: It walks the jaxpr to convert those checks into a returned `Error` value.

Previously, integer error codes were assigned during the tracing stage using a process-wide global counter (`it.count(1)`). When multiple threads traced checkified functions concurrently, this shared counter led to non-deterministic code numbers. Because these numbers are baked as constants into the jaxpr, concurrent tracing caused spurious cache misses and recompilations.

Now, error code assignment is moved to the discharging stage:
- Tracing stage: `checkify.check()` no longer assigns error codes, using a placeholder code (-1) instead.
- Discharging stage: `checkify` uses a local `CheckifyContext` that numbers checks sequentially starting at 0 for each checkified block.
- Higher-order primitives: For subexpressions like `jit`, `scan`, `while`, and `cond`, the context is threaded through recursively so each subexpression receives disjoint codes. The `start_code` is included in the subexpression cache key to ensure cached jaxprs have unique codes.
- Nested checkify / re-injected errors: When `checkify.check_error(err)` is used to re-inject an `Error` produced by an inner `checkify`, `check_discharge_rule` preserves the inner error's existing code and metadata rather than treating it as an unassigned placeholder.

PiperOrigin-RevId: 971884514