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

Returns
        -------
        regularized_rot_vec : array-like, shape=[..., 3]
            Rotation vector.
        """
        n_rot_mats, _, _ = rot_mat.shape

        trace = gs.trace(rot_mat, axis1=1, axis2=2)
        trace = gs.to_ndarray(trace, to_ndim=2, axis=1)
        trace_num = gs.clip(trace, -1, 3)
        angle = gs.arccos(0.5 * (trace_num - 1))
        rot_mat_transpose = gs.transpose(rot_mat, axes=(0, 2, 1))
        rot_vec_not_pi = self.vector_from_skew_matrix(
            rot_mat - rot_mat_transpose)
        mask_0 = gs.cast(gs.isclose(angle, 0.), gs.float32)
        mask_pi = gs.cast(gs.isclose(angle, gs.pi, atol=1e-2), gs.float32)
        mask_else = (1 - mask_0) * (1 - mask_pi)

        numerator = 0.5 * mask_0 + angle * mask_else
        denominator = (1 - angle ** 2 / 6) * mask_0 + 2 * gs.sin(
            angle) * mask_else + mask_pi

        rot_vec_not_pi = rot_vec_not_pi * numerator / denominator

        vector_outer = 0.5 * (gs.eye(3) + rot_mat)
        gs.set_diag(
            vector_outer, gs.maximum(
                0., gs.diagonal(vector_outer, axis1=1, axis2=2)))
        squared_diag_comp = gs.diagonal(vector_outer, axis1=1, axis2=2)
        diag_comp = gs.sqrt(squared_diag_comp)
        norm_line = gs.linalg.norm(vector_outer, axis=2)

Convert rotation vector to rotation matrix.
        """
        assert self.belongs(rot_vec, point_type='vector')
        rot_vec = self.regularize(rot_vec, point_type='vector')
        n_rot_vecs, _ = rot_vec.shape

        if self.n == 3:
            angle = gs.linalg.norm(rot_vec, axis=1)
            angle = gs.to_ndarray(angle, to_ndim=2, axis=1)

            skew_rot_vec = self.skew_matrix_from_vector(rot_vec)

            coef_1 = gs.zeros_like(angle)
            coef_2 = gs.zeros_like(angle)

            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)

Returns
        -------
        exp : array-like, shape=[..., dim + 1]
            Point on the hypersphere equal to the Riemannian exponential
            of tangent_vec at the base point.
        """
        _, extrinsic_dim = base_point.shape
        n_tangent_vecs, _ = tangent_vec.shape

        hypersphere = Hypersphere(dim=extrinsic_dim - 1)
        proj_tangent_vec = hypersphere.to_tangent(
            tangent_vec, base_point)
        norm_tangent_vec = self.embedding_metric.norm(proj_tangent_vec)
        norm_tangent_vec = gs.to_ndarray(norm_tangent_vec, to_ndim=1)

        mask_0 = gs.isclose(norm_tangent_vec, 0.)
        mask_non0 = ~mask_0

        coef_1 = gs.zeros((n_tangent_vecs,))
        coef_2 = gs.zeros((n_tangent_vecs,))
        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)

log : array-like, shape=[..., dim]
            Tangent vector at the base point equal to the Riemannian logarithm
            of point at the base point.
        """
        add_base_point = self.mobius_add(-base_point, point)
        norm_add =\
            gs.expand_dims(gs.linalg.norm(
                           add_base_point, axis=-1), axis=-1)

        norm_base_point =\
            gs.expand_dims(gs.linalg.norm(
                           base_point, axis=-1), axis=-1)

        log = (1 - norm_base_point**2) * gs.arctanh(norm_add)

        mask_0 = gs.isclose(gs.squeeze(norm_add, axis=-1), 0.)
        mask_non0 = ~mask_0
        add_base_point = gs.assignment(
            add_base_point,
            gs.zeros_like(add_base_point[mask_0]),
            mask_0)
        add_base_point = gs.assignment(
            add_base_point,
            add_base_point[mask_non0] / norm_add[mask_non0],
            mask_non0)

        log = gs.einsum(
            '...i,...j->...j', log, add_base_point)
        return log

