`Nx.Serving.batched_run` not working - always hangs

[GenericJam]

I could not get batched_run to work under any circumstances using EMLX. I had to use Nx.Serving.run instead.

To reproduce, run this notebook - https://hexdocs.pm/bumblebee/stable_diffusion.html

Modified like this:

Mix.install([
  {:emlx, github: "elixir-nx/emlx", branch: "main"},
  {:bumblebee, github: "elixir-nx/bumblebee"},
  {:nx, "~> 0.10.0"},
  {:kino, "~> 0.18.0"}
])

Nx.global_default_backend({EMLX.Backend, device: :gpu})

With testing it hangs right at the spot where batched_run is called. The EMLX readme describes a potential NIF deadlock under some circumstances which would describe the behaviour well.

[polvalente]

The NIF deadlock comment was outdated and has been updated in the documentation, so I don't think that is the root cause.
The set of modifications you suggest by itself does not work at the very least due to the opts = [params_variant: "fp16", type: :bf16, backend: EXLA.Backend] line, which should also be modified.

Do you have the full version of the livebook you tried to run?

[polvalente]

The fix needed for this is in Nx, so let's keep this open until Nx is updated for EMLX