Change metric compute() return type to jnp.ndarray to match upcoming … #6037

	name: build

	on: [push]

	jobs:
	build:
	runs-on: ubuntu-latest
	steps:
	- uses: actions/checkout@v2
	- name: Set up Python
	uses: actions/setup-python@v4
	with:
	python-version: '3.9.x'
	cache: 'pip'
	cache-dependency-path: setup.py
	- name: Install dependencies
	run: \|
	pip install -e .[test] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
	- name: Test with pytest
	# TODO(adarob): Re-enable once tests are updated.
	run: \|
	export FLAX_LAZY_RNG=no
	pytest
	# The below step just reports the success or failure of tests as a "commit status".
	# This is needed for copybara integration.
	- name: Report success or failure as github status
	if: always()
	shell: bash
	run: \|
	status="${{ job.status }}"
	lowercase_status=$(echo $status \| tr '[:upper:]' '[:lower:]')
	curl -sS --request POST \
	--url https://api.github.com/repos/${{ github.repository }}/statuses/${{ github.sha }} \
	--header 'authorization: Bearer ${{ secrets.GITHUB_TOKEN }}' \
	--header 'content-type: application/json' \
	--data '{
	"state": "'$lowercase_status'",
	"target_url": "https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }}",
	"description": "'$status'",
	"context": "github-actions/build"
	}'

Provide feedback