AutoTVM tree RNN model


I read the “Learning to Optimize Tensor Programs” paper, which describes the design of AutoTVM.
Though the focus is on the GBT cost model, I’m interested in replicating and exploring the TreeGRU model.

Does anyone have advice on embedding loop levels?
How I can access the tree structure of the AST to encode each level?
Can I leverage itervar features and curve sample features used by the GBT model?


Here is the related code. It is on my branch and not upstreamed.
You can find related tree-structure extraction code.