Hi Michael,
Thanks for your reply!
The problem ended up being that grid_sample performs an inverse warping, which means that passing an affine_grid for the matrix A actually corresponds to the transformation A^(-1). So in my example above, the transformation with B followed by A actually corresponds to A^(-1)B^(-1) = (BA)^(-1), which means I should use C = BA and not C = AB as would be the case if the original matrix was used by grid_sample.
The reason why it worked in 2D and not 3D is because 2D rotation matrices commute whereas 3D rotation matrices do not.
It looks like you basically found the same fix by using world coordinates!
Thanks again for your replies and good to hear you fixed your problems too.