equi2equi TODO: Implement numpy Implement torch Implement torch with batches Fix rotation axis Implement c++ with cuda