So in essence: Normally a complete TE graph will be lowered to TIR representation. Your assumptions is that doing small changes in part of the TE graph should not propagate throughout all the TIR AST and you want to somehow “cache” the part of the TIR which is independent of changes of this subgraph. I guess since you are describing a complete network in TE lowering takes some time and you want to save this time. Correct?
Sadly I don’t know how to help you. The only thing that comes to mind is to do it at TIR level, but this would basically mean you need to reimplement the scheduling primitives (from TE) directly on TIR level. Maybe leverage the functionality from [RFC] TensorIR: A schedulable IR for TVM can help you?