Change how `take_along_axis` is computed in DeBERTa to stop confusing XLA
The previous code for `take_along_axis()` in DeBERTa used dynamic TF shapes like `tf.shape()` and `tf.rank()` in conditionals. This is a data-dependent conditional, which is forbidden in XLA.

Replacing these with the static shape equivalents `x.shape` and `x.shape.rank` works fine, and now DeBERTa can be compiled successfully with XLA.
_The documentation is not available anymore as the PR was closed or merged._<|||||>@gante The original torch code used `take_along_axis`, so I guess this is a complete TF reimplementation of it! That approach makes way more sense, though - let me make some changes!<|||||>> The original torch code used take_along_axis

That would explain it!