yaxlib 0.1.0__tar.gz

yaxlib-0.1.0/.gitignore

@@ -0,0 +1,77 @@
+# ===== Environnements virtuels =====
+.venv/
+venv/
+env/
+ENV/
+.conda/
+.python-version
+
+# ===== Artefacts de build / cache Python =====
+__pycache__/
+*.py[cod]
+*$py.class
+build/
+dist/
+*.egg
+*.egg-info/
+.eggs/
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+.ipynb_checkpoints/
+
+# ===== Node / JS (au cas où) =====
+node_modules/
+
+# ===== Éditeurs & assistants IA =====
+.idea/
+.vscode/
+.claude/
+.gemini/
+.agent/
+.agents/
+*.swp
+*~
+.DS_Store
+
+# ===== Archives =====
+*.zip
+
+# ===== Médias lourds — sauvegardés par Drive, hors git =====
+*.mp4
+*.mov
+*.webm
+*.mkv
+*.avi
+*.mp3
+*.wav
+*.flac
+*.m4a
+*.ogg
+# (images gardées dans git)
+
+# ===== Données scientifiques lourdes — sauvegardées par Drive =====
+*.npy
+*.npz
+# Tables EOS supernova (données CompOSE; attention: pas 'eos.*' qui
+# ignorerait aussi les eos.py du code)
+**/eos.thermo
+**/eos.compo
+**/eos.micro
+**/eos.t
+**/eos.nb
+**/eos.yq
+# (tout autre fichier > 50 Mo est bloqué par le hook pre-commit)
+
+# ===== Fichiers-pointeurs Google Drive (liens, pas du vrai contenu) =====
+*.gdoc
+*.gsheet
+*.gslides
+*.gdraw
+*.gform
+*.gmap
+*.gsite
+*.gtable
+*.glink
+*.gjam
+*.gnote

yaxlib-0.1.0/PKG-INFO

@@ -0,0 +1,72 @@
+Metadata-Version: 2.5
+Name: yaxlib
+Version: 0.1.0
+Summary: Un mini-framework de réseaux de neurones pour jax, à but pédagogique (import : yax)
+Author: Vincent Vigon
+License: MIT
+Requires-Python: >=3.10
+Requires-Dist: jax>=0.4.30
+Requires-Dist: numpy
+Requires-Dist: optax>=0.2
+Provides-Extra: dev
+Requires-Dist: equinox; extra == 'dev'
+Requires-Dist: matplotlib; extra == 'dev'
+Requires-Dist: pytest; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# yaxlib
+
+Un mini-framework de réseaux de neurones pour jax, à but pédagogique.
+Distribution `yaxlib`, import `yax` :
+
+```bash
+pip install yaxlib          # ou : pip install <url du wheel>
+```
+
+```python
+import jax.random as jr
+import yax
+
+model = yax.MLP((2, 32, 32, 1), "tanh", jr.key(0))
+```
+
+## Principes
+
+- **Un modèle est un pytree.** `yax.Module` range les tableaux dans les
+  feuilles et tout le reste (`yax.StaticField`) dans la structure :
+  `jax.grad(loss)(model)`, `jax.jit` et `optimizer.init(model)` acceptent le
+  modèle tel quel, sans machinerie de filtrage. Les champs dynamiques ne
+  peuvent contenir que des tableaux jax, des sous-modules ou des conteneurs de
+  ceux-ci — tout écart est une erreur immédiate et explicite à la
+  construction. Un `StaticField` peut contenir un tableau : il devient une
+  constante du modèle (encodage positionnel, grille figée), invisible pour les
+  gradients.
+- **Signature uniforme `apply(x, rkey=None)`**, écrite pour UN échantillon
+  (le batch vient de `jax.vmap`). `rkey` est une *source d'aléatoire*
+  (dropout, échantillonnage), jamais un mode.
+- **Le mode se bascule par `model = model.set_inference(True/False)`**
+  (récursif, immuable). En pratique on n'y touche pas : le `Trainer` entraîne
+  en `False`, valide et rend le meilleur modèle en `True`.
+- **Immutabilité** : on « modifie » un module avec `yax.tree_at`.
+
+## Contenu
+
+- `yax.layers` : Linear, MLP, Dropout, LayerNorm, Embedding, Conv_layer,
+  RNN_layer (GRU/LSTM), MultiHeadAttention, TransformerBlock,
+  MessagePassing_layer, encodage positionnel.
+- `yax.models` : UNet, MiniYOLO (références des mini-projets).
+- `yax.training` : Trainer (checkpoints par `mother_folder`), History, pertes
+  (`loss_fn(model, x, y, rkey)`).
+- `yax.image` : augmentation différentiable et vmap-able.
+- `demos/` : une démonstration synthétique par famille de modèles, qui
+  converge en quelques secondes sur CPU.
+
+## Tests
+
+```bash
+pip install yaxlib[dev]
+pytest tests/
+```
+
+Les couches réimplémentées (GRU, LSTM, convolution) sont vérifiées
+numériquement contre `equinox.nn`, qui ne sert qu'à cela.

