Nx for Erlang, LFE and other BEAM languages

I honestly don’t know. I have mainly so far done it for personal interest.

From what I can see and through some simple tests Nx seems to work well with Erlang. I have at least testade the simple examples given in the “Nx for absolute beginners” and “Up and running Nx”. I made a very simple Erlang module called nx which just passes all the calls through to Nx. As there a re no macros in Nx there are no problems calling it from Erlang.

The thing to remember is that Elixir runs on top of Erlang/BEAM so everything it does goes through them so it can all be done directly from Erlang. The question is more what the interface looks like. As soon as it includes Elixir macros then it becomes more complex to handle. And the module names don’t help.

I must also admit I have another interest. I have been testing from LFE This is about the same as for Erlang but if I ever get to the stage of needing macros then LFE is much easier than Erlang. And I like the parentheses.

We shall see how it goes with the book.

I think the only thing that you miss out from non-Elixir is the defn and friends macros that are defined in Nx.Defn.

Perhaps there’s a way to avoid needing them by using a more verbose and manual approach, though.
Feel free to ping us at machine-learning in the EEF Slack to discuss this. José probably has the answer from the top of his head (if it’s possible to avoid defn macros in favor of manual non-macro calls), but I can try to help as well!

Worst case scenario you can get the same effect by manually compiling anonymous functions with the equivalent of:

iex(13)> defmodule NonDefn do
...(13)>   def add(tensor, i) do
...(13)>     Nx.Defn.jit_apply(fn a, b -> IO.inspect(Nx.add(a, b)) end, [tensor, i])
...(13)>   end
...(13)> end
warning: redefining module NonDefn (current version defined in memory)
  iex:13

{:module, NonDefn,
 <<70, 79, 82, 49, 0, 0, 6, 132, 66, 69, 65, 77, 65, 116, 85, 56, 0, 0, 0, 236,
   0, 0, 0, 23, 14, 69, 108, 105, 120, 105, 114, 46, 78, 111, 110, 68, 101, 102,
   110, 8, 95, 95, 105, 110, 102, 111, 95, ...>>, {:add, 2}}
iex(14)> NonDefn.add(10, 11)                                                        
#Nx.Tensor<
  s64
  
  Nx.Defn.Expr
  parameter a:0   s64
  parameter b:1   s64
  c = add a, b    s64
>
#Nx.Tensor<
  s64
  21
>
iex(15)> NonDefn.add(1, 2)  
#Nx.Tensor<
  s64
  
  Nx.Defn.Expr
  parameter a:0   s64
  parameter b:1   s64
  c = add a, b    s64
>
#Nx.Tensor<
  s64
  3
>

There are some features that you would miss out, like while, cond and case which rely on macros.
However, the main effect of having a condensed Nx.Defn.Expr graph that can be compiled with the compiler of choice is still there.

6 comments

/nx #split-thread #lfe

19 1391 6

2023-03-30 21:34:15 UTC

Most Liked

rvirding

Creator of Erlang

That is very interesting. Is it documented how the defn macros are expanded? With a bit of trickery you could probably a lot of that as well using parse transforms. For LFE it macros should be able to handle that.

It also starts to explain a jit/2 function defined in the erlang nx module I have been working on. It was originally written by Duncacn McGreggor and I have playing with it.

As I said I don’t know how much interest there in the erlang community.

Post #2

josevalim

Creator of Elixir

Moved to a separate thread so we do not disturb the book readers.

@polvalente, I edited your comments to use jit_apply because __runtime__ is precisely that (and I even changed main to be the same). You are right that while/cond are the only private APIs at the moment but we can cross that bridge when necessary.

Post #6

polvalente

Nx Core Team

I don’t think there is a lot of documentation around that. But they are basically regular Elixir macros, and some of the macros defined in Nx.Defn.Kernel check that they are being called inside a defn/defnp-defined function.

Other than these restrictions on Nx.Defn.Kernel, defn is basically a way for people to write regular elixir code that is passed into the Nx.Defn compiler (which I circumvented in my example by using jit, which does kind of the same thing).

I do think that putting some effort into documenting this even if through comments would benefit the BEAM community. I imagine that having an erlang-supported way to use the defn-specific macros would be the only missing link (or at least the most important one).

Post #3