Your code is not a minimal working example, it seems to use TensorFlow (is that what `tf` is in there?). I do not have TensorFlow installed.
1.- Send an example that uses only mpi4py and numpy.
2.- Do not paste the code in the email body. Put it in a main.py script and attach it. If you paste, things do not go well, I get some weird quotation marks, if I copy and paste back, then the code does not run (indentation errors).
3.- Make sure that the code in main.py is complete and self-contained. The code you sent was missing import statements. Maybe these are trivial imports, but why should I waste time figuring them out?
4.- PLEASE RUN YOUR CODE before sending it, so that you are sure the code runs with no import/syntax error and fails as expected.
5.- Remember to confirm how many MPI processes are you using to run the code.
Please understand that here we provide help for free and in good will. Do not make the life of those trying to help miserable by forcing them to guess missing information and fight with trivial bugs (missing import statements, weird quotation marks in the code, using heavy dependencies like TensorFlow when it is obviously not needed, indentation errors, etc.). If you do not follow these rules, the outcome is slower resolution of your issue.
All that being said, I still think that you code has some logic bugs that lead to deadlocks.
I said it before and I'm going to say it again. Irecv().Wait() is the same as Recv(), things will work the same, and there is no "eliminate waiting" argument to make.
How much experience do you have with MPI? Have you used it before? Have you ever coded something similar to what you are trying to do now? The code, as written the last time, seems to have send-to-self calls. Is that intended?
You are using Isend() and discarding the result Request objects. You should collect them all in a list and eventually call MPI.Request.Waitall(list_of_requests)