Make ModCirc Jax compliant

Word on the street is that [Autoemulate](https://github.com/alan-turing-institute/autoemulate) is moving to become JAX compliant.

[JAX](https://docs.jax.dev/en/latest/index.html) is NumPy-compatible Python library that adds automatic differentiation, GPU acceleration, and JIT compilation to scientific computing tasks, making it faster and more powerful than SciPy for machine learning and complex numerical computations, though less comprehensive in terms of built-in scientific algorithms. More information on JAX at the bottom.

For the DT pipline which is inporocess ensuring AD from beggining to end is vital for any method that one can hope to employ, involving gradients and this package.

**JAX**
JAX stands out as a powerful tool for modeling and simulation, particularly in the realms of neural networks and probabilistic computing. At its core, JAX's automatic differentiation capabilities make it exceptionally well-suited for implementing custom neural networks and gradient-based optimization tasks. This feature seamlessly extends to probabilistic modeling, where JAX's built-in support for probability distributions and random sampling makes it natural to implement methods like MCMC and variational inference.

The library's performance characteristics are equally impressive. Through JIT compilation, models run at near-native speeds, while automatic parallelization across GPUs and TPUs enables scaling to larger problems. JAX's vectorization capabilities, implemented through vmap, allow efficient batching of operations, making it possible to handle large-scale simulations that would be impractical in pure Python environments.


Start by auditing the ModCirc DT pipeline and its modelling methods to identify operations that prevent JAX compatibility and end-to-end automatic differentiation. The issue names no files or tests, so first map the relevant entry points and establish how gradient support will be verified. Done means the applicable methods run through the pipeline with automatic differentiation.

Written by the indexing model from the issue text.