@Hauke first of all, I happy to see your PR and your question here. I see that you implemented it without tf.eye usage as a composition of another available tf ops. There are two ways: doing nothing about tf.eye or create an extension or util function which do the same (but built on a few tf ops). I never extract the same pieces of code as tf extensions, but faced with the same issue in a few moments around gradient clipping, or making fake concat op via padding. So you could be a first contributor here. If you need a separate issue and separate PR for that for tracking it separately feel free to do it.
There are some analogues of Keras functions (used as an analogue of tf.moments) and
here you could find some useful extension functions for the TF Tensor. You could add your implementation here, in "extension" package