Simulation probabilistic classical computation on quantum computers

{Edit: I found a much easier way to do this without resorting to linear system solving. I have included both methods in the answer}

Finding such a $U$ certainly seems possible as with the enlarged Hilbert space there is much more freedom to choose your parameters. I will outline how to do this. First a bit of notation. Let $$ langle j,b |U|i,a rangle = U_{N(i-1) + a ,~ N(j-1) + b}.$$

$U$ is an $N^2 times N^2$ unitary matrix.

Apart form the usual conditions for unitarity, the rows $1, ~N+1, ~2N+1 ldots$ will be affected by the constraint given in your question.

Now lets construct the first row corresponding to $i=1$ and $a = 1$. We know that,

$$p_{1,j} = sum_{b=1} ^ N |U_{1,N(j-1)+b}|^2.$$

$large{textbf{Method 1}}$

Choose $U_{1,1} = sqrt{p_{1,1}}$, $U_{1,N+1}= sqrt{p_{1,2}}$,... ,$U_{1,N(N-1)+1}= sqrt{p_{1,N}}$ to satisfy this condition. The condition $sum_j p_{1,j} = 1$ ensures that the norm of the first row is unity.

Next let us construct the row $N+1$. This row vector has to be orthogonal to the first row and should satisfy the condition, $$p_{2,j} = sum_{b=1} ^ N |U_{N+1,N(j-1)+b}|^2. $$.

So choose $U_{N+1,2} = sqrt{p_{2,1}}$, $U_{N+1,N+2}= sqrt{p_{2,2}}$... $U_{N+1,N(N-1)+2}= sqrt{p_{2,N}}$. You can easily verify that the inner product condition is satisfied. Next you can select the row $2N +1$. Here again you need to set the values as above but shifted by one element in the horizontal direction so that the orthogonality conditions are satisfied. Do this till you reach row $N(N-1) +1$. Now your matrix has $N$ rows filled according to the constraint in the question. The rest of the rows can be filled by vectors orthonormal to the ones you already constructed using the Gram-Schmidt procedure.

This is much easier than the next method, which involves linear system solving.

$large{textbf{Method 2}}$

Select the first $N$ elements in the first row so that the sum of squares their absolute values is equal to $p_{11}$, the second $N$ such that the sum of their absolute values is equal to $p_{12}$ and so on till you complete the first row. The condition $sum_j p_{1,j} = 1$ ensures that the norm of the first row is unity.

Next we will fill row $N+1$ corresponding to $i=2 $ and $~a = 1$. The constraint here is, $$p_{2,j} = sum_{b=1} ^ N |U_{N+1,N(j-1)+b}|^2.$$

Choose the first $N$ elements such that their inner product with the first $N$ elements of the first row is zero,$$ sum_{b=1} ^ N U_{1,b}U_{N+1,b }^* = 0.$$ You can see this as an under constrained system of linear equations. Choose a solution from the family of solutions you get and then scale it appropriately so that sum of square of their absolute values add up to $p_{2,1}$. Do this for the next $N$ elements in the row and so on till you fill that row. Again the condition $sum_j p_{2,j} = 1$ ensures that the norm of the row is unity.

Now go to the row $2N+1$. This time you will have to solve a $N times 2$ system of linear equations. As you keep going down the size of the system you need to solve will go up. I believe solutions will always exist as the previous rows are always orthogonal to each other and your system will always have maximum rank .You will solve $N times N$ system at the row $N(N-1)+1$. Now your matrix has $N$ rows filled according to the constraint in the question. The rest of the rows can be filled by vectors orthonormal to the ones you already constructed using the Gram-Schmidt procedure.