Here's a vibe coding experiment in porting to JAX, where I was able to guide the agent from an expert perspective and point out bugs based just on output / results
https://github.com/verdverm/pge-jax