13 ms·
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
by 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 agotf.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 agoYes! 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...