The Reparameterisation Trick

Sampling $z \sim \mathcal N(\mu, \sigma^2)$ isn't differentiable in $\mu$ and $\sigma$. The reparameterisation trick writes $z = \mu + \sigma\,\epsilon$ with $\epsilon \sim \mathcal N(0, 1)$, so gradients flow through $\mu$ and $\sigma$.

VAEs predict $\log\sigma^2$ ("logvar") for stability. Implement:

• sample_z(mu, logvar, eps): $\mu + e^{\frac12 \text{logvar}}\,\epsilon$ (the noise eps is passed in)
• sample_grads(mu, logvar, eps, dz): given upstream gradient dz, return (dmu, dlogvar)