It's been a little while since I sent this mail, but we are just about ready to push this code into TF addons. Since it has been a while since we discussed this, I want to check again if there are any issues.
The layer itself now (our internal version) only depends on public symbols in TF 2.8. As well, once this layer is in TF addons, we will add unit tests on the CloudTPU side to ensure that there is no breakage. We also guarantee continued support for this layer going forward.
I've tagged Ziyin on this chain, who will be sending the pull request at some point in the next few weeks.