Tensorflow addons compatibility.