Remove take_along_axis and take

[josevalim]

take can directly call gather, no need for optional callbacks. take_along_axis should have an optional callback.

[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]

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.

[Benjamin-Philip]

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

[polvalente]

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