Update JAX to 0.5.0

[sramakintel]

Description

Updates JAX XPU support to 0.5.0

Related Issue

Changes Made

• The code follows the project's coding standards.
• No Intel Internal IP is present within the changes.
• The documentation has been updated to reflect any changes in functionality.

Validation

• I have tested any changes in container groups locally with test_runner.py with all existing tests passing, and I have added new tests where applicable.

[github-actions]

Dependency Review

✅ No vulnerabilities or license issues or OpenSSF Scorecard issues found.

OpenSSF Scorecard

Package	Version	Score	Details
pip/flax	0.8.5	🟢 6.2	Details Check Score Reason Maintained 🟢 10 30 commit(s) and 21 issue activity found in the last 90 days -- score normalized to 10 Code-Review 🟢 9 Found 17/18 approved changesets -- score normalized to 9 Packaging ⚠️ -1 packaging workflow not detected CII-Best-Practices ⚠️ 0 no effort to earn an OpenSSF best practices badge detected Dangerous-Workflow 🟢 10 no dangerous workflow patterns detected Token-Permissions ⚠️ 0 detected GitHub workflow tokens with excessive permissions License 🟢 10 license file detected Binary-Artifacts 🟢 10 no binaries found in the repo Fuzzing ⚠️ 0 project is not fuzzed Signed-Releases ⚠️ -1 no releases found Branch-Protection ⚠️ -1 internal error: error during branchesHandler.setup: internal error: githubv4.Query: Resource not accessible by integration Security-Policy 🟢 10 security policy file detected Pinned-Dependencies ⚠️ 0 dependency not pinned by hash detected -- score normalized to 0 SAST ⚠️ 2 SAST tool is not run on all commits -- score normalized to 2 Vulnerabilities 🟢 6 4 existing vulnerabilities detected
pip/intel-extension-for-openxla	0.5.0	🟢 6.4	Details Check Score Reason Packaging ⚠️ -1 packaging workflow not detected Dangerous-Workflow 🟢 10 no dangerous workflow patterns detected Maintained 🟢 8 10 commit(s) and 0 issue activity found in the last 90 days -- score normalized to 8 Security-Policy 🟢 10 security policy file detected SAST ⚠️ 0 no SAST tool detected Code-Review ⚠️ 0 Found 0/30 approved changesets -- score normalized to 0 Token-Permissions 🟢 10 GitHub workflow tokens follow principle of least privilege Binary-Artifacts 🟢 10 no binaries found in the repo CII-Best-Practices ⚠️ 0 no effort to earn an OpenSSF best practices badge detected Pinned-Dependencies ⚠️ 1 dependency not pinned by hash detected -- score normalized to 1 License 🟢 10 license file detected Signed-Releases ⚠️ -1 no releases found Fuzzing ⚠️ 0 project is not fuzzed Branch-Protection ⚠️ -1 internal error: error during branchesHandler.setup: internal error: githubv4.Query: Resource not accessible by integration Vulnerabilities 🟢 10 0 existing vulnerabilities detected
pip/jax	0.4.30	Unknown	Unknown
pip/jaxlib	0.4.30	Unknown	Unknown
pip/optax	0.2.4	Unknown	Unknown

Scanned Files

• jax/requirements.txt

[github-advanced-security]

This pull request sets up GitHub code scanning for this repository. Once the scans have completed and the checks have passed, the analysis results for this pull request branch will appear on this overview. Once you merge this pull request, the 'Security' tab will show more code scanning analysis results (for example, for the default branch). Depending on your configuration and choice of analysis tool, future pull requests will be annotated with code scanning analysis results. For more information about GitHub code scanning, check out the documentation.

[github-actions]

Integration Test Results

Groups Tested: jax/tests

Results

Test-Group	Test	Status
jax/tests	jax-import-idp	PASS
jax/tests	jax-import-pip	PASS
jax/tests	jax-import-jupyter-idp	PASS
jax/tests	jax-import-jupyter-pip	PASS
jax/tests	jax-xpu-example-idp	PASS
jax/tests	jax-xpu-example-pip	PASS

Overall Result: PASS ✅