The JAX/FLAX version of this model was
contributed by kamalkraj.