Support compute_dtype in HowToQuantize in QWIX

[copybara-service]

Support compute_dtype in HowToQuantize in QWIX

Decouples the logical quantization format (qtype) from the physical compute/storage representation (compute_dtype).

Rationale and use cases:

• Preserve experimental accumulation precision: Allows simulation and ablation studies to explicitly control the matrix multiplication accumulation path (e.g. evaluating hardware FP8 accumulation versus BF16/FP32 dequantization) independently of the quantization discretization grid.
• Hardware acceleration for sub-byte and integer formats: Accelerates INT4 and FP4 inputs on hardware with native FP8 matrix multiplication units (e.g. TPU v6e / Ghostfish and modern GPUs). Because all 16 values of INT4 ([-8, 7]) and all 15 finite values of FP4 (E2M1) are exact representable values in float8_e4m3fn (zero numerical precision loss), setting compute_dtype=jnp.float8_e4m3fn allows INT4, FP4, and multi-pass INT4 (emulating INT8) to flow directly through native hardware FP8 GEMM units without requiring specialized sub-byte hardware kernels.

Changes:

• Add compute_dtype field to HowToQuantize and quantize_api.
• In quantize_with_scale_zero_point, cast qvalue (and zero_point if present) to compute_dtype after discretization while preserving the logical qtype on QArray.
• Add unit test in qarray_test verifying logical qtype preservation and compute_dtype casting.