Thanks @Mathieu - incredibly helpful and appreciate you taking the time to put together the example! I’ll give your #2 route a go with my application since I’m already training a Random Forest with multiple trees.
Also thanks for the link to Wagner et al paper - I was aware of the technique as I’ve seen reference in the R ranger package and this Scikit Learn contrib package 1, but hadn’t seen the source paper so look forward to taking a look to understand more deeply.
Sorry for asking on both GitHub and the forum - after posting there I wasn’t sure if GitHub was only for code focused issues and the forum for help requests so thought here might be more appropriate.
Will add a feature request to GitHub. Do you prefer that they sit at the TFDF level? In this case I think it would probably involve changes to YDF (i.e. a custom reducer similar to here 1) but it looks like most activity is on the TFDF repository.