exp_backward_kernel has wrong offset for grad?

[zjcs]

As I know, exp operation will convert lie vector to lie group, so the grad of exp should be on the tangent space, a Group::K vector, while the offset in exp_backward_kernel is Group::N. Is it wrong?

Current implementation:

        Grad dX(grad + i*Group::N);

Modified implementation:

        Grad dX(grad + i*Group::K);

Here is the code in lietorch_gpu.cu:

template <typename Group, typename scalar_t>
EIGEN_DEVICE_FUNC void exp_forward_kernel(const scalar_t* a_ptr, scalar_t* X_ptr) {
    // exponential map forward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;
    
    GPU_1D_KERNEL_LOOP(i) {
        Tangent a(a_ptr + i*Group::K);
        Eigen::Map<Data>(X_ptr + i*Group::N) = Group::Exp(a).data();
    }
}

template <typename Group, typename scalar_t>
EIGEN_DEVICE_FUNC void exp_backward_kernel(const scalar_t* grad, const scalar_t* a_ptr, scalar_t* da) {
    // exponential map backward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Grad = Eigen::Matrix<scalar_t,1,Group::K>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;

    GPU_1D_KERNEL_LOOP(i) {
        Tangent a(a_ptr + i*Group::K);
        Grad dX(grad + i*Group::N);
        Eigen::Map<Grad>(da + i*Group::K) = dX * Group::left_jacobian(a);
    }
}