Allocate intermediate tensors in the same device as the input tensor

[polvalente]

Some functions in Torchx.Backend create intermediate Iotas and similar tensors. These should be allocated in the same device as the function's input tensor.

[josevalim]

We should probably create a helper like we did for vectorization, where do:

with_device([a, b, c], fn a, b, c ->

end)

So we can wrap our backend functions by such helper. Pinging @grzuy because he likely needs similar for candlex.

[ffloyd]

Hey!

Any update on this?

I think this is an important one because:

• a lot of developers have Apple Silicon (and many has it provided by their companies)
• due to unified memory Mac Book Pro with 30+ GB of memory is capable of ~22b model serving. I tested it with ollama and codestral 22b for example. It's not blazing fast, but my 3080 RTX cannot even run such models.
• GPU setup with 30+ Gb is either more expensive than macbook (A100) or somewhat tricky (build a separate machine with several old Tesla GPUs, nvlink them, pray that you LLM rig will work).

Considering this, a huge audience who would like to experiment with LLMs will be forced to use Python or pay for hosted GPUs. :(

[josevalim]

For Apple Silicon, there is: https://github.com/elixir-nx/emlx - if you would like to use Torchx to target Metal, a pull request is welcome.