Replace `jax.sharding.use_mesh` with `jax.set_mesh`. `jax.set_mesh` can act as a global setter or a context manager. #133

	name: Run Unit Tests

	on:
	push:
	branches: [ "main", "master" ]
	pull_request:
	branches: [ "main", "master" ]
	workflow_dispatch:

	jobs:
	run_tests:
	name: Run Tests
	runs-on: ubuntu-latest
	steps:
	- name: Checkout code
	uses: actions/checkout@v4
	- name: Install Dependencies
	run: \|
	python -m pip install --upgrade pip
	pip install -r requirements.txt # If you have a requirements.txt file
	- name: Run Pytest
	run: \|
	export KERAS_BACKEND=jax
	python -m pytest -v

Provide feedback