[JAX] Update simple accuracy tests to support optional XLA optimizations

[wpietka]

Type of Change

bug_fix

Description

With xla_cpu_experimental_onednn_custom_call=true, layer's output may differ due to different computation path after JIT optimization, which results in failing test. This PR adds alternative way of computing expected values in the test so it matches the ones that we have with xla_cpu_experimental_onednn_custom_call

Additionally XLA with xla_allow_excess_precision=true (Default) increases computation accuracy by dropping unnecessary casts from higher to lower precision floats. This PR adjusts tests also for that

[Copilot]

🟡 Changes recommended

The oneDNN reference loses precision during bfloat16 zero-point subtraction, making expected outputs inaccurate.

1 open issue

• Convert bfloat16 codes and zero point before subtraction · New

What changed in this PR

Updates the JAX accuracy-test reference calculations to account for oneDNN execution and XLA excess precision.

Changes:

• Parses XLA flags to select the reference computation path.
• Adds oneDNN-style accumulation and scaling.
• Adjusts intermediate precision and final output casting.

File	Description
test/​jax/​jax_test_utility.py	Adds flag-aware quantized dense reference calculations.

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

Compiler-dependent numerical assumptions need validation against JIT outputs with oneDNN enabled and excess precision both enabled and disabled.

No open issues 🎉

1 previously missed issue

In code that hasn't changed since last review

Add accuracy coverage for oneDNN and excess-precision branches

test/​jax/​jax_test_utility.py:150

The existing accuracy test compares this helper with JIT output, but it does not vary XLA_FLAGS, and .azure-pipelines/scripts/ut/run_3x_jax.sh:33 runs it only once with inherited flags. This leaves the new oneDNN and excess-precision branches without explicit regression coverage. Add separate accuracy-test processes for the default route and for oneDNN enabled with xla_allow_excess_precision=true and false. Set the flags before JAX initializes so each expected-output path is checked against the corresponding compiled model.

1 resolved since last review

• Convert bfloat16 codes and zero point before subtraction

_{Review effort: Balanced}