yaxlib-0.1.0/README.md

@@ -0,0 +1,56 @@
+# yaxlib
+
+Un mini-framework de réseaux de neurones pour jax, à but pédagogique.
+Distribution `yaxlib`, import `yax` :
+
+```bash
+pip install yaxlib          # ou : pip install <url du wheel>
+```
+
+```python
+import jax.random as jr
+import yax
+
+model = yax.MLP((2, 32, 32, 1), "tanh", jr.key(0))
+```
+
+## Principes
+
+- **Un modèle est un pytree.** `yax.Module` range les tableaux dans les
+  feuilles et tout le reste (`yax.StaticField`) dans la structure :
+  `jax.grad(loss)(model)`, `jax.jit` et `optimizer.init(model)` acceptent le
+  modèle tel quel, sans machinerie de filtrage. Les champs dynamiques ne
+  peuvent contenir que des tableaux jax, des sous-modules ou des conteneurs de
+  ceux-ci — tout écart est une erreur immédiate et explicite à la
+  construction. Un `StaticField` peut contenir un tableau : il devient une
+  constante du modèle (encodage positionnel, grille figée), invisible pour les
+  gradients.
+- **Signature uniforme `apply(x, rkey=None)`**, écrite pour UN échantillon
+  (le batch vient de `jax.vmap`). `rkey` est une *source d'aléatoire*
+  (dropout, échantillonnage), jamais un mode.
+- **Le mode se bascule par `model = model.set_inference(True/False)`**
+  (récursif, immuable). En pratique on n'y touche pas : le `Trainer` entraîne
+  en `False`, valide et rend le meilleur modèle en `True`.
+- **Immutabilité** : on « modifie » un module avec `yax.tree_at`.
+
+## Contenu
+
+- `yax.layers` : Linear, MLP, Dropout, LayerNorm, Embedding, Conv_layer,
+  RNN_layer (GRU/LSTM), MultiHeadAttention, TransformerBlock,
+  MessagePassing_layer, encodage positionnel.
+- `yax.models` : UNet, MiniYOLO (références des mini-projets).
+- `yax.training` : Trainer (checkpoints par `mother_folder`), History, pertes
+  (`loss_fn(model, x, y, rkey)`).
+- `yax.image` : augmentation différentiable et vmap-able.
+- `demos/` : une démonstration synthétique par famille de modèles, qui
+  converge en quelques secondes sur CPU.
+
+## Tests
+
+```bash
+pip install yaxlib[dev]
+pytest tests/
+```
+
+Les couches réimplémentées (GRU, LSTM, convolution) sont vérifiées
+numériquement contre `equinox.nn`, qui ne sert qu'à cela.

yaxlib-0.1.0/demos/cnn_demo_classif.py

