Paquete instalable con pip que fija un conjunto probado de bibliotecas de JAX (Flax, Optax, Orbax, Chex y Grain) para IA y computación numérica.
JAX AI Stack es un punto de entrada único para instalar un conjunto de bibliotecas basadas en JAX que funcionan en conjunto. Al ejecutar pip install jax-ai-stack se instala JAX junto con proyectos como Flax, para crear redes neuronales; Optax, para procesar gradientes y optimizar; Orbax, para gestionar checkpoints; Chex, para escribir código fiable con JAX; Grain, para cargar datos; y ml_dtypes, que amplía los tipos de datos de NumPy. Las versiones se fijan en combinaciones verificadas por las pruebas de integración del repositorio.
JAX es una biblioteca de computación con arrays que ofrece transformaciones de programas como jit, vmap y grad. Incluye pocas funciones específicas para redes neuronales, y el ecosistema que la rodea se desarrolla en paquetes separados. Este stack reúne las bibliotecas que usan los desarrolladores de Google para modelos de IA generativa, de modo que los usuarios no tengan que resolver manualmente qué versiones son compatibles.
Los extras opcionales añaden TensorFlow Datasets, y se pueden instalar en el mismo comando compilaciones de JAX específicas para el hardware, como las destinadas a GPU CUDA o TPU.
Instala la pila de IA de JAX desde PyPI
pip install jax-ai-stackOpcionalmente, añade el extra tfds para habilitar la compatibilidad con TensorFlow Datasets
pip install jax-ai-stack[tfds]Opcionalmente, instala con compatibilidad con GPU y añade el extra JAX CUDA en el mismo comando
pip install jax-ai-stack "jax[cuda]"Resumen generado por IA a partir del README del proyecto. Las estadísticas de GitHub se actualizan a diario.
Solvers numéricos de ecuaciones diferenciales para JAX, con soporte para ODE, SDE y CDE, autodiferenciación y ejecución en GPU.
Pax es un framework basado en JAX para configurar y ejecutar experimentos de machine learning a gran escala, principalmente con modelos de lenguaje, en TPU.
Backend de PyTorch que permite ejecutar modelos y operaciones de PyTorch en las unidades de procesamiento tensorial (TPU) de Google.
Keras 3 es un framework de deep learning que funciona con JAX, TensorFlow, PyTorch u OpenVINO para inferencia.
Una biblioteca de Python que convierte código de machine learning entre frameworks como PyTorch, TensorFlow, JAX y NumPy mediante ivy.transpile.
Una biblioteca de redes neuronales para JAX centrada en la flexibilidad, con la API Flax NNX para crear, inspeccionar y depurar modelos.