Jax is definitely installed, yet JaxSolver does not run

[cyasing]

Just as the title says, I can't seem to run JaxSolver() even with Jax installed.

I tried the pybamm_install_jax, but that gives this: NameError: name 'pybamm_install_jax' is not defined.
What version of Jax is compatible with PyBaMM?

[Saransh-cpp]

Supported jax and jaxlib versions:

PyBaMM/pybamm/util.py

Lines 24 to 26 in 4f03df4

	# versions of jax and jaxlib compatible with PyBaMM
	JAX_VERSION = "0.2.12"
	JAXLIB_VERSION = "0.1.70"

I think you are executing pybamm_install_jax in a python file? pybamm_install_jax should be executed in your terminal.

Also, I see that the docstring is a bit misleading. The docstring uses >>>, misleading users into executing it in the python interpreter (a PR fixing it would be very welcome!) -

PyBaMM/pybamm/util.py

Lines 330 to 341 in 4f03df4

	def install_jax(arguments=None): # pragma: no cover
	"""
	Install compatible versions of jax, jaxlib.
	Command Line Interface:
	-----------------------
	>>> pybamm_install_jax
	optional arguments:
	-h, --help show help message
	-f, --force force install compatible versions of jax and jaxlib
	"""

[Saransh-cpp]

Also, https://pybamm.readthedocs.io/en/latest/source/api/util.html misses the API documentation for install_jax. The function can be added here: https://github.com/pybamm-team/PyBaMM/blob/develop/docs/source/api/util.rst.

[cyasing]

Thanks! That worked. However, I am encountering this error now:
raise TracerArrayConversionError(self) jax._src.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ShapedArray(float64[1,1])>with<DynamicJaxprTrace(level=1/2)> (https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError)

I used method=BDF, because for some reason the default one does not work. I had to use 'model.events=[]' to remove events before attempting to solve using BDF.