Riemannian SGD for spherical embeddings?

1

I would like to run SGD on the unit sphere in order to learn spherical word embeddings (vectors of norm 1).

For general statements:

and more specifically, Equation (7) in Meng et al. 2019:

https://papers.nips.cc/paper/9031-spherical-text-embedding.pdf

Is there a simpler way of implementing the update rule (eq. (7) in Meng et al. 2019) paper than by re-writing C code similar to the update rule for plain SGD as in

2

Update: problem solved. I managed running the algorithm by writing the riemannian gradient steps in python using gradients computed inside of autograd.record() and setting values of modified rows weights only manually inside the training loop. Note: the set_data method does not allow to use sparse gradients efficiently, I had to implement this myself