4 ms·

Why is that difficult in TensorFlow? Wouldn’t you just call tf.gradients on the output of tf.gradients? At least that is what I usually do and it’s very simple.

• duckworthd 8y ago
  There are two reasons, 1) Take a closer look at tf.hessian()'s implementation. Notice the tf.while() loop, calculating tf.gradient() for each coordinate of the first gradient? That doesn't happen in parallel! Plus, you need to know about TensorArrays! 2) Per-example gradients means an outside tf.while() loop over each example. Another linear slow down! The difficulty I refer to is in wrangling conditionals in TF, and the trickiness of obtaining partial derivatives (rather than total derivatives) which JAX makes trivial. There may be a
  • albertzeyer 8y ago
    tf.while code will run in parallel as far as possible (if it does not depend on each other) (up to some configurable degree).
    • duckworthd 8y ago
      Yes! Though this answer has some subtlety. tf.while() will run several iterations in parallel but this is not the same as _batching_ those same iterations. For that, you'll need to use the experimental parallel_for feature [0]. Using this should get you into roughly the same speed as JAX. Tricky! [0]: https://github.com/tensorflow/tensorflow/tree/b3e007394680801113f492fa1f5a9784e8502f19/tensorflow/python/ops/parallel_for https://github.com/tensorflow/tensorflow/tree/b3e00739468080...