fix shard again? this time with input from yash #292

	name: Run Tests

	on: [push, pull_request]

	jobs:
	build:

	runs-on: ubuntu-latest

	steps:
	- uses: actions/checkout@v3
	- name: Set up Python 3.10.11
	uses: actions/setup-python@v4
	with:
	python-version: 3.10.11
	- name: Install dependencies
	run: \|
	python -m pip install --upgrade pip
	pip install flake8 pytest
	pip install --upgrade "jax[cpu]==0.4.25" "jaxlib[cpu]==0.4.25"
	pip install .[dev]
	- name: Test with pytest
	run: \|
	XLA_FLAGS=--xla_force_host_platform_device_count=8 PYTHONPATH=tests:src:. pytest tests

Provide feedback