The only solutions that did not break when tested with OOD shapes were made by experts who know a lot about GPU programming and that did not create 25k lines of CUDA but followed and adjusted their solution in reasonable bounds.
The takeaway from this is that these approaches will always solve for specificity, but it's a much harder task to steer the model into making general solutions. So if you're an inference provider for some specific model shape, fantastic, go for it. If you are a maintainer of a open-source library, this is not useful.