@@ -0,0 +1,123 @@
+"""Classification d'images synthétiques par un petit CNN.
+
+Trois classes de formes — disque, carré, croix — posées à position et taille
+aléatoires sur un fond bruité. Le jeu se génère en une ligne de vmap, converge
+en quelques secondes, et l'invariance par translation de la convolution y est
+réellement mise à l'épreuve (la position varie d'une image à l'autre).
+
+Le sous-échantillonnage est fait par stride=2 (l'alternative au pooling,
+chapitre 11) : 16x16 -> 8x8 -> 4x4, puis tête dense sur l'aplati.
+"""
+
+import os
+
+from yax.core import Module
+import jax
+import jax.numpy as jnp
+import jax.random as jr
+
+from yax.layers.Conv_layer import Conv_layer
+from yax.layers.Linear import Linear
+from yax.training.Trainer import train, load_run, OUT_FOLDER
+from yax.training.configs import TrainConfig
+from yax.training.losses import softmax_ce_loss, batch_apply
+
+IMG_SIZE = 16
+NB_CLASSES = 3  # 0: disque, 1: carré, 2: croix
+
+
+def make_image(rkey, label):
+    """UNE image (1, 16, 16) : une forme `label` sur fond bruité.
+
+    Les rayons restent assez grands (4.5 à 5.5) pour qu'un disque et un carré
+    diffèrent de plusieurs pixels aux coins : avec des formes plus petites,
+    seuls les coins les distinguent, la tâche devient réellement ambiguë sous
+    le bruit et la confusion disque/carré plafonne la justesse."""
+    k_pos, k_radius, k_noise = jr.split(rkey, 3)
+    center = jr.uniform(k_pos, (2,), minval=6.0, maxval=IMG_SIZE - 6.0)
+    radius = jr.uniform(k_radius, minval=4.5, maxval=5.5)
+
+    rows = jnp.arange(IMG_SIZE)[:, None]
+    cols = jnp.arange(IMG_SIZE)[None, :]
+    dr, dc = rows - center[0], cols - center[1]
+
+    disque = dr**2 + dc**2 < radius**2
+    carre = jnp.maximum(jnp.abs(dr), jnp.abs(dc)) < radius
+    croix = ((jnp.abs(dr) < 1.0) | (jnp.abs(dc) < 1.0)) & (jnp.abs(dr) + jnp.abs(dc) < 2 * radius)
+
+    shape = jnp.stack([disque, carre, croix])[label]
+    img = jnp.where(shape, 1.0, 0.0) + 0.25 * jr.normal(k_noise, (IMG_SIZE, IMG_SIZE))
+    return img[None, :, :]  # channels-first : (1, H, W)
+
+
+def make_data(rkey, nb):
+    k_labels, k_images = jr.split(rkey)
+    labels = jr.randint(k_labels, (nb,), 0, NB_CLASSES)
+    X = jax.vmap(make_image)(jr.split(k_images, nb), labels)
+    return X, labels
+
+
+def present_data():
+    from matplotlib import pyplot as plt
+
+    X, Y = make_data(jr.key(0), 12)
+    fig, axs = plt.subplots(3, 4, figsize=(8, 6))
+    for ax, img, label in zip(axs.flat, X, Y):
+        ax.imshow(img[0], cmap="gray")
+        ax.set_title(["disque", "carré", "croix"][int(label)])
+        ax.axis("off")
+    plt.tight_layout()
+    plt.show()
+
+
+class CNN_classifier(Module):
+    conv1: Conv_layer
+    conv2: Conv_layer
+    head: Linear
+
+    def __init__(self, rkey):
+        k1, k2, k3 = jr.split(rkey, 3)
+        self.conv1 = Conv_layer(1, 8, 3, k1, stride=2)    # (1,16,16) -> (8,8,8)
+        self.conv2 = Conv_layer(8, 16, 3, k2, stride=2)   # -> (16,4,4)
+        self.head = Linear(16 * 4 * 4, NB_CLASSES, k3)
+
+    def apply(self, x, rkey=None):
+        x = jax.nn.relu(self.conv1.apply(x))
+        x = jax.nn.relu(self.conv2.apply(x))
+        return self.head.apply(x.reshape(-1))  # logits (NB_CLASSES,)
+
+
+def accuracy(model, x, y):
+    logits = batch_apply(model, x, None)
+    return float(jnp.mean(jnp.argmax(logits, axis=-1) == y))
+
+
+def train_classif(verbose=True):
+    key_train, key_val, key_model, key_shuffle = jr.split(jr.key(0), 4)
+    X_train, Y_train = make_data(key_train, 1024)
+    X_val, Y_val = make_data(key_val, 512)
+
+    model = CNN_classifier(key_model)
+    acc_before = accuracy(model, X_val, Y_val)
+
+    config = TrainConfig(batch_size=64, learning_rate=0.003, n_epoch=100,
+                         lr_final_ratio=0.01)
+    mother_folder = os.path.join(OUT_FOLDER, "cnn_classif")
+    best_model, best_loss, folder = train(
+        model, mother_folder, config, X_train, Y_train, X_val, Y_val,
+        rkey=key_shuffle, loss_fn=softmax_ce_loss, verbose=verbose)
+
+    acc_after = accuracy(best_model, X_val, Y_val)
+    print(f"\naccuracy {acc_before:.3f} -> {acc_after:.3f}   (best val CE = {best_loss:.3g})")
+    return folder, best_loss, acc_after
+
+
+if __name__ == "__main__":
+    from matplotlib import pyplot as plt
+
+    present_data()
+    folder, best_loss, acc = train_classif()
+    history = load_run(folder, "history")["history"]
+    history.plot(title=f"CNN formes : CE={best_loss:.2g}, accuracy={acc:.3f}")
+    plt.tight_layout()
+    plt.show()

