This should be easy to implement by using von Mises Fisher distribution (or the qualitatively similar Projected Normal distribution) instead of the Normal distribution for sampling a latent angle. One will also need to modify the outputs of the encoder NN such that it outputs a "concentration" parameter in addition to z_loc and z_scale.
This should be easy to implement by using von Mises Fisher distribution (or the qualitatively similar Projected Normal distribution) instead of the Normal distribution for sampling a latent angle. One will also need to modify the outputs of the encoder NN such that it outputs a "concentration" parameter in addition to z_loc and z_scale.