The required kernel changes were implemented almost 2 years ago in 80da0a63, but we forgot to change the Python gradient registry. PiperOrigin-RevId: 185790703