tfa.register.register_custom_kernels

View source on GitHub