Attempt at reproducing "exponential-map-sum-of-radial-flow" results (EMSRE Table 1) of Normalizing Flows on Tori and Spheres by Rezende et al., done in JAX.

The target density on $\mathbb{S}^2$ is

To train the flow with hyperparameters $^\dagger$ specified in the paper, run below adjusting N and K as desired.

$ python3 main.py --N=1 --K=12

Comparing the authors' results (average) and ours (single run as is), with the convention of theirs / ours:

Scenario N=24, K=1 was also attempted, though unsuccessfully in 20,000 iterations claimed.

$^\dagger$ OTOH, parameter initialisation and constraint enforcing are not specified so had to be assumed.