The first two “Out:” blocks of that section seem identical to me, aside from variable/parameter names, even though the second one supposedly is contains the branching logic that the first one is missing. So, am I going crazy, or is there something subtle happening there, or is that doc incorrect?
The outputs don’t make the difference obvious, but the warning raised by tracing the model would give you some information, what could fail using this approach:
TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
if x.sum() > 0:
This warning should be raised while tracing the MyDecisionGate.
The documentation is incorrect because the control flow is in the innermost module and that is not the one for which code is being printed. I have had a PR up to fix this for a while but couldn’t get a review for a while and I haven’t merged it since it’s been accepted. I will get this merged today.