Interactive coding notebooks with Elixir, Nx and Livebook.

27.09.2023Raffael Schneider
Tech Elixir Livebook Numerical Elixir TensorFlow GPU PyTorch ML-Ops

!— !—

Since its humble beginnings, Elixir has evolved into a popular programming language that focuses on the needs of modern concurrent distributed systems. We’ve written about Elixir and the Elixir ecosystem several times here at b-nova, and find that it solves several problem domains particularly well. For example, web applications can be efficiently and easily mapped to Elixir’s well-known Phoenix web framework.

There has been an intention within the ecosystem and its community for some time to conquer the world of scientific programming and data analysis. The cover of this intention is the so-called project Nx, which simply stands for Numerical Elixir. Under Nx fall various sub-projects, including Livebook, a sort of alternative to the well-known Jupyter Notebook. In this blog article we will go into more detail about Nx as well as Livebook and what you can do with it.

Elixir also understands numerical computing thanks to Nx

Elixir’s foray into the field of numerical computing was neither sudden nor unexpected. The introduction of projects like Matrex, an Elixir library for matrix computations that featured a C-based CBLAS implementation, showed Elixir’s potential in this niche and inspired the community to strive for a more comprehensive and robust solution.

Thus was born the Nx project. Nx, as mentioned above, stands for Numerical Elixir and is, simply put, a library for numerical computing. Especially in recent years, when ML and AI have become particularly well known, not least through ChatGPT, the project naturally aims to establish Elixir in areas such as machine learning, artificial intelligence, and data science.

At the heart of Nx are tensors, which play a central role in many modern scientific and machine learning applications. Without getting too deep into the subject here, recall that a tensor is simply a multidimensional array. Nx provides a number of operations as well as abstractions to operate on these tensors in a CPU-efficient manner. Nx can also perform automatic differentiation, which makes it useful for machine learning and similar applications, similar to what NumPy, PyTorch or TensorFlow do.

José Valim, the creator of Elixir who is also heavily involved in the Nx project, summarized Nx on Hacker News as follows:

