[QwixInterception] Selective disable_jit for XLA and ODML providers.

[copybara-service]

[QwixInterception] Selective disable_jit for XLA and ODML providers.

This CL refines the interception logic to only use disable_jit when absolutely necessary, distinguishing between high-level function interception (XLA) and primitive-level interception (ODML).

XLA Quantization (QtProvider, PtqProvider):

• Doesn't disable_jit by default.
• Intercepts high-level Functions and PjitFunctions (e.g., jnp.dot).
• Qwix uses code-object patching to hook these functions during the initial JIT trace.
• PjitFuncitons will use simple attribute patching because code object is bypassed in jitted environments.
• Removing disable_jit for XLA models(large) significantly reduces compile times.

ODML Quantization (ODML Providers):

• Disable_jit by default.
• Intercepts high-level Functions, PjitFunctions (e.g., jnp.dot), and low-level JAX primitives via Primitive.bind for fine-grained metadata propagation.
• Visibility: When JIT is enabled, JAX dispatches primitives through a C++ path that bypasses Python-level bind overrides. disable_jit=True is required to force the Eager path where these primitive calls are visible to the Qwix interceptor.
• Recursion Avoidance: For PjitFunctions like jnp.dot, the interceptor calls the globally patched version. disable_jit ensures the recursion-guarded 'original' call runs eagerly, preventing nested JIT traces that might lead to inconsistent metadata states.
• Name Stability: Forcing Eager mode ensures consistent op ID generation (e.g., 'dot') between QAT and conversion, avoiding the shift to primitive names (e.g., 'dot_general') that would occur if the high-level Python wrappers were bypassed during JIT tracing.