HOWTO: Per example gradient

[rolandgvc]

This howto exemplifies how per-example gradients can be retrieved using jax.vmap.

[codecov-commenter]

Codecov Report

Merging #306 into master will not change coverage.
The diff coverage is n/a.

@@           Coverage Diff           @@
##           master     #306   +/-   ##
=======================================
  Coverage   79.33%   79.33%           
=======================================
  Files          34       34           
  Lines        2255     2255           
=======================================
  Hits         1789     1789           
  Misses        466      466           

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 8901bc0...0134435. Read the comment docs.

[avital]

Need to check compute_metrics. It probably computes the loss a second time. Instead, we should use jax.value_and_grad where vmap_loss_grad is computed.

[avital]

We've redone our HOWTO system. Here's an example of the new format:

Rendered in ReadTheDocs: https://flax.readthedocs.io/en/latest/howtos/ensembling.html
Document source: https://github.com/google/flax/blob/master/docs/howtos/ensembling.rst (see the :testsetup, :testcode and :codediff directives)

We'd happy take pull requests revamping this HOWTO to the new system.

[avital]

Looks like this PR is inactive, so I'll close it. We're tracking the HOWTO request here: #858