Once we pin a relatively recent version of PyMC we can remove this, as it is now part of Aesara https://github.com/pymc-labs/roche-preclinical-hd/blob/24b710b7df0e929eba54c4397a9bdca5cd756c3d/preclinpack/jax_utils.py#L20-L35