Nx for Erlang, LFE and other BEAM languages

Elixir Framework Forums Nx Forum
, ,
1

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.

4 Likes
2

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.

3 Likes
3

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).

2 Likes
4

I did a quick test of this in both Erlang and LFE and it worked. Are the jit and jit_apply calls expanded in the macro expansion at compile time or is the reultant code basically what you showed?

2 Likes
5

I took a closer look into what defn actually boils down to.
And I think you can get away with doing something like:

defmodule MyMod do
  def f(x, y) do
    Nx.add(x, y)
  end

  def f_compiled_like_defn(x, y) do
    Nx.Defn.jit_apply(fn x, y -> f(x, y) end, [x, y], on_conflict: :reuse)
  end
end

This only leaves the hole of the Defn.Kernel macros like while and friends that would
require some form of non-macro public interface :slight_smile:

Note that the example above uses no macros (aside from def and defmodule :stuck_out_tongue:), and can be extended so that f is any combination of Nx and Nx.LinAlg functions, or functions built through them.

Also, defining f isn’t required, you could pass an anonymous function directly (or function captures, if those exist in Erlang)

2 Likes
6

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.

3 Likes
7

Thanks!

I guess cond is mostly a syntax sugar/optimization given Nx.select. while is a bit trickier, but vectorization does fill that hole a bit too :wink:

1 Like