How to get a coarse-grained op-level graph in tensorflow

Hi, tensorflow community,
I want to use tensorflow to get the full computation graph (including forward, backward and parameter update). I tried tf.functions, but the graph I got is too fine-grained, as many ops (Adam for example) are splited into smaller operators (add, mul, div etc.). So is there any methods that I can get a coarse-grained op-level graph? Thanks a lot!

Jiangfei