Fix test.sh, which uses set -u, so that it works when JAX_VERSION is unset.

[carlosgmartin]

Currently, running

sh ./test.sh

as described in the readme yields the following error:

./test.sh: line 45: JAX_VERSION: unbound variable

This is because the script uses nounset (set -u), which treats an unset variable as an error (not an empty string).

This forces one to do

JAX_VERSION= sh ./test.sh

This PR fixes test.sh so that it works when JAX_VERSION is unset.

[fabianp]

Thanks!