Repolode
Enviar
google

google/flax

Una biblioteca de redes neuronales para JAX centrada en la flexibilidad, con la API Flax NNX para crear, inspeccionar y depurar modelos.

Ver en GitHubSitio web
Estrellas
7,366
Forks
849
Issues abiertos
526
Último commit
anteayer

¿Qué es?

Flax es una biblioteca y un ecosistema de redes neuronales basados en JAX. Su nueva API Flax NNX permite definir modelos con objetos Python convencionales que admiten compartir referencias y mutabilidad, lo que facilita crear, inspeccionar, depurar y analizar redes. La API anterior, Flax Linen, se publicó en 2020 y es la predecesora de la que evolucionó NNX.

La biblioteca ofrece capas habituales como Linear, Conv, BatchNorm, LayerNorm, GroupNorm, atención multi-cabeza, celdas LSTM y GRU, y Dropout. También incluye utilidades para el entrenamiento replicado, la serialización y la gestión de checkpoints, las métricas y la precarga de datos en dispositivos. Los ejemplos abarcan MNIST y la inferencia con el modelo de lenguaje Gemma.

Flax está diseñado para que los usuarios modifiquen directamente el bucle de entrenamiento, en lugar de depender de funciones añadidas al framework. Se puede instalar desde PyPI y requiere Python 3.8 o posterior.

Funciones clave

  • 01API Flax NNX que utiliza objetos Python convencionales con referencias compartidas y mutabilidad
  • 02Capas integradas como Linear, Conv, capas de normalización, atención, celdas recurrentes y Dropout
  • 03Utilidades para el entrenamiento replicado, la gestión de checkpoints, las métricas y la precarga de datos en dispositivos
  • 04Ejemplos de MNIST e inferencia con el modelo de lenguaje Gemma
  • 05Instalable desde PyPI sobre JAX

Casos de uso

  • Investigar y crear prototipos de nuevas arquitecturas de redes neuronales en JAX
  • Escribir bucles de entrenamiento personalizados para experimentos sin las limitaciones de un framework
  • Entrenar modelos en CPU, GPU o TPU mediante JAX
  • Aprender deep learning con JAX mediante los ejemplos de MNIST y Gemma

Inicio rápido

  1. 1

    Primero, instala JAX siguiendo las instrucciones oficiales de instalación de JAX para tu hardware

  2. 2

    Instala Flax desde PyPI

    pip install flax

Antes de usarlo

  • Flax depende de JAX, así que también debes instalar y configurar JAX para tu hardware.

Temas

#jax#neural-network#deep-learning#machine-learning#python#research

Preguntas frecuentes

¿Qué es flax?
Una biblioteca de redes neuronales para JAX centrada en la flexibilidad, con la API Flax NNX para crear, inspeccionar y depurar modelos. Flax es una biblioteca y un ecosistema de redes neuronales basados en JAX. Su nueva API Flax NNX permite definir modelos con objetos Python convencionales que admiten compartir referencias y mutabilidad, lo que facilita crear, inspeccionar, depurar y analizar redes. La API anterior, Flax Linen, se publicó en 2020 y es la predecesora de la que evolucionó NNX.
¿flax es gratuito? ¿Qué licencia tiene?
flax es open source bajo la licencia Apache-2.0. Revisa los términos de la licencia antes de usarlo en productos comerciales.
¿flax sigue mantenido?
Mantenimiento activo. El último commit de flax fue el 5 nov 2026.
¿Para qué puedo usar flax?
Investigar y crear prototipos de nuevas arquitecturas de redes neuronales en JAX Escribir bucles de entrenamiento personalizados para experimentos sin las limitaciones de un framework Entrenar modelos en CPU, GPU o TPU mediante JAX Aprender deep learning con JAX mediante los ejemplos de MNIST y Gemma
¿Cómo se instala flax?
Primero, instala JAX siguiendo las instrucciones oficiales de instalación de JAX para tu hardware Instala Flax desde PyPI: pip install flax
¿Dónde consigo flax?
El código fuente y las instrucciones de instalación están en GitHub: https://github.com/google/flax

Resumen generado por IA a partir del README del proyecto. Las estadísticas de GitHub se actualizan a diario.

Proyectos relacionados