def partial_trace(self, keep_axes: set[int]) -> np.ndarray:
        """Traces out all qubits except keep_axes.

        Args:
            keep_axes: The set of axes that are left after computing the
                partial trace. For example, if we have a circuit for 3 qubits
                and this parameter only has one qubit, the entire density matrix
                would be 8x8, but this function returns a 2x2 matrix.

        Returns:
            An array that contains the partial trace.
        """

        contracted_inds = set(map(self.i_str, set(range(len(self._qid_shape))) - keep_axes))

        conj_pfx = "conj_"

        tensor_network = qtn.TensorNetwork(self._M)

        # Rename the internal indices to avoid collisions. Also rename the qubit
        # indices that are kept. We do not rename the qubit indices that are
        # traced out.
        conj_tensor_network = tensor_network.conj()
        reindex_mapping = {}
        for M in conj_tensor_network.tensors:
            for ind in M.inds:
                if ind not in contracted_inds:
                    reindex_mapping[ind] = conj_pfx + ind
        conj_tensor_network.reindex(reindex_mapping, inplace=True)
        partial_trace = conj_tensor_network @ tensor_network

        forward_inds = list(map(self.i_str, keep_axes))
        backward_inds = [conj_pfx + forward_ind for forward_ind in forward_inds]
        return partial_trace.to_dense(forward_inds, backward_inds)
Traces out all qubits except keep_axes.

        Args:
            keep_axes: The set of axes that are left after computing the
                partial trace. For example, if we have a circuit for 3 qubits
                and this parameter only has one qubit, the entire density matrix
                would be 8x8, but this function returns a 2x2 matrix.

        Returns:
            An array that contains the partial trace.
        
cirq-core/cirq/contrib/quimb/mps_simulator.py