How to use the geomstats.backend.cos function in geomstats

def __init__(self, n_meridians=40, n_circles_latitude=None,
                 points=None):
        if n_circles_latitude is None:
            n_circles_latitude = max(n_meridians / 2, 4)

        u, v = gs.meshgrid(
            gs.arange(0, 2 * gs.pi, 2 * gs.pi / n_meridians),
            gs.arange(0, gs.pi, gs.pi / n_circles_latitude))

        self.center = gs.zeros(3)
        self.radius = 1
        self.sphere_x = self.center[0] + self.radius * gs.cos(u) * gs.sin(v)
        self.sphere_y = self.center[1] + self.radius * gs.sin(u) * gs.sin(v)
        self.sphere_z = self.center[2] + self.radius * gs.cos(v)

        self.points = []
        if points is not None:
            self.add_points(points)

rot_vec = self.regularize(rot_vec)

        angle = gs.linalg.norm(rot_vec, axis=1)
        angle = gs.to_ndarray(angle, to_ndim=2, axis=1)

        mask_0 = gs.isclose(angle, 0.)
        mask_not_0 = ~mask_0

        rotation_axis = gs.divide(
            rot_vec,
            angle
            * gs.cast(mask_not_0, gs.float32)
            + gs.cast(mask_0, gs.float32))

        quaternion = gs.concatenate(
            (gs.cos(angle / 2),
             gs.sin(angle / 2) * rotation_axis[:]),
            axis=1)

        return quaternion