yaxlib-0.1.0/demos/gnn_demo_nodes.py

@@ -0,0 +1,134 @@
+"""Classification de nœuds sur UN graphe fixe, par message passing.
+
+Un graphe à deux communautés (modèle à blocs stochastiques) : les liens sont
+denses à l'intérieur d'une communauté, rares entre les deux. Chaque nœud porte
+une caractéristique très bruitée de sa communauté — si bruitée qu'un nœud seul
+ne dit presque rien : l'optimum théorique nœud par nœud est ~69 %, et le MLP,
+qui sur-apprend ses nœuds d'entraînement dans les 7 dimensions de bruit pur,
+fait pire encore. Le message passing moyenne l'information sur le voisinage,
+qui est surtout de la même communauté : le bruit s'écrase, le GNN dépasse
+90 %. C'est l'argument du chapitre en une expérience.
+
+La moitié des nœuds est étiquetée, la justesse est mesurée sur les autres.
+Boucle d'entraînement full-batch écrite à la main (comme mlp_demo) : un seul
+graphe, pas de mini-batchs — le Trainer n'a rien à mélanger ici.
+"""
+
+from yax.core import Module
+import jax
+import jax.numpy as jnp
+import jax.random as jr
+import numpy as np
+import optax
+
+from yax.layers.Linear import Linear
+from yax.layers.MLP import MLP
+from yax.layers.MessagePassing_layer import MessagePassing_layer
+
+NB_PER_BLOCK = 60
+DIM_FEATURES = 8
+SIGMA = 2.0          # écart-type du bruit sur les caractéristiques
+NB_LABELED = 30      # nœuds étiquetés par communauté (la moitié)
+
+
+def make_graph(seed=0):
+    """Rend (features, labels, senders, receivers, train_idx, test_idx).
+    Génération en numpy : les arêtes sont des tableaux concrets, construits une
+    fois pour toutes hors de tout jit — le graphe est FIXE."""
+    rng = np.random.default_rng(seed)
+    nb_nodes = 2 * NB_PER_BLOCK
+    labels = np.repeat([0, 1], NB_PER_BLOCK)
+
+    # liens : p_in dans une communauté, p_out entre communautés
+    p = np.where(labels[:, None] == labels[None, :], 0.2, 0.02)
+    upper = np.triu(rng.random((nb_nodes, nb_nodes)) < p, k=1)
+    i, j = np.nonzero(upper)
+    senders = np.concatenate([i, j])      # un lien non orienté = deux arêtes
+    receivers = np.concatenate([j, i])
+
+    # caractéristiques : +/-1 sur la première coordonnée, noyé dans le bruit
+    mu = np.zeros((2, DIM_FEATURES))
+    mu[0, 0], mu[1, 0] = 1.0, -1.0
+    features = mu[labels] + SIGMA * rng.normal(size=(nb_nodes, DIM_FEATURES))
+
+    labeled = np.concatenate([rng.choice(np.arange(b * NB_PER_BLOCK, (b + 1) * NB_PER_BLOCK),
+                                         NB_LABELED, replace=False) for b in range(2)])
+    test_idx = np.setdiff1d(np.arange(nb_nodes), labeled)
+
+    to_j = lambda a, dtype: jnp.asarray(a, dtype=dtype)
+    return (to_j(features, jnp.float32), to_j(labels, jnp.int32),
+            to_j(senders, jnp.int32), to_j(receivers, jnp.int32),
+            to_j(labeled, jnp.int32), to_j(test_idx, jnp.int32))
+
+
+class GNN_node_classifier(Module):
+    mp1: MessagePassing_layer
+    mp2: MessagePassing_layer
+    head: Linear
+
+    def __init__(self, rkey):
+        k1, k2, k3 = jr.split(rkey, 3)
+        # agrégation "mean" : l'échelle de l'agrégat ne dépend pas du degré,
+        # qui varie beaucoup d'un nœud à l'autre dans ce graphe
+        self.mp1 = MessagePassing_layer(DIM_FEATURES, 16, 16, k1, aggregation="mean")
+        self.mp2 = MessagePassing_layer(16, 16, 16, k2, aggregation="mean")
+        self.head = Linear(16, 2, k3)
+
+    def apply(self, h, rkey=None, *, senders, receivers):
+        h = self.mp1.apply(h, senders=senders, receivers=receivers)
+        h = self.mp2.apply(h, senders=senders, receivers=receivers)
+        return self.head.apply(h)  # logits (nb_noeuds, 2)
+
+
+def train_nodes(logits_fn, model, labels, train_idx, nb_steps=300, lr=0.01):
+    """logits_fn(model) -> (nb_noeuds, 2). Perte : CE sur les seuls nœuds
+    étiquetés — l'indexation par entiers reste valide sous jit."""
+    optimizer = optax.adam(lr)
+    opt_state = optimizer.init(model)
+
+    def loss_fn(model):
+        logits = logits_fn(model)
+        ce = optax.softmax_cross_entropy_with_integer_labels(
+            logits[train_idx], labels[train_idx])
+        return jnp.mean(ce)
+
+    @jax.jit
+    def step(model, opt_state):
+        loss, grads = jax.value_and_grad(loss_fn)(model)
+        updates, opt_state = optimizer.update(grads, opt_state)
+        return optax.apply_updates(model, updates), opt_state, loss
+
+    for _ in range(nb_steps):
+        model, opt_state, loss = step(model, opt_state)
+    return model, float(loss)
+
+
+def accuracy_on(logits_fn, model, labels, idx):
+    pred = jnp.argmax(logits_fn(model)[idx], axis=-1)
+    return float(jnp.mean(pred == labels[idx]))
+
+
+def run_comparison(seed=0, verbose=True):
+    X, labels, senders, receivers, train_idx, test_idx = make_graph(seed)
+    k_gnn, k_mlp = jr.split(jr.key(seed))
+
+    # référence : un MLP nœud par nœud, aveugle au graphe
+    mlp = MLP((DIM_FEATURES, 32, 2), "relu", k_mlp)
+    mlp_logits = lambda m: m.apply(X)
+    mlp, _ = train_nodes(mlp_logits, mlp, labels, train_idx)
+    acc_mlp = accuracy_on(mlp_logits, mlp, labels, test_idx)
+
+    # le GNN : mêmes caractéristiques, plus la structure du graphe
+    gnn = GNN_node_classifier(k_gnn)
+    gnn_logits = lambda m: m.apply(X, senders=senders, receivers=receivers)
+    gnn, _ = train_nodes(gnn_logits, gnn, labels, train_idx)
+    acc_gnn = accuracy_on(gnn_logits, gnn, labels, test_idx)
+
+    if verbose:
+        print(f"MLP (sans le graphe) : {acc_mlp:.3f}")
+        print(f"GNN (2 couches de message passing) : {acc_gnn:.3f}")
+    return acc_mlp, acc_gnn
+
+
+if __name__ == "__main__":
+    run_comparison()

