Repolode
Enviar
jax-ml

jax-ml/jax-ai-stack

Um pacote instalável via pip que fixa um conjunto testado de bibliotecas JAX (Flax, Optax, Orbax, Chex, Grain) para IA e computação numérica.

Ver no GitHub
Estrelas
315
Forks
51
Issues abertas
38
Último commit
há 25 dias

O que é?

JAX AI Stack é um ponto de entrada único para instalar um conjunto de bibliotecas baseadas em JAX que funcionam em conjunto. Ao executar pip install jax-ai-stack, você instala JAX e projetos como Flax, para criar redes neurais; Optax, para processamento de gradientes e otimização; Orbax, para salvar e restaurar checkpoints; Chex, para escrever código JAX confiável; Grain, para carregamento de dados; e ml_dtypes, para extensões de tipos de dados do NumPy. As versões são fixadas em combinações verificadas pelos testes de integração do repositório.

JAX é uma biblioteca de computação com arrays que oferece transformações de programas como jit, vmap e grad. Ela inclui poucos recursos específicos para redes neurais, e o ecossistema ao redor é desenvolvido em pacotes separados. Esta stack reúne as bibliotecas que os desenvolvedores do Google usam para modelos de IA generativa, para que os usuários não precisem resolver manualmente quais versões são compatíveis.

Extras opcionais adicionam TensorFlow Datasets, e builds do JAX específicos para hardware, como GPUs CUDA ou TPUs, podem ser instalados no mesmo comando.

Recursos principais

  • 01Instala JAX com Flax, Optax, Orbax, Chex, Grain e ml_dtypes em versões testadas em conjunto.
  • 02Oferece transformações do JAX, como jit, vmap e grad.
  • 03Disponibiliza o extra opcional tensorflow-datasets via pip.
  • 04Oferece suporte a instalações do JAX específicas para hardware, como GPUs CUDA e TPUs.

Casos de uso

  • Configurar um ambiente JAX para treinar redes neurais com Flax e Optax.
  • Salvar e restaurar checkpoints de modelos com Orbax em projetos JAX.
  • Executar cargas de trabalho JAX em GPUs NVIDIA ou Google TPUs com um conjunto de dependências compatíveis e testadas.

Início rápido

  1. 1

    Instale a stack de IA do JAX pelo PyPI

    pip install jax-ai-stack
  2. 2

    Opcionalmente, adicione o extra tfds para habilitar o suporte ao TensorFlow Datasets

    pip install jax-ai-stack[tfds]
  3. 3

    Opcionalmente, instale com suporte a GPU, adicionando o extra JAX CUDA no mesmo comando

    pip install jax-ai-stack "jax[cuda]"

Antes de usar

  • Este é um conjunto de outras bibliotecas, portanto não oferece funcionalidades independentes além das fornecidas pelos pacotes fixados.
  • O README informa que a documentação e os tutoriais ainda estão em desenvolvimento.

Em coleções

Treine e execute IA em GPUs

Tópicos

#jax#machine-learning#deep-learning#python#numerical-computing#gpu#tpu

Perguntas frequentes

O que é jax-ai-stack?
Um pacote instalável via pip que fixa um conjunto testado de bibliotecas JAX (Flax, Optax, Orbax, Chex, Grain) para IA e computação numérica. JAX AI Stack é um ponto de entrada único para instalar um conjunto de bibliotecas baseadas em JAX que funcionam em conjunto. Ao executar pip install jax-ai-stack, você instala JAX e projetos como Flax, para criar redes neurais; Optax, para processamento de gradientes e otimização; Orbax, para salvar e restaurar checkpoints; Chex, para escrever código JAX confiável; Grain, para carregamento de dados; e ml_dtypes, para extensões de tipos de dados do NumPy. As versões são fixadas em combinações verificadas pelos testes de integração do repositório.
jax-ai-stack é gratuito? Qual é a licença?
jax-ai-stack é open source sob a licença Apache-2.0. Confira os termos da licença antes de usá-lo em produtos comerciais.
jax-ai-stack ainda é mantido?
Manutenção ativa. O último commit de jax-ai-stack foi em 13 de out. de 2026.
Para que posso usar jax-ai-stack?
Configurar um ambiente JAX para treinar redes neurais com Flax e Optax. Salvar e restaurar checkpoints de modelos com Orbax em projetos JAX. Executar cargas de trabalho JAX em GPUs NVIDIA ou Google TPUs com um conjunto de dependências compatíveis e testadas.
Como instalar jax-ai-stack?
Instale a stack de IA do JAX pelo PyPI: pip install jax-ai-stack Opcionalmente, adicione o extra tfds para habilitar o suporte ao TensorFlow Datasets: pip install jax-ai-stack[tfds] Opcionalmente, instale com suporte a GPU, adicionando o extra JAX CUDA no mesmo comando: pip install jax-ai-stack "jax[cuda]"
Onde consigo jax-ai-stack?
O código-fonte e as instruções de instalação estão no GitHub: https://github.com/jax-ml/jax-ai-stack

Resumo gerado por IA a partir do README do projeto. As estatísticas do GitHub são atualizadas diariamente.

Projetos relacionados