I have searched the forums for this and ended up baking my own recipe: just a 2 cents abstraction in jit.gen if anyone needs it (there is Jasch's 3Dcartopol object which works great, but I needed a thing on matrices with more than 1*1 dim)

I am sure there is a more elegant way, but may be a starting point