Remove take_along_axis and take about nx

[Benjamin-Philip]

Yeah! To be clear, take can be replaced by gather (look at how EXLA does it) at the Nx level. take_along_axis also can be, but as an optional callback (again look at EXLA). Torchx has a simpler operation for take_along_axis

Thanks! I understood that from José's original wording as well.

[polvalente]

Torchx should have the specific implementation. EXLA and BinaryBackend can use the default that delegates to gather

[Benjamin-Philip]

Could I have a stab at this?

[polvalente]

Yeah! To be clear, take can be replaced by gather (look at how EXLA does it) at the Nx level. take_along_axis also can be, but as an optional callback (again look at EXLA). Torchx has a simpler operation for take_along_axis

[Benjamin-Philip]

Quick question: for which backend(s) do you want to implement the take_along_axis optional axis callback?