Gradient Tape in Jax

[eadadi]

Hey all,

I know that in TF there was the notion of gradient tape that allowed efficient computation like:

""" Computing And Preserving Shared-Gradients """
with tf.GradientTape(persistent=True) as tape:
.... .... ...

return traj

""" Later on, use  tape two times """
with tape:
....
...
#Then Optimize first component using tape..
self.opt(tape, loss, A)

"""  Later, use tape again to  optimize second component """
with tape:
... ...
# Optimize the second component
self.opt(tape, loss, B)

This allows computing the first part of gradient once for both components. Is there anything similar in Jax?