pip prod(deps): bump jax from 0.4.31 to 0.4.32 #12195

	# This workflow will install Python dependencies, run tests and lint with a single version of Python
	# For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions

	name: CI - Python-based Testing

	on:
	# Trigger the workflow on push or pull request,
	# but only for the master branch
	push:
	branches:
	- main
	pull_request:

	jobs:
	build:

	runs-on: ubuntu-latest
	strategy:
	matrix:
	python-version: ['3.10', '3.11', '3.12']

	steps:
	- uses: actions/checkout@v4

	- name: Install native dependencies
	uses: awalsh128/cache-apt-pkgs-action@latest
	with:
	packages: mesa-vulkan-drivers
	version: 1.0

	- name: Set up Python
	uses: actions/setup-python@v5
	with:
	python-version: ${{ matrix.python-version }}
	cache: 'pip' # caching pip dependencies
	cache-dependency-path: \|
	setup.py
	requirements.txt

	- name: Upgrade pip
	run: \|
	pip install --upgrade pip

	- name: Install the package locally
	run: pip install -r requirements.txt

	- name: Test with pytest
	run: \|
	pytest -W error

	- name: Execute lit tests
	run: \|
	export PYTHONPATH=$(pwd)
	lit -v tests/filecheck/
	lit -v docs/Toy/examples/

Provide feedback