Update attention.py #447

	name: Tests

	on:
	- push
	- pull_request

	jobs:
	test:
	runs-on: ${{ matrix.os }}
	strategy:
	matrix:
	os: [ubuntu-latest]
	python-version: ["3.8", "3.9", "3.10", "3.11"]

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

	- name: Install dependencies
	run: \|
	python -m pip install --upgrade pip
	python -m pip install pytreeclass>=0.4.0
	python -m pip install keras_core>=0.1.1
	python -m pip install pytest wheel optax jaxlib coverage kernex
	- name: Pytest Check
	run: \|
	python -m pip install .
	coverage run -m pytest tests

	- name: Codecov
	uses: codecov/[email protected]

Provide feedback