Nx with multiple processes

[bill huang]

Hi,

I'm coming from Python/JAX, and quite new to Elixir. The overall user experience of Elixir/Nx has been great. What I care most is that Elixir offers distributed computations of out the box. But Nx's document doesn't touch too much on how things interact when using multiple processes with Nx.

Here are some problems:

- Is it zero-copy to send Nx tensors to other processes? I know BinaryBackend enables zero-copy, but could other backends do the same thing?

- Is there a way to easily control some commonly used environment variables for each process, like CUDA_VISIBLE_DEVICES, so that I can assign one process for each GPU.

- Does Nx/EXLA preallocates memory like JAX does? In JAX, you need to set environment variables like XLA_PYTHON_CLIENT_PREALLOCATE to allow multiple JAX processes to work together. Since EXLA is using XLA, I think Nx/EXLA might have similar behavior to JAX/Tensorflow, but I'm not so sure.

Thanks

[José Valim]

All excellent questions. I will add notes to the docs. Answers also inline:

1. Yes.

2. You can use Nx.default_backend({EXLA.Backend, client: :cuda, device: 1}) and set a particular backend for that process. The jit/compile functions also support client/device pairs as options

3. Yes. This one is briefly mentioned here: https://hexdocs.pm/exla/EXLA.html#module-client-options

--
You received this message because you are subscribed to the Google Groups "Numerical Elixir (Nx)" group.
To unsubscribe from this group and stop receiving emails from it, send an email to elixir-nx+...@googlegroups.com.
To view this discussion on the web visit https://groups.google.com/d/msgid/elixir-nx/1dd230fc-e67b-4a8b-ac01-daf7378ad85fn%40googlegroups.com.

[José Valim]

Apologies, it should be device_id on the second answer.

[bill huang]

Thanks for the quick reply. It's good to know these.