Remove workaround for `take_along_axis`

[mattdangerw]

Our backend module is now a fairly clean Keras 3 API, but we have one major work around remaining. During our port to Keras 3 we had to hack up take_along_axis on the tensorflow backend to support dynamic shape.

With a Keras 3 local development environment, with the tensorflow backend:

• [ ] Delete the take_along_axis hack.
• [ ] Some tests should begin fail because of dynamic shape issues.
• [ ] Either, move to another ops usage that avoids this issue, or upstream the work-around into Keras 3 on the tf backend (add tests, make more robust).
• [ ] Observe test now pass. If we had to upstream a Keras 3 change as part of the work, restore the hack, but leave a TODO to delete with the next Keras release.

If we upstream the change into Keras 3, we may have to keep the compat shim around until we get a new Keras release with our fix. That is ok.