Run the Jax solver on Google Colab GPU

[wigging]

I'm testing the PyBaMM Jax solver in a Google Colab notebook using the following example:

import matplotlib.pylab as plt
import pybamm

# Load the SPMe model
model = pybamm.lithium_ion.SPMe(options={'timescale': 1.0})
model.convert_to_format = 'jax'
model.events = []

# Create geometry
geometry = model.default_geometry

# Load parameter values and process model and geometry
param = model.default_parameter_values
param.process_model(model)
param.process_geometry(geometry)

# Set the mesh and discretize the model
mesh = pybamm.Mesh(geometry, model.default_submesh_types, model.default_var_pts)
discrete = pybamm.Discretisation(mesh, model.default_spatial_methods)
discrete.process_model(model)

# Solve model and get solution object
solver = pybamm.JaxSolver()
solution = solver.solve(model, [0, 3600])

# Plot results
_, ax = plt.subplots()
ax.plot(solution['Time [s]'].entries, solution['Terminal voltage [V]'].entries)
ax.set_xlabel('Time [s]')
ax.set_ylabel('Voltage [V]')

plt.show()

The version of Jax on Google Colab is v0.3.1 which interfaces with the Nvidia GPU that is available through the notebook environment. However, PyBaMM requires Jax v0.2.12 and when I use the pybamm_install_jax command it appears to restrict Jax to use only the CPU. Is there a way to use the PyBaMM Jax solver with the Google Colab GPU? Also, there are no PyBaMM Jax solver examples in the GitHub repo or documentation. So let me know if there's something else I need to implement to properly use the Jax solver in my example.

[valentinsulzer]

Jax constantly makes breaking updates hence the pinned version. I don't see any problem with changing the pinned version to 0.3.1, but to do so will likely require fixing bugs that arise in JaxSolver. I agree we should add an example

[wigging]

I can submit my example to the repository if you think it's suitable for sharing with others.

[wigging]

But related to running on the GPU, is the PyBaMM Jax solver restricted to CPU?

[valentinsulzer]

No idea, question for @martinjrobins

[wigging]

@martinjrobins Do you know if the PyBaMM Jax solver is restricted to the CPU or can it utilize the GPU too?