I've been using JAX for a while now. A paper I'm an author on (https://arxiv.org/abs/1806.09597, a follow-up to https://news.ycombinator.com/item?id=18633215) resulted in an algorithm that required taking second-derivatives on a per-example basis. This is extremely difficult in TensorFlow, but with JAX it was a 2-liner. Even better, it's _super_ fast, thanks to XLA's compile-to-GPU and JAX's auto-batching mechanics.
I highly recommend JAX to power users. It's nowhere near as feature-complete from a neural network sense as, say, PyTorch, but it is very good at what it does, and its core developers are second to none in responsiveness.
About this time last year I had an optimization problem involving linear combinations of Voronoi cell centroids and I basically used random search (I can't remember why differential evolution didn't work).
I don't know much, say, Keras, but I think I understand how to implement backprop-like behaviors with Autograd. Maybe even sandwich concept-specific models between fully-connected linear+sigmoid layers to give them some oomph.
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.
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.
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.
Comments
My $0.02.
I've been using JAX for a while now. A paper I'm an author on (https://arxiv.org/abs/1806.09597, a follow-up to https://news.ycombinator.com/item?id=18633215) resulted in an algorithm that required taking second-derivatives on a per-example basis. This is extremely difficult in TensorFlow, but with JAX it was a 2-liner. Even better, it's _super_ fast, thanks to XLA's compile-to-GPU and JAX's auto-batching mechanics.
I highly recommend JAX to power users. It's nowhere near as feature-complete from a neural network sense as, say, PyTorch, but it is very good at what it does, and its core developers are second to none in responsiveness.
Interesting.
About this time last year I had an optimization problem involving linear combinations of Voronoi cell centroids and I basically used random search (I can't remember why differential evolution didn't work).
I don't know much, say, Keras, but I think I understand how to implement backprop-like behaviors with Autograd. Maybe even sandwich concept-specific models between fully-connected linear+sigmoid layers to give them some oomph.
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.
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
tf.while code will run in parallel as far as possible (if it does not depend on each other) (up to some configurable degree).
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/b3e00739468080...