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

def test_projection(self):
        # Test 3D and nD cases
        for n in self.n_seq:
            group = self.so[n]
            rot_mat = gs.eye(n)
            delta = 1e-12 * gs.random.rand(n, n)
            rot_mat_plus_delta = rot_mat + delta
            result = group.projection(rot_mat_plus_delta)
            expected = rot_mat
            self.assertTrue(gs.allclose(result, expected))

def test_transpose(self):
        tr = self.space.transpose
        ar = gs.array
        a = gs.eye(3, 3, 1)
        b = gs.eye(3, 3, -1)
        self.assertAllClose(tr(a), b)
        self.assertAllClose(tr(ar([a, b])), ar([b, a]))

# Diagonal left and right invariant metrics
        diag_mat_at_identity = gs.eye(group.dim)

        left_diag_metric = InvariantMetric(
            group=group,
            inner_product_mat_at_identity=None,
            left_or_right='left')
        right_diag_metric = InvariantMetric(
            group=group,
            inner_product_mat_at_identity=diag_mat_at_identity,
            left_or_right='right')

        # General left and right invariant metrics
        # FIXME (nina): This is valid only for bi-invariant metrics
        sym_mat_at_identity = gs.eye(group.dim)

        left_metric = InvariantMetric(
            group=group,
            inner_product_mat_at_identity=sym_mat_at_identity,
            left_or_right='left')

        right_metric = InvariantMetric(
            group=group,
            inner_product_mat_at_identity=sym_mat_at_identity,
            left_or_right='right')

        matrix_left_metric = InvariantMetric(group=matrix_so3)

        matrix_right_metric = InvariantMetric(
            group=matrix_so3,
            left_or_right='right')

def test_matrix_from_tait_bryan_angles_extrinsic_zyx(self):
        tait_bryan_angles = gs.array([0., 0., 0.])
        result = self.group.matrix_from_tait_bryan_angles_extrinsic_zyx(
            tait_bryan_angles)
        expected = gs.eye(3)

        self.assertAllClose(result, expected)

        angle = gs.pi / 6.
        cos_angle = gs.cos(angle)
        sin_angle = gs.sin(angle)

        tait_bryan_angles = gs.array([angle, 0., 0.])
        result = self.group.matrix_from_tait_bryan_angles_extrinsic_zyx(
            tait_bryan_angles)
        expected = gs.array([[1., 0., 0.],
                             [0., cos_angle, - sin_angle],
                             [0., sin_angle, cos_angle]])

        self.assertAllClose(result, expected)

def __init__(self, dim, default_point_type='vector', **kwargs):
        super(LieGroup, self).__init__(
            dim=dim, default_point_type=default_point_type, **kwargs)

        self.left_canonical_metric = InvariantMetric(
            group=self,
            inner_product_mat_at_identity=gs.eye(self.dim),
            left_or_right='left',
        )

        self.right_canonical_metric = InvariantMetric(
            group=self,
            inner_product_mat_at_identity=gs.eye(self.dim),
            left_or_right='right',
        )

        self.metrics = []

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

            rot_mat = self.projection(rot_mat)

        else:
            skew_mat = self.skew_matrix_from_vector(rot_vec)
            rot_mat = self.embedding_manifold.group_exp_from_identity(skew_mat)

jacobian_block_line_1 = gs.concatenate(
            [jacobian_rot, block_zeros_1], axis=2)

        if left_or_right == 'left':
            rot_mat = self.rotations.matrix_from_rotation_vector(
                rot_vec)
            jacobian_trans = rot_mat
            block_zeros_2 = gs.zeros(
                (n_points, dim_translations, dim_rotations))
            jacobian_block_line_2 = gs.concatenate(
                [block_zeros_2, jacobian_trans], axis=2)

        else:
            inv_skew_mat = - self.rotations.skew_matrix_from_vector(
                rot_vec)
            eye = gs.to_ndarray(gs.eye(self.n), to_ndim=3)
            eye = gs.tile(eye, [n_points, 1, 1])
            jacobian_block_line_2 = gs.concatenate(
                [inv_skew_mat, eye], axis=2)

        jacobian = gs.concatenate(
            [jacobian_block_line_1, jacobian_block_line_2], axis=-2)
        return jacobian[0] if (len(point) == 1 or point.ndim == 1) \
            else jacobian

quat_sq_norm = quat_vec_norm ** 2 + quat_scalar ** 2

        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.einsum('ni,nj->nij', quat_vec, quat_vec)
                               / quat_vec_norm * quat_vec_norm)
                            + 2 * quat_arctan2 / quat_vec_norm * gs.eye(3))

        differential_scalar_t = gs.transpose(differential_scalar, axes=(1, 0))

        upper_left_block = gs.hstack(
            (differential_scalar_t, differential_vec[0]))
        upper_right_block = gs.zeros((3, 3))
        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)
    return grad

# This avoids division by 0.
                angle += mask_pi_float * 1.
                coef_1 += mask_else_float * (angle / 2) / gs.tan(angle / 2)
                coef_2 += mask_else_float * (1 - coef_1) / angle ** 2

                jacobian = gs.zeros((n_points, self.dimension, self.dimension))
                for i in range(n_points):
                    mask_i_float = get_mask_i_float(i, n_points)

                    sign = - 1
                    if left_or_right == 'left':
                        sign = + 1

                    jacobian_i = (
                        coef_1[i] * gs.eye(self.dimension)
                        + coef_2[i] * gs.outer(point[i], point[i])
                        + sign * self.skew_matrix_from_vector(point[i]) / 2)
                    jacobian_i = gs.squeeze(jacobian_i, axis=0)

                    jacobian += gs.einsum(
                        'n,ij->nij', mask_i_float, jacobian_i)

            else:
                if left_or_right == 'right':
                    raise NotImplementedError(
                        'The jacobian of the right translation'
                        ' is not implemented.')
                jacobian = self.matrix_from_rotation_vector(point)

            assert gs.ndim(jacobian) == 3

y_true_pose = gs.hstack([y_true_rot_vec, y_true[:, 4:]])
        grad = lie_group.grad(y_pred_pose, y_true_pose, SE3, metric)

        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

        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.einsum('ni,nj->nij', quat_vec, quat_vec)
                               / quat_vec_norm * quat_vec_norm)
                            + 2 * quat_arctan2 / quat_vec_norm * gs.eye(3))

        differential_scalar_t = gs.transpose(differential_scalar, axes=(1, 0))

        upper_left_block = gs.hstack(
            (differential_scalar_t, differential_vec[0]))
        upper_right_block = gs.zeros((3, 3))
        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)