jax

92e5f93a - tweak docstrings in mnist examples