Ben Cassese

This is not currently affecting any active projects, but I do still think it's a concern for datasets near the saturation limit.

Dittoing @ekpass, that change was added to jax/_src/interpreters/ad.py in [this commit](https://github.com/jax-ml/jax/commit/ba9b2ca5f6917444a82792a773f26e818366c9c2), right before the release of jax v0.5.0 at the beginning of this year. Downgrading to 0.4.38 solved the problem...