Lower TensorFlow compute graph to MLIR at train time

I wish to explore the pipeline to lower a model written in TensorFlow to MLIR and from there to LLVM IR.
The idea is while Training a model written in TensorFlow 2, how to lower the compute graph to MLIR so that custom processing may be done on it?

What are the APIs in TensorFlow 2 and in MLIR to achieve the same?

1 Like