The key reason of this issue is:
After the flatten_if enabled, the if statements are unrolled to select, therefore the code path that not be walked through (e.g., with condition false) in the forward pass will also be transformed and executed when in backward pass. In the case div, the denominator might be a variable which only appears in the path with condition false. The variable doesn't get any assignment in the forward pass and ends up with a default value zero, which causes nan in backward pass.
Related issue: https://github.com/taichi-dev/taichi/issues/7103
This is temporary fix.
The key reason of this issue is: After the
flatten_ifenabled, the if statements are unrolled toselect, therefore the code path that not be walked through (e.g., with condition false) in the forward pass will also be transformed and executed when in backward pass. In the casediv, the denominator might be a variable which only appears in the path with condition false. The variable doesn't get any assignment in the forward pass and ends up with a default value zero, which causesnanin backward pass.