[LSC] change uses of jax.random.KeyArray and jax.random.PRNGKeyArray … #108

	name: Unittests & Auto-publish

	# Allow to trigger the workflow manually (e.g. when deps changes)
	on: [push, workflow_dispatch]

	jobs:
	test-job:
	runs-on: ubuntu-latest
	timeout-minutes: 10

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

	steps:
	- uses: actions/checkout@v3

	# Install deps
	- uses: actions/setup-python@v4
	with:
	python-version: 3.8
	# Uncomment to cache of pip dependencies (if tests too slow)
	# cache: pip
	# cache-dependency-path: '**/pyproject.toml'

	- run: pip --version
	- run: pip install -e .[test]
	- run: pip freeze

	# Run tests (in parallel)
	- name: amos_test
	run: python3 jestimator/amos_test.py

	- name: amos_helper_test
	run: python3 jestimator/amos_helper_test.py

	# Auto-publish when version is increased
	publish-job:
	# Only try to publish if:
	# * Repo is self (prevents running from forks)
	# * Branch is `main`
	if: \|
	github.repository == 'google-research/jestimator'
	&& github.ref == 'refs/heads/main'
	needs: test-job # Only publish after tests are successful
	runs-on: ubuntu-latest
	permissions:
	contents: write
	timeout-minutes: 10

	steps:
	# Publish the package (if local `__version__` > pip version)
	- uses: etils-actions/pypi-auto-publish@v1
	with:
	pypi-token: ${{ secrets.PYPI_API_TOKEN }}
	gh-token: ${{ secrets.FOR_JESTIMATOR_RELEASE }}
	parse-changelog: true

Provide feedback