You wrote this...

  All the comparisons like > are better written using their TensorFlow equivalent (e.g tf.greater). Autograph can convert them (you could write >), but it’s less idiomatic and I recommend to do not relying upon the automatic conversion, for having full control.

...but I'm not sure you realized that the for loop and the if statement in your code are being transparently compiled to dataset.map() and tf.cond() for you by Autograph :)