Refactor rotations

Created by: APJansen

This is a small refactor of the rotation layer, now requiring a rotation_axis argument that will be used to take a tensor product with that axis from the input and the 0th axis of the rotation matrix. (the resulting axis is by default put at the end, so this is transposed back afterward). It's more intuitive than the previous axes which specified the amount of axes to contract (starting from the end of the input and the start of the rotation matrix). More importantly, it now requires no change when axes are added to the pdf.

I've added a regression test to make sure results are still the same.