the function computes its complementary in 2pi and
        inverts the direction of the rotation axis.
        """
        if point_type is None:
            point_type = self.default_point_type

        if point_type == 'vector':
            point = gs.to_ndarray(point, to_ndim=2)
            assert self.belongs(point, point_type)
            n_points, _ = point.shape

            regularized_point = gs.copy(point)
            if self.n == 3:
                angle = gs.linalg.norm(regularized_point, axis=1)

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

                mask_0_float = gs.cast(mask_0, gs.float32)
                mask_not_0_float = gs.cast(mask_not_0, gs.float32)
                mask_pi_float = gs.cast(mask_pi, gs.float32)

                k = gs.floor(angle / (2 * gs.pi) + .5)

                norms_ratio = gs.zeros_like(angle)

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

                norms_ratio += mask_not_0_float * (
                    1. - 2. * gs.pi * k / angle)

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

        trace = gs.trace(rot_mat, axis1=1, axis2=2)
        trace = gs.to_ndarray(trace, to_ndim=2, axis=1)
        trace_num = gs.clip(trace, -1, 3)
        angle = gs.arccos(0.5 * (trace_num - 1))
        rot_mat_transpose = gs.transpose(rot_mat, axes=(0, 2, 1))
        rot_vec_not_pi = self.vector_from_skew_matrix(
            rot_mat - rot_mat_transpose)
        mask_0 = gs.cast(gs.isclose(angle, 0.), gs.float32)
        mask_pi = gs.cast(gs.isclose(angle, gs.pi, atol=1e-2), gs.float32)
        mask_else = (1 - mask_0) * (1 - mask_pi)

        numerator = 0.5 * mask_0 + angle * mask_else
        denominator = (1 - angle ** 2 / 6) * mask_0 + 2 * gs.sin(
            angle) * mask_else + mask_pi

        rot_vec_not_pi = rot_vec_not_pi * numerator / denominator

        vector_outer = 0.5 * (gs.eye(3) + rot_mat)
        gs.set_diag(
            vector_outer, gs.maximum(
                0., gs.diagonal(vector_outer, axis1=1, axis2=2)))
        squared_diag_comp = gs.diagonal(vector_outer, axis1=1, axis2=2)
        diag_comp = gs.sqrt(squared_diag_comp)
        norm_line = gs.linalg.norm(vector_outer, axis=2)
        max_line_index = gs.argmax(norm_line, axis=1)

def is_symmetric(mat, tolerance=TOLERANCE):
    """Check if a matrix is symmetric."""
    mat = gs.to_ndarray(mat, to_ndim=3)
    n_mats, _, _ = mat.shape
    mat_transpose = gs.transpose(mat, axes=(0, 2, 1))

    mask = gs.isclose(mat, mat_transpose, atol=tolerance)
    mask = gs.all(mask, axis=(1, 2))

    return mask

point = self.regularize(point)

        rotations = self.rotations
        dim_rotations = rotations.dim

        rot_vec = point[:, :dim_rotations]
        angle = gs.linalg.norm(rot_vec, axis=1)
        angle = gs.to_ndarray(angle, to_ndim=2, axis=1)

        translation = point[:, dim_rotations:]

        skew_rot_vec = rotations.skew_matrix_from_vector(rot_vec)
        sq_skew_rot_vec = gs.matmul(skew_rot_vec, skew_rot_vec)

        mask_close_0 = gs.isclose(angle, 0.)
        mask_close_pi = gs.isclose(angle, gs.pi)
        mask_else = ~mask_close_0 & ~mask_close_pi

        mask_close_0_float = gs.cast(mask_close_0, gs.float32)
        mask_close_pi_float = gs.cast(mask_close_pi, gs.float32)
        mask_else_float = gs.cast(mask_else, gs.float32)

        mask_0 = gs.isclose(angle, 0., atol=1e-6)
        mask_0_float = gs.cast(mask_0, gs.float32)
        angle += mask_0_float * gs.ones_like(angle)

        coef_1 = - 0.5 * gs.ones_like(angle)
        coef_2 = gs.zeros_like(angle)

        coef_2 += mask_close_0_float * (
            1. / 12. + angle ** 2 / 720.
            + angle ** 4 / 30240.