? Pourquoi combiner Dart et Jax pour l'apprentissage automatique?

Lors de la création d'applications, la sélection des bons outils est cruciale. Vous voulez des performances élevées, un développement facile et un déploiement multiplateforme transparent. Les cadres populaires offrent des compromis:

• c fournit de la vitesse mais peut ralentir le développement.
• dart (avec Flutter) est plus lent mais simplifie la gestion de la mémoire et le développement multiplateforme.

Mais voici le Catch: La plupart des frameworks manquent de support robuste Machine Learning (ML) . Cet écart existe parce que ces cadres sont antérieurs à la flèche de l'IA. La question est:

Comment pouvons-nous intégrer efficacement ML dans les applications?

Les solutions communes comme onnx runtime permettent d'exporter des modèles ML pour l'intégration des applications, mais ils ne sont pas optimisés pour les CPU ou suffisamment flexibles pour les algorithmes généralisés.

Entrée jax , une bibliothèque Python qui:

• permet d'écrire des algorithmes ML optimisés et à usage général.
• propose une exécution autochtone sur les processeurs, les GPU et les TPU.
• prend en charge des fonctionnalités de pointe comme autograd et compilation jit .

Dans cet article, nous vous montrerons comment:

1. Écrivez des programmes JAX dans Python.
2. générer des spécifications XLA.
3. Déployer le code jax optimisé dans dart en utilisant c ffi .

? Qu'est-ce que Jax?

jax est comme Numpy sur les stéroïdes. Développé par Google, c'est une bibliothèque de bas niveau et haute performance qui rend ML accessible mais puissant.

• plate-forme agnostic : le code s'exécute sur des processeurs, des GPU et des TPU sans modification.
• Speed ​​: alimenté par le compilateur XLA, Jax optimise et accélère l'exécution.
• flexibilité : parfait pour les modèles ML et les algorithmes généraux.

Voici un exemple comparant Numpy et Jax:

# NumPy version
import numpy as np  
def assign_numpy():  
  a = np.empty(1000000)  
  a[:] = 1  
  return a  

# JAX version
import jax.numpy as jnp  
import jax  

@jax.jit  
def assign_jax():  
  a = jnp.empty(1000000)  
  return a.at[:].set(1)  

Benchmarking dans Google Colab révèle les performances de Jax:

• cpu & gpu : jax est plus rapide que numpy.
• tpu : les accélérations deviennent perceptibles pour les grands modèles en raison des coûts de transfert de données.

Cette flexibilité et cette vitesse rendent Jax idéal pour les environnements de production où les performances sont essentielles.

? ️ amener Jax en production

Microservices cloud par rapport au déploiement local

• cloud : Les microservices Python conteneurisés sont idéaux pour le calcul basé sur le cloud.
• local : l'expédition d'un interprète Python n'est pas idéal pour les applications locales.

Solution: Tirez parti de la compilation XLA de Jax

jax traduit le code python en hlo (Optimiseur de haut niveau) Spécifications , qui peuvent être compilées et exécutées à l'aide de C xla Libraries . Cela permet:

1. Écriture d'algorithmes dans Python.
2. les exécuter nativement via une bibliothèque C.
3. intégrant avec Dart via ffi (interface de fonction étrangère) .

Intégration étape par étape

1. Générer un proto HLO

Écrivez votre fonction JAX et exportez sa représentation HLO. Par exemple:

import jax.numpy as jnp  

def fn(x, y, z):  
  return jnp.dot(x, y) / z  

pour générer le HLO, utilisez le script jax_to_ir.py du référentiel jax:

python jax_to_ir.py \
  --fn jax_example.prog.fn \
  --input_shapes '[("x", "f32[2,2]"), ("y", "f32[2,2")]' \
  --constants '{"z": 2.0}' \
  --ir_format HLO \
  --ir_human_dest /tmp/fn_hlo.txt \
  --ir_dest /tmp/fn_hlo.pb

Placez les fichiers résultants (fn_hlo.txt et fn_hlo.pb) dans le répertoire des actifs de votre application.

2. Construisez une bibliothèque dynamique C

Modifier le code C. de Jax

Clone le référentiel JAX et accédez à jax / exemples / jax_cpp .

• Ajoutez un fichier d'en-tête main.h:

#ifndef MAIN_H  
#define MAIN_H  

extern "C" {  
  int bar(int foo);  
}  

#endif  

• Mettez à jour le fichier build pour créer une bibliothèque partagée:

cc_shared_library(  
   name = "jax",  
   deps = [":main"],  
   visibility = ["//visibility:public"],  
)  

Compiler avec bazel:

bazel build examples/jax_cpp:jax  

vous trouverez le libjax.dylib compilé dans le répertoire de sortie.

3. Connectez Dart avec C en utilisant FFI

Utilisez le package FFI de Dart pour communiquer avec la bibliothèque C. Créer un fichier jax.dart:

import 'dart:ffi';  
import 'package:dynamic_library/dynamic_library.dart';  

typedef FooCFunc = Int32 Function(Int32 bar);  
typedef FooDartFunc = int Function(int bar);  

class JAX {  
  late final DynamicLibrary dylib;  

  JAX() {  
    dylib = loadDynamicLibrary(libraryName: 'jax');  
  }  

  Function get _bar => dylib.lookupFunction('bar');  

  int bar(int foo) {  
    return _bar(foo);  
  }  
}  

Incluez la bibliothèque dynamique de votre répertoire de projet. Testez-le avec:

final jax = JAX();  
print(jax.bar(42));  

Vous verrez la sortie de la bibliothèque C dans votre console.

? Étapes suivantes

Avec cette configuration, vous pouvez:

• Optimiser les modèles ML avec JAX et XLA.
• Exécutez localement des algorithmes puissants.

Les cas d'utilisation potentiels incluent:

• Rechercher des algorithmes (par exemple, a *).
• optimisation combinatoire (par exemple, planification).
• Traitement d'images (par exemple, détection de bord).

Jax comble l'écart entre le développement basé sur Python et les performances au niveau de la production, permettant aux ingénieurs ML de se concentrer sur les algorithmes sans se soucier du code C de bas niveau.

Nous construisons une plate-forme d'IA de pointe avec des jetons de chat illimités et une mémoire à long terme, assurant des interactions transparentes et consacrées à un contexte qui évoluent avec le temps.

il est entièrement gratuit, et vous pouvez aussi l'essayer dans votre ide actuel.