There’s a lot going on in this question.
But you shouldn’t need to go that low level to do this.
Given Alex’s response 1 I hope we can get that issue fixed eventually.
3 things that may help here in the mean time:


GradientTape.gradient accepts an output_gradients argument. To run your gradient in stages you can do something like this:

with tf.GradientTape() as tape1:
   x = chunk1(input)

with tf.GradientTape() as tape2:
  tape2.watch(x)
  y = chunk2(x)
  loss = loss_fn(y, y_true)

g2, gx = tape2.gradient(loss, [chunk2.variables, x])

g2 = tape1.gradient(x, [chunk1.variables], output_gradient=gx)

I should add an example of this to the Advanced automatic differentiation guide 1
I’m going to guess that if you wrap your with tf.device over the with tf.GradientTape then the gradient operations go on that device.


Another thing that could work would be to use @tf.custom_gradient. I’m pretty sure you can use a with tf.device(), tf.GradientTape() as tape: in the custom gradient function.


For model parallelism mesh approaches have been gaining popularity. Where instead of assigning layers to devices you split layers across devices. But I don’t have an off-the-shelf solution for that.