? Por que combinar Dart e Jax para aprendizado de máquina?

Ao criar aplicativos, a seleção das ferramentas certas é crucial. Você deseja um alto desempenho, desenvolvimento fácil e implantação contínua de plataforma cruzada. As estruturas populares oferecem trade-offs:

• c fornece velocidade, mas pode diminuir o desenvolvimento.
• Dart (com flutter) é mais lento, mas simplifica o gerenciamento da memória e o desenvolvimento de plataformas cruzadas.

Mas aqui está o problema: a maioria das estruturas não tem robusto nativo de aprendizado de máquina (ml) suporte. Essa lacuna existe porque essas estruturas são anteriores ao boom da IA. A pergunta é:

Como podemos integrar com eficiência ML em aplicativos?

soluções comuns como Onnx RunTime permitem exportar modelos ML para integração de aplicativos, mas não são otimizados para CPUs ou flexíveis o suficiente para algoritmos generalizados.

enter Jax , uma biblioteca Python que:

• permite escrever algoritmos ML otimizados e de uso geral.
• oferece execução de agnóstico da plataforma nas CPUs, GPUs e TPUs.
• suporta recursos de ponta como AutoGRAD e jit compilation .

Neste artigo, mostraremos como:

1. Escreva programas JAX em Python.
2. gerar especificações do XLA.
3. Implante o código Jax otimizado no DART usando c ffi .

? O que é Jax?

Jax é como Numpy em esteróides. Desenvolvido pelo Google, é uma biblioteca de baixo nível e alto desempenho que torna a ML acessível, mas poderosa.

• AGNOSTIC : o código é executado nas CPUs, GPUs e TPUs sem modificação.
• Speed ​​: alimentado pelo compilador XLA, Jax otimiza e acelera a execução.
• flexibilidade : Perfeito para modelos ML e algoritmos gerais.

aqui está um exemplo comparando Numpy e 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 no Google Colab revela a vantagem de desempenho de Jax:

• cpu & gpu : Jax é mais rápido que Numpy.
• TPU : Speed-UPs tornam-se perceptíveis para modelos grandes devido a custos de transferência de dados.

essa flexibilidade e velocidade tornam o JAX ideal para ambientes de produção onde o desempenho é chave.

? Aste trazer Jax para a produção

Microsserviços em nuvem vs. implantação local

• Cloud : Os microsserviços Python de contêiner são ótimos para computação baseada em nuvem.
• local : enviando um intérprete Python não é ideal para aplicativos locais.

Solução: alavancar a compilação XLA de Jax

Jax traduz o código Python em hlo (otimizador de alto nível) especificações , que podem ser compiladas e executadas usando C XLA Libraries . Isso permite:

1. escrevendo algoritmos em python.
2. executando -os nativamente através de uma biblioteca C.
3. integrando com o DART via ffi (interface da função estranha) .

✍️ Integração passo a passo

1. Gere um proto HLO

Escreva sua função JAX e exporte sua representação HLO. Por exemplo:

import jax.numpy as jnp  

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

Para gerar o HLO, use o script jax_to_ir.py do repositório 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

Coloque os arquivos resultantes (fn_hlo.txt e fn_hlo.pb) no diretório de ativos do seu aplicativo.

2. Construa uma biblioteca dinâmica C

Modifique o código de exemplo C de Jax

clone o repositório Jax e navegue para Jax/Exemplos/Jax_cpp .

• Adicione um arquivo de cabeçalho main.h:

#ifndef MAIN_H  
#define MAIN_H  

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

#endif  

• Atualize o arquivo build para criar uma biblioteca compartilhada:

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

compilar com Bazel:

bazel build examples/jax_cpp:jax  

você encontrará o libjax.dylib compilado no diretório de saída.

3. Conecte o DART com C usando FFI

use o pacote ffi do Dart para se comunicar com a biblioteca C. Crie um arquivo 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);  
  }  
}  

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);  
  }  
}