Hi,
The main problem with conditionals is that they are handled on the python side and so the values needs to be on the CPU.
So if you use an accelerator like GPU or TPU, the CPU has to wait for the value to be computed and then it can go through the conditional and continue the execution. This is sub-optimal but cannot really be worked around.
cc @ailzhang there might be other details about TPU that I don’t know about?