Returns
        -------
        christoffel : array-like, shape=[..., contravariant index, 1st
                                         covariant index, 2nd covariant index]
            Christoffel symbols at point.
        """
        if self.dim != 2 or point_type != 'spherical':
            raise NotImplementedError(
                'The Christoffel symbols are only implemented'
                ' for spherical coordinates in the 2-sphere')

        point = gs.to_ndarray(point, to_ndim=2)
        christoffel = []
        for sample in point:
            gamma_0 = gs.array(
                [[0, 0], [0, - gs.sin(sample[0]) * gs.cos(sample[0])]])
            gamma_1 = gs.array([[0, gs.cos(sample[0]) / gs.sin(sample[0])],
                                [gs.cos(sample[0]) / gs.sin(sample[0]), 0]])
            christoffel.append(gs.stack([gamma_0, gamma_1]))

        christoffel = gs.stack(christoffel)
        if gs.ndim(christoffel) == 4 and gs.shape(christoffel)[0] == 1:
            christoffel = gs.squeeze(christoffel, axis=0)
        return christoffel

----------
        tait_bryan_angles : array-like, shape=[n_samples, 3]

        Returns
        -------
        rot_mat : array-like, shape=[n_samples, 3, 3]
        """
        n_tait_bryan_angles, _ = tait_bryan_angles.shape

        rot_mat = gs.zeros((n_tait_bryan_angles,) + (self.n,) * 2)
        angle_1 = tait_bryan_angles[:, 0]
        angle_2 = tait_bryan_angles[:, 1]
        angle_3 = tait_bryan_angles[:, 2]

        for i in range(n_tait_bryan_angles):
            cos_angle_1 = gs.cos(angle_1[i])
            sin_angle_1 = gs.sin(angle_1[i])
            cos_angle_2 = gs.cos(angle_2[i])
            sin_angle_2 = gs.sin(angle_2[i])
            cos_angle_3 = gs.cos(angle_3[i])
            sin_angle_3 = gs.sin(angle_3[i])

            column_1 = [[cos_angle_1 * cos_angle_2],
                        [cos_angle_2 * sin_angle_1],
                        [- sin_angle_2]]
            column_2 = [[(cos_angle_1 * sin_angle_2 * sin_angle_3
                          - cos_angle_3 * sin_angle_1)],
                        [(cos_angle_1 * cos_angle_3
                          + sin_angle_1 * sin_angle_2 * sin_angle_3)],
                        [cos_angle_2 * sin_angle_3]]
            column_3 = [[(sin_angle_1 * sin_angle_3
                          + cos_angle_1 * cos_angle_3 * sin_angle_2)],

- Y(angle_2) is a rotation of angle angle_2 around axis y.
        - Z(angle_3) is a rotation of angle angle_3 around axis z.
        """
        assert self.n == 3, ('The Tait-Bryan angles representation'
                             ' does not exist'
                             ' for rotations in %d dimensions.' % self.n)
        tait_bryan_angles = gs.to_ndarray(tait_bryan_angles, to_ndim=2)
        n_tait_bryan_angles, _ = tait_bryan_angles.shape

        rot_mat = gs.zeros((n_tait_bryan_angles,) + (self.n,) * 2)
        angle_1 = tait_bryan_angles[:, 0]
        angle_2 = tait_bryan_angles[:, 1]
        angle_3 = tait_bryan_angles[:, 2]

        for i in range(n_tait_bryan_angles):
            cos_angle_1 = gs.cos(angle_1[i])
            sin_angle_1 = gs.sin(angle_1[i])
            cos_angle_2 = gs.cos(angle_2[i])
            sin_angle_2 = gs.sin(angle_2[i])
            cos_angle_3 = gs.cos(angle_3[i])
            sin_angle_3 = gs.sin(angle_3[i])

            column_1 = [[cos_angle_2 * cos_angle_3],
                        [(cos_angle_1 * sin_angle_3
                          + cos_angle_3 * sin_angle_1 * sin_angle_2)],
                        [(sin_angle_1 * sin_angle_3
                          - cos_angle_1 * cos_angle_3 * sin_angle_2)]]

            column_2 = [[- cos_angle_2 * sin_angle_3],
                        [(cos_angle_1 * cos_angle_3
                          - sin_angle_1 * sin_angle_2 * sin_angle_3)],
                        [(cos_angle_3 * sin_angle_1

norm2 = norm_tangent_vec[mask_0]**2
        norm4 = norm2**2
        norm6 = norm2**3

        coef_1 = gs.assignment(
            coef_1,
            1. - norm2 / 2. + norm4 / 24. - norm6 / 720.,
            mask_0)
        coef_2 = gs.assignment(
            coef_2,
            1. - norm2 / 6. + norm4 / 120. - norm6 / 5040.,
            mask_0)

        coef_1 = gs.assignment(
            coef_1,
            gs.cos(norm_tangent_vec[mask_non0]),
            mask_non0)
        coef_2 = gs.assignment(
            coef_2,
            gs.sin(
                norm_tangent_vec[mask_non0]) /
            norm_tangent_vec[mask_non0],
            mask_non0)

        exp = (gs.einsum('...,...j->...j', coef_1, base_point)
               + gs.einsum('...,...j->...j', coef_2, proj_tangent_vec))

        return exp

def main():
    y_pred = gs.array([1., 1.5, -0.3, 5., 6., 7.])
    y_true = gs.array([0.1, 1.8, -0.1, 4., 5., 6.])

    loss_rot_vec = loss(y_pred, y_true)
    grad_rot_vec = grad(y_pred, y_true)

    logging.info('The loss between the poses using rotation '
                 'vectors is: {}'.format(loss_rot_vec))
    logging.info('The Riemannian gradient is: {}'.format(grad_rot_vec))

    angle = gs.array(gs.pi / 6)
    cos = gs.cos(angle / 2)
    sin = gs.sin(angle / 2)
    u = gs.array([1., 2., 3.])
    u = u / gs.linalg.norm(u)
    scalar = gs.array(cos)
    vec = sin * u
    translation = gs.array([5., 6., 7.])
    y_pred_quaternion = gs.concatenate([[scalar], vec, translation], axis=0)

    angle = gs.array(gs.pi / 7)
    cos = gs.cos(angle / 2)
    sin = gs.sin(angle / 2)
    u = gs.array([1., 2., 3.])
    u = u / gs.linalg.norm(u)
    scalar = gs.array(cos)
    vec = sin * u
    translation = gs.array([4., 5., 6.])

is computed.
        base_point : array-like, shape=[..., dim + 1]
            Point on the hypersphere.

        Returns
        -------
        transported_tangent_vec: array-like, shape=[..., dim + 1]
            Transported tangent vector at exp_(base_point)(tangent_vec_b).
        """
        theta = gs.linalg.norm(tangent_vec_b, axis=-1)
        normalized_b = gs.einsum('..., ...i->...i', 1 / theta, tangent_vec_b)
        pb = gs.einsum('...i,...i->...', tangent_vec_a, normalized_b)
        p_orth = tangent_vec_a - gs.einsum('..., ...i->...i', pb, normalized_b)
        transported = \
            - gs.einsum('..., ...i->...i', gs.sin(theta) * pb, base_point)\
            + gs.einsum('..., ...i->...i', gs.cos(theta) * pb, normalized_b)\
            + p_orth
        return transported

mask_0 = gs.isclose(angle, 0.)
            mask_0_float = gs.cast(mask_0, gs.float32)

            coef_1 += mask_0_float * (1 - (angle ** 2) / 6)
            coef_2 += mask_0_float * (1 / 2 - angle ** 2)

            mask_else = ~mask_0
            mask_else_float = gs.cast(mask_else, gs.float32)

            # This avoids division by 0.
            angle += mask_0_float * 1.

            coef_1 += mask_else_float * (gs.sin(angle) / angle)
            coef_2 += mask_else_float * (
                (1 - gs.cos(angle)) / (angle ** 2))

            term_1 = gs.zeros((n_rot_vecs,) + (self.n,) * 2)
            term_2 = gs.zeros_like(term_1)

            coef_1 = gs.squeeze(coef_1, axis=0)
            term_1 = (gs.eye(self.dimension)
                      + gs.einsum('n,njk->njk', coef_1, skew_rot_vec))

            term_2 = (coef_2
                      + gs.einsum('nij,njk->nik', skew_rot_vec, skew_rot_vec))
            #for i in range(n_rot_vecs):
            #    term_1[i] = (gs.eye(self.dimension)
            #                 + coef_1[i] * skew_rot_vec[i])
            #    term_2[i] = (coef_2[i]
            #                 * gs.matmul(skew_rot_vec[i], skew_rot_vec[i]))
            rot_mat = term_1 + term_2

subsphere = Hypersphere(dimension=dim-1)

    # Define north pole
    north_pole = np.zeros(dim+1)
    north_pole[dim] = 1.0
    for j in range(NN):
        # Sample n points from the uniform distrib on a subsphere
        # of radius theta (i.e cos(theta) in ambiant space)
        data = gs.zeros((n, dim+1), dtype=gs.float64)
        # For sampling on a subsphere, use RandomUniform(dim-1)
        directions = subsphere.random_uniform(n)
        # directions = my_random_sample(n, dim-1)
        for i in range(n):
            for j in range(dim):
                data[i,j] = gs.sin(theta) * directions[i,j]
            data[i,dim] = gs.cos(theta)
        ## Compute empirical Frechet mean of the n-sample
        current_mean = sphere.metric.adaptive_gradientdescent_mean(data, n_max_iterations=64, init_points=[north_pole])
        var.append(sphere.metric.squared_dist(north_pole, current_mean))
    return (np.mean(var), 2* np.std(var) / np.sqrt( NN ) )