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

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...