Is it possible to use tf.data with tf operations while utilizing jax or torch as the backend?

Is it possible to use tf.data with tf operations while utilizing jax or torch as the backend? #2433