Fix jax vit unit tests

[bkowalskiINTEL]

This PR changes how the accuracy is evaluated in jax vit unit tests.
Previously: exact top-5 match with hardcoded labels
Now: exact top-1 match with label calculated on the original model + top-1 label propability error threshold

Also, now the input images for vit are in uint8 instead of float32

[Copilot]

Pull request overview

Updates the JAX ViT quantization unit test to use a more robust accuracy check by comparing top‑1 predictions (label + probability) against the original (unquantized) model, and switches test image inputs to uint8 to match the ViT example/input pipeline.

Changes:

• Switch ViT test evaluation from hardcoded top‑5 label matching to top‑1 label matching plus a probability-delta threshold.
• Add int8 to the tested quantization dtypes for the ViT classification test matrix.
• Simplify load_image() in JAX test utilities to always return uint8 (no optional normalization path).

Reviewed changes

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

File	Description
test/jax/test_vit_quantization.py	Adjusts ViT quantization test to compare top‑1 label and probability between original and quantized models; adds int8 coverage.
test/jax/jax_test_utility.py	Changes load_image() to always return unnormalized uint8 image tensors.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

[wpietka]

LGTM

[anko-intel]

LGTM