kl_div_to_std_normal

behavenet.fitting.losses.kl_div_to_std_normal(mu, logvar)

Compute element-wise KL(q(z) || N(0, 1)) where q(z) is a normal parameterized by mu, logvar.

Parameters

• mu (torch.Tensor) – mean parameter of shape (n_frames, n_dims)

• logvar – log variance parameter of shape (n_frames, n_dims)

Returns

KL divergence summed across dims, averaged across batch

Return type

torch.Tensor