3c082fc0821ff9f9b1284a486a5cf60f9aec238e,examples/loss_and_gradient_se3.py,,grad,#Any#Any#Any#Any#,39

Before Change

        differential = gs.zeros((1, 6, 7))

        upper_left_block = gs.zeros((1, 3, 4))
        lower_right_block = gs.zeros((1, 3, 3))
        quat_scalar = y_pred[:, :1]
        quat_vec = y_pred[:, 1:4]

        quat_vec_norm = gs.linalg.norm(quat_vec, axis=1)
        quat_sq_norm = quat_vec_norm ** 2 + quat_scalar ** 2
        // TODO(nina): check that this sq norm is 1?

        quat_arctan2 = gs.arctan2(quat_vec_norm, quat_scalar)
        differential_scalar = - 2 * quat_vec / (quat_sq_norm)
        differential_vec = (2 * (quat_scalar / quat_sq_norm
                                 - 2 * quat_arctan2 / quat_vec_norm)
                            * gs.outer(quat_vec, quat_vec) / quat_vec_norm ** 2
                            + 2 * quat_arctan2 / quat_vec_norm * gs.eye(3))

        upper_left_block[0, :, :1] = differential_scalar.transpose()
        upper_left_block[0, :, 1:] = differential_vec
        lower_right_block[0, :, :] = gs.eye(3)

        differential[0, :3, :4] = upper_left_block
        differential[0, 3:, 4:] = lower_right_block

        grad = gs.matmul(grad, differential)

After Change

        lower_right_block = gs.eye(3)
        lower_left_block = gs.zeros((3, 4))

        top = gs.hstack((upper_left_block, upper_right_block))
        bottom = gs.hstack((lower_left_block, lower_right_block))

        differential = gs.vstack((top, bottom))
        differential = gs.expand_dims(differential, axis=0)

        grad = gs.einsum("ni,nij->ni", grad, differential)

    grad = gs.squeeze(grad, axis=0)

Instances