In the first expression, a * b is int32 multiplication. In the second expression, a * b is float32 multiplication. I’m guessing that the type differences are messing up the gradients. From my understanding the gradient tape is watching for operations, so the order of operations and casting matters for backprop.
I didn’t do a deep dive, but the documentation on Tensorflow type promotion mentions dunder (double underline) operations where the math goes wrong due to bit-widening.
Did you read my example thread (linked several times earlier on this thread) where I showed some of the possible ways to get errors here?
Yes! That thread helped me get to a successful submission for my assignment.