Moins d’un an après que l’équipe de développement ait présenté pour la première fois le projet Nx, la première version publique de la bibliothèque est désormais disponible sur GitHub. Nx 0.1.0 doit permettre aux développeurs d’Elixir de combiner la programmation fonctionnelle avec les compilateurs Tensor afin d’ouvrir de nouveaux champs d’application tels que le calcul numérique et l’apprentissage automatique.

Exploiter Elixir et la VM Erlang

Bien que l’équipe de développement Nx de José Valim et Sean Moriarity ait été consciente du fait qu’un langage de programmation fonctionnel tel qu’Elixir pouvait poser problème lors de calculs numériques, en particulier lorsqu’il s’agit de l’immutabilité (Immutability) en travaillant avec de grands blocs de mémoire, les avantages potentiels leur ont semblé séduisants. En effet, le langage de programmation Elixir, qui a récemment introduit la recompilation sémantique avec la version 1.13, et sa VM Erlang sous-jacente promettent une faible latence lors du fonctionnement dans des systèmes distribués et tolérants aux pannes, ce qui permet également de construire des applications évolutives de manière flexible et plus faciles à maintenir.

La bibliothèque Nx, conçue pour les tableaux multidimensionnels (tenseurs) et la compilation juste à temps de l’élixir numérique, tant sur les CPU que sur les GPU, s’oriente pour cette raison vers des modèles tels que NumPy et JAX de l’écosystème Python. JAX, en particulier, est fondamentalement orienté vers la programmation fonctionnelle et les principes d’immutabilité.

Nx doit créer sur cette base les conditions nécessaires pour travailler également à des niveaux d’abstraction plus élevés, comme par exemple les définitions numériques (defn) ou des graphiques de calcul, les avantages de la programmation fonctionnelle. Au sein de defnqui est un sous-ensemble d’Elixir adapté au calcul numérique, Nx permet également d’utiliser les opérateurs réguliers d’Elixir et de les traduire en opérations Tensor équivalentes. Les développeurs disposent ainsi de fonctions de langage et de types de données tels que les macros, l’opérateur Pipe, le Pattern Matching ou les Maps.

defmodule MyModule do
  import Nx.Defn

  defn softmax
    normalized = t - Nx.max
    Nx.exp(normalized) / Nx.sum(Nx.exp(normalized))
  end
end

Après l’appel, le code du listing précédent compile les types et les formes des arguments, par exemple via le compilateur XLA de Google, en un code optimisé qui peut être exécuté sur des CPU, des GPU ou même des Cloud TPU. Les tenseurs multidimensionnels peuvent être créés avec Nx à partir d’entiers non signés (u8, u16, u32, u64) et signés (s8, s16, s32, s64) ainsi que de nombres flottants non signés (f16, f32, f64) et signés (bf16).

Nouvelles fonctions pour les projets ML

Afin de développer Nx pour une utilisation productive dans des projets d’apprentissage automatique, l’équipe a notamment implémenté ces derniers mois des fonctions de streaming qui permettent de charger une application dans des GPU et des TPU pendant qu’elle reçoit un lot d’entrées. Ces fonctions devraient contribuer à une exécution plus efficace de l’inférence, mais aussi se prêter à l’apprentissage distribué.

En vue des réseaux neuronaux récurrents (RNN) et des modèles pour la reconnaissance vocale, l’analyse sémantique, la traduction en langue des signes et les domaines d’application comparables, les créateurs de Nx ont ajouté aux définitions numériques while-ont été ajoutées aux boucles. Celles-ci doivent permettre le désentrelacement statique et dynamique des boucles.

Axon : Réseaux neuronaux basés sur Nx avec Elixir

Afin de pouvoir tester la bibliothèque Nx et son moteur de dérivation automatique contre de nombreux réseaux neuronaux, Sean Moriarity a développé Axon. Axon met à disposition une API fonctionnelle qui fonctionne comme une API de bas niveau pour les définitions numériques (defn) et crée en même temps les conditions pour d’autres API découplées les unes des autres pour la création de modèles, l’optimisation et l’entraînement, afin de pouvoir utiliser des réseaux neuronaux basés sur Nx avec Elixir. Le listing suivant montre le modèle d’un réseau neuronal convolutionnel (CNN) implémenté avec Axon pour l’apprentissage et la classification du jeu de données CIFAR-10 :

Axon.input(input_shape)
|> Axon.conv(32, kernel_size: {3, 3}, activation: :relu)
|> Axon.batch_norm()
|> Axon.max_pool(kernel_size: {2, 2})
|> Axon.conv(64, kernel_size: {3, 3}, activation: :relu)
|> Axon.batch_norm()
|> Axon.max_pool(kernel_size: {2, 2})
|> Axon.flatten()
|> Axon.dense(64, activation: :relu)
|> Axon.dropout(rate: 0.5)
|> Axon.dense(10, activation: :softmax)

Pour plus d’informations sur Nx 0.1.0 et sur l’ensemble du projet, veuillez consulter le billet de blog consacré à la sortie de la première version prête pour la production ainsi que l’annonce initiale de Nx (Numerical Elixir). Les dépôts publics du projet peuvent être consultés sur GitHub.