Nx.Random.uniform(key) error when using mps backend in Torchx.

[stewartj3]

Calling Nx.Random.uniform(key) throws an error when using mps backend in Torchx.

Mix.install(
  [
    {:nx, "~> 0.9.2"},
    {:torchx, "~> 0.9.2"}
  ]
)

:erlang.system_flag(:backtrace_depth, 100)
Nx.default_backend({Torchx.Backend, device: :mps})
key = Nx.Random.key(42)
Nx.Random.uniform(key)

Error

** (ArgumentError) cannot perform operation across devices mps and cpu
    (torchx 0.9.2) lib/torchx.ex:469: anonymous fn/2 in Torchx.prepare_tensors_list!/2
    (elixir 1.18.2) lib/enum.ex:1714: Enum."-map/2-lists^map/1-1-"/2
    (elixir 1.18.2) lib/enum.ex:1714: Enum."-map/2-lists^map/1-1-"/2
    (torchx 0.9.2) lib/torchx.ex:464: Torchx.prepare_tensors_list!/2
    (elixir 1.18.2) lib/enum.ex:1840: Enum."-map_reduce/3-lists^mapfoldl/2-0-"/3
    (torchx 0.9.2) lib/torchx.ex:244: Torchx.concatenate/2
    (torchx 0.9.2) lib/torchx/backend.ex:344: Torchx.Backend.concatenate/3
    (nx 0.9.2) lib/nx/defn/evaluator.ex:441: Nx.Defn.Evaluator.eval_apply/4
    (nx 0.9.2) lib/nx/defn/evaluator.ex:256: Nx.Defn.Evaluator.eval/3
    (elixir 1.18.2) lib/enum.ex:1840: Enum."-map_reduce/3-lists^mapfoldl/2-0-"/3
    (nx 0.9.2) lib/nx/defn/evaluator.ex:419: Nx.Defn.Evaluator.eval_apply/4
    (nx 0.9.2) lib/nx/defn/evaluator.ex:256: Nx.Defn.Evaluator.eval/3
    (nx 0.9.2) lib/nx/defn/tree.ex:213: Nx.Defn.Tree.apply_args/4
    (nx 0.9.2) lib/nx/defn/evaluator.ex:419: Nx.Defn.Evaluator.eval_apply/4
    (nx 0.9.2) lib/nx/defn/evaluator.ex:256: Nx.Defn.Evaluator.eval/3
    (elixir 1.18.2) lib/enum.ex:1840: Enum."-map_reduce/3-lists^mapfoldl/2-0-"/3
    (nx 0.9.2) lib/nx/defn/evaluator.ex:419: Nx.Defn.Evaluator.eval_apply/4
    (nx 0.9.2) lib/nx/defn/evaluator.ex:256: Nx.Defn.Evaluator.eval/3
    (elixir 1.18.2) lib/enum.ex:1840: Enum."-map_reduce/3-lists^mapfoldl/2-0-"/3
    (nx 0.9.2) lib/nx/defn/evaluator.ex:419: Nx.Defn.Evaluator.eval_apply/4
    (nx 0.9.2) lib/nx/defn/evaluator.ex:256: Nx.Defn.Evaluator.eval/3
    (elixir 1.18.2) lib/enum.ex:1840: Enum."-map_reduce/3-lists^mapfoldl/2-0-"/3
    (elixir 1.18.2) lib/enum.ex:1840: Enum."-map_reduce/3-lists^mapfoldl/2-0-"/3
    (nx 0.9.2) lib/nx/container.ex:113: Nx.Container.Tuple.traverse/3
    (nx 0.9.2) lib/nx/defn/evaluator.ex:366: Nx.Defn.Evaluator.eval_apply/4
    (nx 0.9.2) lib/nx/defn/evaluator.ex:256: Nx.Defn.Evaluator.eval/3
    (nx 0.9.2) lib/nx/defn/evaluator.ex:319: Nx.Defn.Evaluator.eval_apply/4
    (nx 0.9.2) lib/nx/defn/evaluator.ex:256: Nx.Defn.Evaluator.eval/3
    (elixir 1.18.2) lib/enum.ex:1840: Enum."-map_reduce/3-lists^mapfoldl/2-0-"/3
    (nx 0.9.2) lib/nx/defn/tree.ex:207: Nx.Defn.Tree.apply_args/4
    (nx 0.9.2) lib/nx/defn/evaluator.ex:419: Nx.Defn.Evaluator.eval_apply/4
    (nx 0.9.2) lib/nx/defn/evaluator.ex:256: Nx.Defn.Evaluator.eval/3
    (nx 0.9.2) lib/nx/defn/tree.ex:213: Nx.Defn.Tree.apply_args/4
    (nx 0.9.2) lib/nx/defn/evaluator.ex:419: Nx.Defn.Evaluator.eval_apply/4
    (nx 0.9.2) lib/nx/defn/evaluator.ex:256: Nx.Defn.Evaluator.eval/3
    (elixir 1.18.2) lib/enum.ex:1840: Enum."-map_reduce/3-lists^mapfoldl/2-0-"/3
    (nx 0.9.2) lib/nx/defn/evaluator.ex:419: Nx.Defn.Evaluator.eval_apply/4
    (nx 0.9.2) lib/nx/defn/evaluator.ex:256: Nx.Defn.Evaluator.eval/3
    (elixir 1.18.2) lib/enum.ex:1840: Enum."-map_reduce/3-lists^mapfoldl/2-0-"/3
    (nx 0.9.2) lib/nx/defn/evaluator.ex:419: Nx.Defn.Evaluator.eval_apply/4
    (nx 0.9.2) lib/nx/defn/evaluator.ex:256: Nx.Defn.Evaluator.eval/3
    (elixir 1.18.2) lib/enum.ex:1840: Enum."-map_reduce/3-lists^mapfoldl/2-0-"/3
    (nx 0.9.2) lib/nx/defn/evaluator.ex:419: Nx.Defn.Evaluator.eval_apply/4
    (nx 0.9.2) lib/nx/defn/evaluator.ex:256: Nx.Defn.Evaluator.eval/3
    (elixir 1.18.2) lib/enum.ex:1840: Enum."-map_reduce/3-lists^mapfoldl/2-0-"/3
    (nx 0.9.2) lib/nx/defn/evaluator.ex:419: Nx.Defn.Evaluator.eval_apply/4
    (nx 0.9.2) lib/nx/defn/evaluator.ex:256: Nx.Defn.Evaluator.eval/3
    (elixir 1.18.2) lib/enum.ex:1840: Enum."-map_reduce/3-lists^mapfoldl/2-0-"/3
    (nx 0.9.2) lib/nx/defn/evaluator.ex:419: Nx.Defn.Evaluator.eval_apply/4
    (nx 0.9.2) lib/nx/defn/evaluator.ex:256: Nx.Defn.Evaluator.eval/3
    (elixir 1.18.2) lib/enum.ex:1840: Enum."-map_reduce/3-lists^mapfoldl/2-0-"/3
    (nx 0.9.2) lib/nx/defn/evaluator.ex:419: Nx.Defn.Evaluator.eval_apply/4
    (nx 0.9.2) lib/nx/defn/evaluator.ex:256: Nx.Defn.Evaluator.eval/3
    (elixir 1.18.2) lib/enum.ex:1840: Enum."-map_reduce/3-lists^mapfoldl/2-0-"/3
    (nx 0.9.2) lib/nx/container.ex:113: Nx.Container.Tuple.traverse/3
    (nx 0.9.2) lib/nx/defn/evaluator.ex:70: anonymous fn/5 in Nx.Defn.Evaluator.__compile__/4
    (nx 0.9.2) lib/nx/defn.ex:452: Nx.Defn.do_jit_apply/3

[polvalente]

@stewartj3 I locked the conversation, but the issue will still be investigated. For some reason the issue was being targeted for spam.