Support torch.distributions for fp16 operations

Been trying to get a VAE to work with APEX.

Ran into a few issues along the way and tried to pull in lessons learned from https://github.com/chainer/chainer/issues/6168 (eg: no exp on logvar as in this example), however something that seems to be missing is proper wrapping of torch.distributions including operands such as the D.kl_divergence.

As an example here is the Normal : Normal KLD which I can't seem get to not blow up.
The only thing I see that could tentatively blow up here is the .pow(2).
(Note that the division of by .scale is not an issue because people typically add some tolerance to the scale, eg: scale += 1e-6 which I have tried to truncate to 3 bits for fp16, but to no avail.)

@register_kl(Normal, Normal)
def _kl_normal_normal(p, q):
    var_ratio = (p.scale / q.scale).pow(2)
    t1 = ((p.loc - q.loc) / q.scale).pow(2)
    return 0.5 * (var_ratio + t1 - 1 - var_ratio.log())

I have got to the point where I can train a VAE without the KL term (i.e. basically an autoencoder with reparameterization but no regularization on the latent variable) but the KL is still an issue.

I have already tried -o1 and -o2 where -o2 fails to work due to an error with returning fp32 values during reparameterization. Forcibly type-casting as follows does not work:

    def _reparametrize_gaussian(self, mu, logvar, force=False):
        """ Internal member to reparametrize gaussian.

        :param mu: mean logits
        :param logvar: log-variance.
        :returns: reparameterized tensor and param dict
        :rtype: torch.Tensor, dict

        """
        if self.training or force:  # returns a stochastic sample for training
            std = logvar.mul(0.5)  # .exp()
            eps = torch.zeros_like(logvar).normal_().type(std.dtype)
            nan_check_and_break(logvar, "logvar")
            reparam_sample = eps.mul(std).add_(mu)
            return reparam_sample, {'z': reparam_sample, 'mu': mu, 'logvar': logvar}

        return mu, {'z': mu, 'mu': mu, 'logvar': logvar}

Start with the torch.distributions Normal:Normal implementation and its register_kl entry point, then trace how APEX applies mixed-precision wrapping to the operations in _kl_normal_normal. Reproduce the VAE fp16 failure described in the issue, including the reparameterization path, and define done as a working KL term without the reported overflow or fp32 return error.