Nx ecosystem 0.12 library updates

News News & Updates
, , , ,
1

Nx, EXLA and Torchx 0.12 have just been published!

This update set comes with the new Nx.block abstraction and the EXLA.CustomCall protocol, which together will allow users to provide native C or CUDA implementations for certain sections of Nx code.

EMLX 0.3 (the official MLX backend for Nx) is also out, with:

  1. Support for the more recent concurrency model provided by MLX
  2. EMLX.Fast for a few fused operations (used in the upcoming emlx_axon library, pending Axon and Bumblebee releases in the next few days)
  3. Support for quantized types
  4. A zero-copy implementation of `Nx.to_binary/1`.
16 Likes
2

:clap: :clap: :clap: :clap: :clap: :clap:

1 Like
3

emlx_axon has been published!

By using it (see this benchmark for an example), we can accelerate certain Bumblebee models without rewriting code for performance when using EMLX! For example, the Qwen3-0.6B-4bit model goes from ~10 tokens/s to ~42 tokens/s on my machine just by using the provided Axon node rewrites!

3 Likes
4

Hi @polvalente, I’m trying to bring Nx + EMLX to iOS. Do you have any suggestions?

Currently, I support inference first, and plan to add basic on-device training later.

5

@manhvu , @GenericJam was able to run EMLX on device, which means all of Nx can run! We’re working to upstream what he needs so that no patches are needed, just config!

2 Likes
6

Thanks for the response. I added support for Nx to my framework, and over the next few days I’ll verify and test everything.

For inference, I currently support both the Nx ecosystem and ONNX Runtime (for cross-platform support), but I’ll add Nx support on Android again soon.

1 Like
7

The answer is EMLX on iOS and Vulkan on Android. These changes should be upstream shortly.

8

EMLX for iOS because it can use Core ML. I’m not sure about Vulkan yet; it doesn’t seem stable enough. I’m still looking for a way to use the mobile GPU on both iOS and Android for on-device training for edge AI.

9

You can try this GitHub - GenericJam/nx_vulkan: Vulkan compute Nx.Backend — cross-platform GPU acceleration for Nx · GitHub
Example lib GitHub - GenericJam/nxeigen_probe: Reference Mob app — cross-platform YOLOv8n via EMLX (iOS) + NxVulkan (Android) · GitHub

1 Like
10

Nice, you’ve done that :ok_hand: I looked for the best solution, but it sounds like I need to separate them. Currently, only WebGPU can run on both platforms, so I’ll try to adapt Nx to it.

11

It looks like @polvalente was right all along and NxIREE is the way to harness Vulkan.

12

Yes, it sounds like he was right.

Currently, I’m trying to bring CubeCL in as a backend for Nx. These days, Rust is quite good for ML/AI and big data, so I want to work with a combination of Elixir and Rust.