yax — des réseaux de neurones au plus près de jax
yax sert à écrire, entraîner et retrouver des réseaux de neurones avec jax,
sans rien apprendre d'autre que jax. Un modèle est une classe Python dont les
tableaux sont les paramètres : jax.grad, jax.jit, jax.vmap et optax
s'appliquent dessus directement. Une fonction, yax.training.train, fait la
boucle d'entraînement et garde le meilleur modèle sur le disque.
pip install yaxlib
En trente secondes
Une sinusoïde bruitée, un perceptron multicouche, un entraînement :
import jax.numpy as jnp
import jax.random as jr
import yax
X = jr.uniform(jr.key(0), (256, 1), minval=-3.0, maxval=3.0)
Y = jnp.sin(X) + 0.1 * jr.normal(jr.key(1), X.shape)
model = yax.layers.MLP(layer_sizes=(1, 32, 32, 1), activation="tanh", rkey=jr.key(2))
config = yax.configs.TrainConfig(learning_rate=1e-2, n_epoch=100)
run = yax.training.train(model, "out/sinus", config,
(X[:200], Y[:200], 32), # entraînement, par batchs de 32
(X[200:], Y[200:]), # validation
rkey=jr.key(3))
print(run.loss) # perte de validation du meilleur modèle
Y_pred = yax.losses.batch_apply(run.trained_model, X, None)
Le dossier out/sinus/0 contient le meilleur modèle, sa configuration et
l'historique des pertes : yax.training.load_run("out/sinus/0") les relit
dans une autre session.
Ce que yax vous donne
- Un modèle est un pytree. Déclarez une classe qui hérite de
yax.Module, ses champs sont les paramètres, sa méthodeapply(x, rkey=None)calcule la sortie pour un exemple. Rien à envelopper, rien à filtrer :jax.grad(loss)(model)rend un modèle de même forme dont les feuilles sont les gradients, etoptimizer.init(model)l'accepte tel quel. - Une seule signature, partout.
applytraite un exemple,jax.vmapfait le batch ;rkeyest la source d'aléatoire quand il y en a (dropout, tirages), et le mode d'évaluation se bascule d'un mot,model.set_inference(True). - Un entraînement qui laisse des traces. Chaque appel à
yax.training.traincrée un dossier numéroté avec le meilleur modèle, l'état de l'optimiseur, la configuration et l'historique. Les runs d'un même dossier se comparent, etyax.training.find_best_rundésigne le meilleur. - Des briques prêtes à l'emploi, toutes écrites comme vous les écririez : couches (MLP, convolutions n-dimensionnelles, GRU et LSTM, attention multi-têtes, blocs transformer, message passing), modèles complets (U-Net, opérateur de Fourier, VAE, flot normalisant, diffusion, mini-YOLO), pertes, optimiseurs d'optax par leur nom, prétraitements de données.
- Des mini-fonctionnalités sympathiques. Essayez par exemple
yax.ipprint(model)dans un notebook.
Pourquoi pas flax ou equinox ?
Par simplicité. Chaque framework se mesure aux concepts qu'il ajoute à jax : flax apporte ses collections de variables, ses scopes et son cycle init/apply ; equinox réduit cela à des modules-pytrees, mais y ajoute sa machinerie de filtrage. yax ajoute le strict minimum pour construire des modèles à base de layers imbriqués. Tout ce qui n'est pas dans yax se code en jax ordinaire, sans friction.
Pour continuer
- Prise en main : un quart d'heure, des données, un modèle, un entraînement, un rechargement.
- Aller plus loin : sous le capot du modèle, l'aléa et le mode inférence, les samplers, la configuration de l'entraînement, les pertes à soi.
- Référence de l'API et démos : une démonstration par famille de modèles, qui converge en quelques secondes sur CPU.
Développement
pip install yaxlib[dev] # ajoute pytest
pytest tests/
1 maintainer
Metadata
1 maintainer
Release files for yaxlib 0.2.1
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | Provenance | |
|---|---|---|---|---|
| yaxlib-0.2.1.tar.gz | 214.5 kB | Without provenance | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| yaxlib-0.2.1-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 282.1 kB
Release files / yaxlib-0.2.1.tar.gz
| Download URL | yaxlib-0.2.1.tar.gz |
|---|---|
| Size | 214.5 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
95c489c772dbefe1d2f27ba3b399bb3c60d272fe37a95375997dd182db8138a0
|
|
BLAKE2b-256 checksum How to use checksums |
d14405f71f0084aec7bdb018f8f7f9070ba8d3977e82b8c57829e7d209b4337d
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
uv/0.12.13 {"installer":{"name":"uv","version":"0.12.13","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"macOS","version":null,"id":null,"libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":null}
|
Release files / yaxlib-0.2.1-py3-none-any.whl
| Download URL | yaxlib-0.2.1-py3-none-any.whl |
|---|---|
| Size | 67.6 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
78a4ca114371634212f7e8fc5e1b6f26d69b0986d51930f5861b50bf4f3d2a65
|
|
BLAKE2b-256 checksum How to use checksums |
b4bd6288cae7cae8ecd05c0aad6871a22286c1d49b1c474b89f2bb1ebafddc00
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
uv/0.12.13 {"installer":{"name":"uv","version":"0.12.13","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"macOS","version":null,"id":null,"libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":null}
|
Security for yaxlib 0.2.1
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
No provenance found
None of the files in this release were uploaded with provenance metadata.
A note on package security
PyPI is a repository for distributing packages, not a service for auditing them. While third-party security researchers and our staff do their best to detect and remove malicious packages, PyPI cannot guarantee the quality, security or suitability of any given package.
To learn more about protecting yourself from malicious actors, visit our package security documentation.
