grad of cond and while

[alexbw]

First and higher order derivatives of functions that use lax.cond and lax.while should be possible.

[alexbw]

@dougalm this includes scan now, correct?

[dougalm]

yes, working on scan (i.e. scan :: (a -> b -> a) -> a -> [b] -> [a]) too.

[matthewdhoffman]

+1 for scan! Once I've got that I can put together an HMC example PR.

[mattjj]

#415 will add lax.cond without reverse-mode autodiff (so it won't close this issue, but is related).

[gnecula]

We have now grad of cond, and jvp of while (but not grad of while).