yaxlib-0.1.0/demos/mlp_demo.py

@@ -0,0 +1,82 @@
+"""Approximation de sinus par un MLP.
+
+Descente de gradient a la main : pas d'optax, pas de filter_grad, pas de
+partition/combine. Tout repose sur le fait qu'un yax.Module est un pytree
+dont les feuilles sont exactement les parametres, a condition d'avoir
+declare en `static` tout ce qui n'est pas un tableau.
+"""
+
+import jax
+import jax.numpy as jnp
+import jax.random as jr
+import matplotlib.pyplot as plt
+from yax.layers.MLP import MLP
+
+
+def make_data(rkey, nb, x_min, x_max):
+    x = jr.uniform(rkey, (nb, 1), minval=x_min, maxval=x_max)
+    y = jnp.sin(x)
+    return x, y
+
+
+def loss_fn(model, x, y):
+    y_pred = jax.vmap(model.apply)(x)
+    return jnp.mean((y_pred - y) ** 2)
+
+
+@jax.jit
+def step(model, x, y, lr):
+    loss, grads = jax.value_and_grad(loss_fn)(model, x, y)
+    new_model = jax.tree.map(lambda p, g: p - lr * g, model, grads)
+    return new_model, loss
+
+
+def train(model, x, y, lr, nb_steps):
+    losses = []
+    for i in range(nb_steps):
+        model, loss = step(model, x, y, lr)
+        losses.append(loss)
+        if i % (nb_steps // 10) == 0:
+            print(f"step {i:5d}   loss = {loss:.3e}")
+    return model, jnp.array(losses)
+
+
+def plot(model, x_train, y_train, losses, x_min, x_max):
+    x_plot = jnp.linspace(x_min, x_max, 400)[:, None]
+    y_plot = jax.vmap(model.apply)(x_plot)
+    fig, (ax0, ax1) = plt.subplots(1, 2, figsize=(11, 4))
+
+    ax0.plot(x_train[:, 0], y_train[:, 0], ".", ms=4, alpha=0.35,
+             color="tab:gray", label="donnees")
+    ax0.plot(x_plot[:, 0], jnp.sin(x_plot[:, 0]), "--", lw=2,
+             color="tab:blue", label="sin(x)")
+    ax0.plot(x_plot[:, 0], y_plot[:, 0], lw=2,
+             color="tab:red", label="MLP")
+    ax0.set_xlabel("x")
+    ax0.set_title("approximation")
+    ax0.legend()
+
+    ax1.semilogy(losses, lw=1, color="tab:red")
+    ax1.set_xlabel("step")
+    ax1.set_title("loss (MSE)")
+    ax1.grid(alpha=0.3)
+
+    fig.tight_layout()
+    return fig
+
+
+if __name__ == "__main__":
+    x_min, x_max = -3.0, 3.0
+    lr = 0.02
+    nb_steps = 10000
+    key_data, key_model = jr.split(jr.PRNGKey(0))
+    x, y = make_data(key_data, 256, x_min, x_max)
+    model = MLP((1, 32, 32, 1), "tanh", key_model)
+
+    print("loss initiale :", loss_fn(model, x, y))
+    model, losses = train(model, x, y, lr, nb_steps)
+    print("loss finale   :", losses[-1])
+
+    fig = plot(model, x, y, losses, x_min, x_max)
+    fig.savefig("mpl_demo.png", dpi=120)
+    plt.show()

yaxlib-0.1.0/demos/rnn_demo_classif.py

@@ -0,0 +1,110 @@
+import os
+import jax.numpy as jnp
+import jax.random as jr
+from matplotlib import pyplot as plt
+from yax.core import Module, StaticField
+import jax
+import optax
+from yax.layers.Linear import Linear
+from yax.layers.RNN_layer import RNN_layer
+from yax.training.Trainer import train, load_run, OUT_FOLDER
+from yax.training.configs import TrainConfig
+from yax.training.losses import bce_logits_loss, batch_apply
+
+
+#C'est l'exemple qui est proposé dans la doc d'equinox
+def get_data(dataset_size:int, rkey):
+    t = jnp.linspace(0, 2 * jnp.pi, 16)
+    offset = jr.uniform(rkey, (dataset_size, 1), minval=0, maxval=2 * jnp.pi)
+    x1 = jnp.sin(t + offset) / (1 + t)
+    x2 = jnp.cos(t + offset) / (1 + t)
+
+    half_dataset_size = dataset_size // 2
+    x1 = x1.at[:half_dataset_size].multiply(-1)
+    y = jnp.ones((dataset_size, 1))
+    y = y.at[:half_dataset_size].set(0)
+    #axis=-1 : on veut (dataset_size, 16, 2), soit 16 pas de temps de 2 features.
+    #Avec axis=1 on obtiendrait (dataset_size, 2, 16), et le scan parcourrait
+    #2 pas de temps de 16 features.
+    x = jnp.stack([x1, x2], axis=-1)
+
+    #les 2 classes sont rangees en bloc, mais batchs_for_one_epoch melange a
+    #chaque epoque, et train/val proviennent de deux appels distincts, chacun
+    #equilibre. Ne PAS decouper un seul get_data en tranches train/val.
+    return x, y
+
+
+def present_data():
+    n_data = 10
+    x, y = get_data(n_data,jr.key(0))
+    print(x.shape, y.shape)
+
+    fig, axs = plt.subplots(n_data, 2, figsize=(8, 1.2 * n_data), sharex="all", sharey="all")
+    for i in range(n_data):
+        axs[i, 0].plot(x[i, :, 0])
+        axs[i, 1].plot(x[i, :, 1])
+        axs[i, 0].set_ylabel(f"label={int(y[i, 0])}")
+
+    axs[0, 0].set_title("feature 0")
+    axs[0, 1].set_title("feature 1")
+    plt.show()
+
+
+class RNN_bin_classifier(Module):
+    rnn_layer: RNN_layer
+    linear_layer: Linear
+
+    dim_hidden:int = StaticField()
+
+    def __init__(self,dim_hidden,rkey,*,cell_type="gru"):
+        self.dim_hidden=dim_hidden
+
+        rk1,rk2=jr.split(rkey)
+        self.rnn_layer = RNN_layer(2,dim_hidden,False,rk1,cell_type=cell_type)
+        self.linear_layer=Linear(dim_hidden,1,rk2)
+
+    def apply(self,xs,rkey=None):
+        # rkey ignoré : modèle déterministe (signature uniforme apply(x, rkey))
+        y=self.rnn_layer.apply(xs)
+        y=self.linear_layer.apply(y)
+        return y #pas de sigmoid: il est ajouté dans la loss
+
+
+def accuracy(model, x, y):
+    logits = batch_apply(model, x, None)
+    return float(jnp.mean((logits > 0) == (y > 0.5)))   #logit>0 <=> sigmoid>0.5
+
+
+def train_classif(cell_type="gru", dim_hidden=16, verbose=True):
+    key_train, key_val, key_model, key_shuffle = jr.split(jr.key(0), 4)
+    X_train, Y_train = get_data(512, key_train)
+    X_val, Y_val = get_data(256, key_val)
+
+    model = RNN_bin_classifier(dim_hidden, key_model, cell_type=cell_type)
+    acc_before = accuracy(model, X_val, Y_val)
+
+    config = TrainConfig(batch_size=32, learning_rate=0.05, n_epoch=60, lr_final_ratio=0.01)
+    mother_folder = os.path.join(OUT_FOLDER, f"rnn_classif_{cell_type}")
+    best_model, best_loss, folder = train(
+        model, mother_folder, config, X_train, Y_train, X_val, Y_val,
+        rkey=key_shuffle, loss_fn=bce_logits_loss, verbose=verbose)
+
+    acc_after = accuracy(best_model, X_val, Y_val)
+    print(f"{cell_type} : accuracy {acc_before:.3f} -> {acc_after:.3f}"
+          f"   (best val BCE = {best_loss:.3g})")
+    return folder, best_loss, acc_after
+
+
+def train_and_plot(cell_type="gru"):
+    folder, best_loss, acc = train_classif(cell_type)
+    history = load_run(folder, "history")["history"]
+    history.plot(title=f"RNN {cell_type} : BCE={best_loss:.2g}, accuracy={acc:.3f}")
+    plt.tight_layout()
+    plt.show()
+
+
+if __name__ == "__main__":
+    train_and_plot("gru")
+
+
+