Predictive-Coding (JAX) Jax version implementation of Predictive Coding algorithms, a biological plausible alternative for back-propagation.