Candlex: an Nx.Backend for candle ML framework

[grzuy]

UPDATE: Released in https://hex.pm/packages/candlex and kept as separate repo in https://github.com/mimiquate/candlex.

---

An experimental work-in-progress exploring the possibility of implementing an Nx.Backend for candle-core rust crate.

Note that https://github.com/huggingface/candle is still in early development and a moving target.

Still a lot of Nx.Backend behavior callbacks pending implementation, so plan it continue making progress here to cover as much as possible.

No CUDA at the moment, only CPU.
UPDATE: Decent CUDA support for a small subset of backend functions.
UPDATE2: Decent CUDA support for slightly more than half backend callbacks.

UPDATE3:

• About ~90 out of 117 Nx.Backend callbacks have some form of implementation (complete or partial).
• Device support: Yes, CPU and GPU/CUDA
• Axon's examples/vision/mnist.exs and examples/generative/fashionmnist_autoencoder.exs are working using Candlex.Backend with similar performance to Torchx both in cpu and cuda. Other examples failing because of still some missing callback or callback option.

Pending:

• Consider rustler_precompiled to avoid dependency on the rust toolchain
• Update github actions to run candlex tests
• Cover more Nx.Backend callbacks (tensor operations)
• Better support for f16 and bf16
• Make it work with more elaborated Axon model examples

---

Updates to upstream candle as part of this effort:
huggingface/candle#741
huggingface/candle#925
huggingface/candle#988
huggingface/candle#996
huggingface/candle#1037

---

UPDATE 4: Also published in a separate repo: https://github.com/mimiquate/candlex.

[grzuy]

This is awesome!

I have dropped some comments and ultimately we will only know if we should merge this or not based on how complete Candle is. If it is missing many operations, then perhaps it is better to wait, but I don't see any problems from a glance. We will need to discuss between us too.

Hi guys, updated this pull request description with a summary of the current status of this effort.

[grzuy]

@josevalim @polvalente Got this PR to a point it might make sense to slow down a bit and wait for some feedback from you guys on whether you see a good fit for that inside Nx or release is separtely for now.

Thanks in advance!

[grzuy]

While this package is in consideration for inclusion as part of this repo, found it valuable to publish separately too (https://github.com/mimiquate/candlex), so it's easier for anyone to navigate it, try it out, open issues, propose improvements via PRs, etc. In case gets accepted here we can just archive the separate repo.

Plan to keep this PR up to date with any updates upstream.

Thanks again!

[josevalim]

Hi @grzuy! I did another pass over it and it looks really solid!

However, after working on Nx over the last few days, i believe we should not merge this. Adding new features to Nx already requires us to add code to both EXLA and Torchx. Although both projects are in C++, it is already time consuming. Adding a Rust project means both additional tooling and an additional language, and I am worried it will harm our productivity.

Therefore I suggest to keep it in a separate project (but feel free to submit a PR to the README to link to it). :)

Have a good one!

[grzuy]

Hi @josevalim,

Reasonable. Sounds good.
Thank you for the feedback!

Will PR to README later.

Thanks!