Ah! Thank you. I am still too stupid to solve this puzzel. But, thank you.

Hmm.. sorry to keep bothering you but reading the discription again, I feel it might be clearer to write "(i.e length of each line should be equal to the total number of lines in the code)". What do you think?

I disagree. The task is challenging enough on its own, discription being vague isn't helpful. Personally, I had to give up without even figuring out the actual task.

"... only allows programs with a square shape ..."

I think this line is not clear. How about something like:

"... only allows programs with a square shape (i.e number of line in code == length of each line) ..."

Also: "he will only let you use functionalities functions from this SDK"

So, what I am understanding from this... in a expression (A and B), if A and B are both True(not 0), B will be returned. Otherwise False is returned which is evaluated as 0, so it passes the last fixed test.

Nice. The count_Kprimes and find_k functions are efficient and elegant.

But the puzzle function could be a lot more optimized. (see g964's solution)

Seems to me when it hit dead ends, which end in 1, it pops twice since range(1-1, 0, -1) skips the for loop. next loop starts with the poped value-1.