Could I have a stab at this?

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

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.

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

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