In your partial_transpose_np
you are missing a dot. Change the line
blkm[Block(i, j)] = trm
to
blkm[Block(i, j)] .= trm
and it will work.
The reason is that you were reassigning each block (which is internally stored as a separate array) to trm
. This way, all the blocks were identical (not just equal) after the loop. Instead you should overwrite the content of each block, hence the dot above.