Cyclic constraints on nn.Modules output

Hello community,

I’m trying to put cyclic constraints on the output of a normalising flow but I am uncertain of how this could be achieved in PyTorch.

Mathematically I would like the flow to satisfy these constraints:

  • f(-π)=-π
  • f(π)=π
  • ∇f(π) = ∇f(-π)

Any idea how this could be done without adding these as soft constraints in the objective?

reference: Rezende, Danilo Jimenez, et al. “Normalizing flows on tori and spheres.” International Conference on Machine Learning. PMLR, 2020.