Initializer does not respect Nx default backend about axon HOT 6 CLOSED

I believe this is an issue in Torchx itself and not in Axon. The Axon initializers are in the end all implemented through Defn, so if the backend options aren't being fetched and passed on properly, chances are this is related to Torchx.

An easy way to test this is to use EXLA in a machine with cuda support and set the default backend to {EXLA.Backend, client: :cuda}.

from axon.

Thanks @polvalente

I have investigated this issue a bit more and it seems like in the initiation process a tensor is being created using Nx.BinaryBacked and then this tensor get's passed to from_nx function in Torchx.Backend. Now based on torchx/backend.ex#L1808, Torchx will call Nx.backend_transfer but it does not know about the default options anymore and will create cpu tensor instead of mps or any other device that the user has specified.

I can see two possible solutions:

1. Investigate why the initializer is making a tensor in Nx.BinaryBackend and force it to use Nx.default_backend instead.
2. Create a new configuration parameter for Torchx which stores the default options for when the from_nx function wants to call backend_transfer(tensor, {Torchx.Backend, default_options}) in torchx/backend.ex#L1808

from axon.

I just ran into the same issue - did you ever try out either of your two proposed solutions?

from axon.

Hi @benswift

I actually changed a line in Torchex Backend to forcefully create tensors with the mps device for my test and it worked as expected.

But I don't yet understand why tensors are created in BinaryBackend in the first place...

from axon.

Ah, ok, thanks for the update :)

from axon.

Yes, this is a torchx issue, so closing for now!

from axon.