[JAX] Add per-channel quantization for kernels

[wpietka]

Type of Change

Feature

Description

• Implements per-channel quantization along specified axis for get_q_params()
• Uses above implementation to add per-channel quantization for dense layers' kernels

Expected Behavior & Potential Risk

Increased quantization accuracy

How has this PR been tested?

Ran smoke tests

Dependency Change?

No changes

[Copilot]

Pull request overview

This PR extends the JAX quantization utilities and layer wrappers to support per-channel weight quantization scales (along a specified axis) and applies it to Dense/EinsumDense/Conv2D kernels, with corresponding test expectation updates.

Changes:

• Add axis support to get_scale() / get_q_params() for per-tensor vs per-channel scale computation.
• Update static and dynamic quantized Dense/EinsumDense/Conv2D layers to compute/store per-channel w_scale (and infer axis for EinsumDense via equation parsing).
• Update JAX accuracy test utilities to compute expected Dense weight scales per output channel.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
test/jax/test_accuracy.py	Updates a comment (quantization → inference) in the test flow.
test/jax/jax_test_utility.py	Adjusts expected Dense weight-scale computation to per-channel.
neural_compressor/jax/utils/utility.py	Adds axis parameter support for scale/params computation (per-channel scaling).
neural_compressor/jax/quantization/layers_static.py	Uses per-channel w_scale for static quantized kernels; adds quant-axis declarations/inference.
neural_compressor/jax/quantization/layers_dynamic.py	Uses per-channel w_scale for dynamic quantized kernels; adds quant-axis declarations/inference and tweaks load detection.

Comments suppressed due to low confidence (1)

neural_compressor/jax/utils/utility.py:226

• get_scale() adds a new axis parameter but the docstring Args: section doesn't document it, which makes the API harder to use correctly.

def get_scale(orig_weight, dtype=ml_dtypes.float8_e4m3, compute_dtype=jnp.float32, axis=None):
    """Compute the quantization scale for a weight tensor.

    Args:
        orig_weight (jnp.ndarray): Weight tensor to analyze.
        dtype (jnp.dtype): Target quantized dtype.
        compute_dtype (jnp.dtype): dtype for scale computation.

[anko-intel]

consider making this if to calculate axis, and shape and have max/min call the same