This is mostly a reminder for myself, but also documentation if someone else ends up doing it.

After reading through the python code (see here), I found that to force function JITing you need to set the _XlaMustCompile and _noinline attributes to true, as is done here.

However, this results in an error. I didn't run it down at the time, but it was reported in tensorflow/tensorflow#50458. Once that is fixed we should try this again.