tensorflow gradienttape with ragged tensor

demo.py

import tensorflow as tf
## note tape.watch() cannot accept ragged tensor as input
rt = tf.ragged.constant([[1.], [2., 3.]])
# rt = rt.to_tensor()
with tf.GradientTape() as t:
    t.watch(rt.values)
    loss = 2*rt[0,0]+rt[1,1]**3
g_flatten = t.gradient(loss, rt.values)
g = rt.with_values(g_flatten)
print(g)