perf: try using fori_loop but unrolling everything #824

	name: Python package

	env:
	PY_COLORS: "1"

	concurrency:
	group: ${{ github.workflow }}-${{ github.ref }}
	cancel-in-progress: true

	on:
	push:
	branches:
	- main
	pull_request:
	workflow_dispatch: null

	jobs:
	build:
	runs-on: ubuntu-latest
	strategy:
	fail-fast: false
	matrix:
	python-version: ["3.10", "3.11", "3.12"]
	group: [1, 2, 3, 4]

	steps:
	- uses: actions/checkout@v4

	- name: Set up Python ${{ matrix.python-version }}
	uses: actions/setup-python@v5
	with:
	python-version: ${{ matrix.python-version }}

	- name: Install dependencies
	run: \|
	python -m pip install --upgrade pip
	python -m pip install pytest pytest-codspeed pytest-split pytest-randomly
	python -m pip install .
	# temp pin until 0.5 is on conda
	python -m pip install "jax<0.5.0"

	- name: Test with pytest
	run: \|
	git submodule update --init --recursive
	pytest -v --durations=0 \
	-k "not test_fpack" \
	--randomly-seed=42 \
	--splits=4 --group=${{ matrix.group }} --splitting-algorithm least_duration

	build-status:
	needs: build
	runs-on: ubuntu-latest

	steps:
	- name: check status
	run: \|
	echo "Builds all passed!"

Provide feedback