Use doctest for new HOWTOs, instead of our old diff system
I split this PR from https://github.com/google/flax/pull/771
as it's orthogonal
This allows us to write embedded code snippets inside .rst files,
that are validated during the doc build. This seems like a better
alternative than using Colabs as documentation, except for cases
where we expicitly want people to fork and run a notebook.