“When we started the Numerical Elixir effort, we were excited about the possibilities of mixing projects like Google XLA’s (from Tensorflow) and LibTorch (from PyTorch) with the Erlang VM capabilities to run concurrent, distributed, and fault-tolerant software.” (Source: https://news.ycombinator.com/item?id=35525661)

This connectivity is provided by Erlang’s NIF (Native Implemented Functions) system, which allows Elixir programs to call native functions in other languages such as C or C++. Specifically, Nx uses the XLA (Accelerated Linear Algebra) library, a specialized linear algebra library developed by Google and released as part of the TensorFlow project. XLA is designed to perform high-performance linear algebra computations on a variety of hardware platforms, including CPUs, GPUs, and even TPUs (Tensor Processing Units).

Short comparison with PyTorch.

Here is a simple code snippet in PyTorch that illustrates the use of a tensor operation. For this example, let’s assume we have two tensors and want to perform some basic operations on them.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
import torch

# Creates two tensors
tensor1 = torch.tensor([[1, 2], [3, 4]])
tensor2 = torch.tensor([[5, 6], [7, 8]])

# Performs some basic operations
sum = tensor1 + tensor2
product = tensor1 * tensor2
difference = tensor1 - tensor2

# Prints the results
print('Sum:\n', sum)
print('Product:\n', product)
print('Difference:\n', difference)

When this Python script is executed, you should get the following output:

1
2
3
4
5
6
7
8
9
sum:
 tensor([[ 6, 8],
        [10, 12]])
Product:
 tensor([[ 5, 12],
        [21, 32]])
Difference:
 tensor([[-4, -4],
        [-4, -4]])

Here is the equivalent code to perform the same operations in Elixir’s Nx:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
defmodule MyNxModule do
  require Nx.Defn

  import Nx.Defn

  defn tensor_operations do
  # Creates two tensors
  tensor1 = Nx.tensor([[1, 2], [3, 4]])
  tensor2 = Nx.tensor([[5, 6], [7, 8]])

      # Performs some basic operations
      sum = Nx.add(tensor1, tensor2)
      product = Nx.multiply(tensor1, tensor2)
      difference = Nx.subtract(tensor1, tensor2)

      {sum, product, difference}
  end
end

In this code, we create two 2x2 tensors, tensor1 and tensor2, using the Nx.tensor() function. We then perform a series of operations on these tensors: addition (Nx.add()), multiplication (Nx.multiply()), and subtraction (Nx.subtract()). Note that these functions do not modify the tensors themselves (since Elixir is an immutability-based language), but return new tensors representing the result of the operations.

Since the Elixir runtime, called BEAM, provides an interactive prompt, here we evaluate the results of MyNxModule as follows:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
iex> MyNxModule.tensor_operations() 
{#Nx.Tensor<
  s64[2][2]
  [
    [6, 8],
    [10, 12]
  ]
>, #Nx.Tensor<
  s64[2][2]
  [
    [5, 12],
    [21, 32]
  ]
>, #Nx.Tensor<
  s64[2][2]
  [
    [-4, -4],
    [-4, -4]
  ]
>}

Here we see that the tensor_operations() function returns a tuple with three tensors representing the sum, product, and difference of the original tensors.

As you can see, the basic operations in Nx are very similar to those in PyTorch. Both use tensors as their basic data structure and provide similar APIs for working with these tensors. The main difference is that PyTorch is written in Python and therefore uses the syntactic and semantic features of that language, while Nx is written in Elixir and uses the features of that language, such as immutability and pattern matching.

The Nx Ecosystem

Nx is primarily the core library for computing numerical problems, especially tensors. But Nx is also an umbrella project for other subprojects.

  • EXLA: EXLA is the Elixir client library for Google’s XLA (Accelerated Linear Algebra), a specialized compiler library for linear algebra and machine learning. EXLA enables Elixir programs to use XLA to efficiently perform Nx operations, especially on hardware accelerators such as GPUs and TPUs.
  • Axon: Axon is a library for neural networks and deep learning built on top of Nx. It provides a high-level API for defining, training, and running neural networks. Axon also supports many common features of deep learning frameworks, such as different types of layers, activation functions, optimization techniques, and loss functions.
  • Explorer: In the context of Elixir and Nx, Explorer could potentially be a tool for data exploration and visualization, allowing one to “explore” datasets and the results of computations and model training.
  • Scholar: Scholar is a library for traditional machine learning tools built on top of Nx. It implements several algorithms for classification, regression, clustering, dimension reduction, metrics, and preprocessing.
  • Bumblebee: Bumblebee is a library that provides pretrained neural network models based on Axon. It includes integration with HuggingFace models, allowing anyone to download models and perform ML tasks with just a few lines of code. With Bumblebee, users can easily use powerful models like BERT for their own applications.
  • Livebook: Livebook is a web-based interface for interactive and collaborative programming with Elixir. It resembles Jupyter notebooks and allows users to write and run Elixir code in an interactive format with text, code, live output, and visualizations. Livebook leverages Nx and its subprojects to provide powerful tools for scientific and data-driven computing in Elixir. Livebook will be our entry point into the world of Elixir and Nx in this TechUp.

We had already mentioned José Valim, he himself summarized the Nx ecosystem as seen in the following screenshot:

!— Screenshot 2023-08-01 at 09.07.52.png !—

From Nx to Livebook

With Nx, the Elixir community has shown its ambitions in the field of numerical computing. But something important was still missing from it, namely tooling, which has been long established in the Python ecosystem. They wanted to create a complete, interactive and user-friendly environment that supports the entire data analysis workflow. Thus, the idea of Livebook was born.

Livebook features

  • Interactive Coding: Like other notebook tools, Livebook allows you to combine code and comments in a single file that is divided into different sections called cells. Code can be written in a cell and then executed to immediately display the result in the same cell. This interactivity enables exploratory programming and makes it easier to debug and test code.

  • Support for Elixir and Erlang/OTP: Livebook is written in Elixir and takes full advantage of that language and the Erlang/OTP system. You can run any Elixir code in Livebook and access the entire Elixir standard library as well as third-party packages. You can also take advantage of features of the Erlang/OTP platform, such as the ability to create concurrent processes and access to OTP behaviours such as GenServer and Supervisor.

  • Integration with Nx: Livebook is designed to work closely with Nx. You can run Nx code directly in Livebook and view the results in an easy-to-understand, formatted form. This makes Livebook an excellent tool for working with numerical data and machine learning.

  • Data Visualization: One of the outstanding features of Livebook is its built-in support for data visualization. Livebook can create different types of charts and graphs to visualize your data and help you see patterns and trends in your data. These visualizations are interactive, so you can zoom, scroll, and perform other actions to view the data from different perspectives.

  • Live updates and multi-session: Livebook supports live updates, so multiple users can work on a notebook at the same time and see each other’s changes in real time. Originally, Livebook apps were intended for long-running applications, with only one instance of a Livebook app running at a time. However, since Livebook has native support for multiple users, all users accessing an app share the same instance. The new release introduces multi-session Livebook apps, where each user gets an exclusive version of the app for themselves. These apps can run indefinitely, similar to single-session apps, but will mostly process user input, execute multiple statements, and then exit.

  • Persistence and portability: Livebook notebooks are simply Elixir script files with the extension .livemd. This versions well with Git and reads easier than Jupyter’s JSON-based file format .ipynb.

  • Native Secrets Management: Livebook has had the ability to store Secrets encrypted in Livebook for some time.

  • Livebook Desktop App: There has also been for a year now the ability to allow non-techies to view and run existing Livebooks without having to set up either Elixir or Erlang/OTP. This brings Livebook to a whole new group of users.

Local setup of Livebook

Installing Livebook is a fairly simple process and can be done on various operating systems such as GNU/Linux, macOS or even Windows. Here is a step-by-step guide to installing Livebook on a system running Elixir and Erlang/OTP:

Step 1: Install Elixir and Erlang/OTP

Livebook is written in Elixir and therefore requires Elixir and Erlang/OTP to run. It is best to check if Elixir is already installed on our system. If you have been through one of the Elixir TechUps, chances are high. If not, you can find out briefly with the following command:

1
$ elixir --version

If Elixir and Erlang/OTP are not already installed, you can install Elixir and the Erlang/OTP as described in the following two links:

ℹ️ Tip: I personally use asdf as a version manager for this, which can be used to easily and elegantly attract exactly the right runtime of a given technology or tech stack. Of course we also have a TechUp for this: Simplify your workflow with a version manager like asdf | b-nova.

Step 2: Clone the Livebook repository from GitHub

Now we clone the livebook repository from GitHub. This can be done with the following command:

1
$ git clone https://github.com/livebook-dev/livebook.git

This command creates a new folder named livebook in the current directory and copies all the files from the livebook repository to this folder.

Step 3: Change to the livebook directory and install the dependencies

Now we change to the livebook directory and install the Elixir dependencies specified by the project (see mix.exs) using the following commands:

1
2
cd livebook
mix deps.get --only prod

Step 4: Start Livebook

Now we can start Livebook with the following command:

1
MIX_ENV=prod mix phx.server

This command starts the Livebook server on the local system. If you see the following output, you have successfully started Livebook. The output indicates that the server is running and at which URL it is accessible,

1
2
3
...
Generated livebook app
[Livebook] Application running at http://localhost:8080/?token=jaiusa5vojhdmejopubxcegdw7kttgyy

The given URL http://localhost:8080 has to be opened in a web browser of your choice. Preferably with the token, otherwise you can simply transfer the token in the web interface. In this case jaiusa5vojhdmejopubxcegdw7kttgyy (will be regenerated at every start and the value will be different).

!— Screenshot 2023-07-31 at 15.18.06.png !—

ℹ️ To note: Livebook is an open source project and it is always being developed on. Best to look at the GitHub repository in case something doesn’t work.

Something else about CUDA and graphics cards

Before we get into filling our Livebooks, it’s worth checking that the machine we’re using has a supported NVIDIA graphics card installed, as this can speed up calculations considerably. As you may already know, numerical operations often use decimal algorithms, which benefit enormously from GPUs. If you have a supported graphics card, you should follow EXLA’s instructions to select XLA_TARGET and install the appropriate CUDA version for your graphics card. After installing CUDA, the correct XLA_TARGET environment variable can be set as follows:

Value Target environment
tpu libtpu
cuda120 CUDA 12.0+, cuDNN 8.8+ (recommended)
cuda118 CUDA 11.8+, cuDNN 8.7+ (recommended)
cuda114 CUDA 11.4+, cuDNN 8.2+
cuda111 CUDA 11.1+, cuDNN 8.0.5+
cuda CUDA x.y, cuDNN (building from source only)
rocm ROCm (building from source only)

The XLA TARGET environment variable should then be set in the livebook. For example with the value cuda120.

If there is no supported NVIDIA graphics card, this is not a problem. Neural networks can also be used on the CPU, but the calculations will be slower then. For example, on a MacBook with an M1 Max CPU, generating two images with Stable Diffusion will take several minutes. However, all other tasks, such as text generation, will complete within a few seconds. So it will be slower, but it will work!

What makes Livebook different

On Hacker News the question came up why the Elixir runtime, the so-called BEAM, is particularly well suited for data processing. José Valim answered that Erlang offers competition within an operating system process and distribution over several nodes. In Livebook, “Concurrency” is used to allow communication with results during execution. “Branched Sections” allows multiple experiments to be run in parallel in the same notebook. Distribution in Erlang is similar to concurrency and he shows an example of how to change an ML model from concurrent to distributed. He emphasizes that it makes sense to default to concurrency because serial processing can be more expensive, especially for data analysis. Functional programming in Erlang is beneficial for data analysis, and livebook notebooks are easy to reproduce.

!— Screenshot 2023-08-01 at 09.15.42.png !—

See also this video on “Meta programmable functional notebooks with Livebook” on Youtube.

Conclusion

In my eyes, the Nx project is a real high-flyer and definitely has what it takes and the maturity to compete with established tooling. It is clear to me that Elixir enjoys a niche existence and is definitely not a widespread ecosystem. And yet, you have to give Nx and especially Livebook credit for a quality and polished feature set. 🔥

This TechUp has been translated by our automatic markdown translator. 🔥