Updates make_functional to use the new improved variants. The new
variants are superior in every way so we're replacing the previous
variants with this.
If someone wants the older variants, they can be found at:
- make_functional_with_buffers_deprecated_v1
- make_functional_deprecated_v1
make_functional*_v2 is superior to the older make_functional. This PR
has all of our examples use it.
This PR also adds a "combine_state_for_ensemble(models)" API.
Coming soon: We're probably going to break BC on make_functional and
replace it with make_functional_v2. That's the nice thing about being a
prototype, we don't have to worry about BC too much.
Introduces a `functional_init` and `functional_init_with_buffers` that
lets one initialize an ensemble of modules more easily than before. This
was done in the sprit of make_functional: the API still looks awkward,
especially when buffers are involved.