minor tweaks for improved use of jax API #164

	name: Build and Test
	permissions: read-all
	on:
	push:
	branches:
	- main
	pull_request:
	branches:
	- main
	concurrency:
	group: ${{ github.workflow }}-${{ github.head_ref \|\| github.ref }}
	cancel-in-progress: true
	jobs:
	build-and-test:
	runs-on:
	labels: ubuntu-20.04-16core
	steps:
	- name: Check out repository code
	uses: actions/checkout@8e5e7e5ab8b370d6c329ec480221332ada57f0ab # pin@v3

	- name: Cache bazel build artifacts
	uses: actions/cache@88522ab9f39a2ea568f7027eddc7d8d8bc9d59c8 # pin@v3.3.1
	with:
	path: \|
	~/.cache/bazel
	key: ${{ runner.os }}-bazel-${{ hashFiles('.bazelversion', '.bazelrc', 'WORKSPACE', 'requirements_dev.txt') }}

	- name: "Run `bazel build`"
	run: \|
	bazel build //...

	- name: "Run `bazel test`"
	run: \|
	bazel test \
	--test_output=errors \
	--test_size_filters=small \
	--test_timeout=1800 \
	--experimental_ui_max_stdouterr_bytes=10485760 \
	//...

Provide feedback