yaxlib 0.1.0__py3-none-any.whl

yax/__init__.py

@@ -0,0 +1,28 @@
+"""yax : un mini-framework de réseaux de neurones pour jax, à but pédagogique.
+
+Trois principes :
+- un modèle est un pytree (`yax.Module`) dont les feuilles sont exactement les
+  paramètres : `jax.grad`, `jax.jit` et `optimizer.init(model)` l'acceptent
+  tel quel, sans machinerie de filtrage ;
+- signature uniforme `apply(x, rkey=None)` : `rkey` est la source d'aléatoire
+  (dropout, échantillonnage), jamais un mode. Le mode évaluation se bascule
+  par `model = model.set_inference(True)` — le Trainer s'en charge ;
+- tout ce qui n'est pas un paramètre est statique (`yax.StaticField`), y
+  compris les constantes-tableaux.
+"""
+
+from yax.core import Module, StaticField, tree_at, tree_pprint
+
+from yax.layers.Linear import Linear
+from yax.layers.MLP import MLP
+from yax.layers.Dropout import Dropout
+from yax.layers.LayerNorm import LayerNorm
+from yax.layers.Embedding import Embedding
+from yax.layers.Conv_layer import Conv_layer
+from yax.layers.RNN_layer import RNN_layer, GRUCell, LSTMCell
+from yax.layers.MultiHeadAttention import MultiHeadAttention, causal_mask, padding_mask
+from yax.layers.TransformerBlock import TransformerBlock
+from yax.layers.MessagePassing_layer import MessagePassing_layer
+from yax.layers.positional_encoding import sinusoidal_positional_encoding
+
+__version__ = "0.1.0"

yax/core.py

@@ -0,0 +1,346 @@
+"""Le coeur de yax : la classe Module, sans dépendance à equinox.
+
+Un yax.Module est un pytree jax. Ses champs se déclarent par annotations de
+classe et se rangent en deux familles :
+
+- les champs **dynamiques** (annotation seule) : les feuilles du pytree — les
+  paramètres. Ils ne peuvent contenir QUE des tableaux jax, des sous-modules,
+  ou des conteneurs (list/tuple/dict) de ceux-ci. Tout le reste est refusé à
+  la construction, avec un message explicite : c'est la discipline qui permet
+  `jax.grad(loss)(model)`, `jax.jit` et `optimizer.init(model)` sans aucune
+  machinerie de filtrage ;
+- les champs **statiques** (`= StaticField(default_value=...)`) : des
+  métadonnées rangées dans la STRUCTURE du pytree, invisibles pour `grad`,
+  `tree.map` et les optimiseurs. Ils font partie de l'identité du modèle
+  (changer un statique = changer de structure = recompilation). Un statique
+  peut être un tableau jax : il devient alors une CONSTANTE du modèle
+  (encodage positionnel, grille figée...), cuite dans le code compilé.
+
+Les modules sont immuables : on « modifie » avec yax.tree_at (reconstruction
+d'une copie). Chaque module porte un drapeau statique `inference` (False par
+défaut) ; `model.set_inference(True)` rend une copie où le drapeau est basculé
+récursivement dans tous les sous-modules. Convention d'usage : le mode se
+change rarement (le Trainer s'en occupe), l'aléatoire passe par l'argument
+`rkey` des `apply`.
+"""
+
+import functools
+
+import jax
+import jax.tree_util as jtu
+import numpy as np
+
+
+class StaticField:
+    """Marqueur de champ statique, avec valeur par défaut optionnelle.
+
+    S'utilise comme valeur de classe : `layer_sizes: tuple = StaticField()`.
+    """
+
+    def __init__(self, default_value=None):
+        self.default_value = default_value
+
+
+class _HashableArray:
+    """Enveloppe hachable pour un tableau rangé dans un champ statique.
+
+    La structure d'un pytree doit être hachable (le cache de jit la hache à
+    chaque appel). Les octets et le hash sont précalculés une fois — le
+    tableau est immuable. À réserver aux petites constantes : l'égalité est
+    comparée à chaque dispatch jitté.
+    """
+
+    __slots__ = ("value", "_bytes", "_hash")
+
+    def __init__(self, value):
+        self.value = value
+        arr = np.asarray(value)
+        self._bytes = (arr.shape, str(arr.dtype), arr.tobytes())
+        self._hash = hash(self._bytes)
+
+    def __eq__(self, other):
+        return type(other) is _HashableArray and self._bytes == other._bytes
+
+    def __hash__(self):
+        return self._hash
+
+
+def _est_tableau(x):
+    return isinstance(x, jax.Array)
+
+
+def _valide_dynamique(cls_name, name, value):
+    """Un champ dynamique : tableau jax, Module, ou conteneur de ceux-ci."""
+    if _est_tableau(value) or isinstance(value, Module):
+        return
+    if isinstance(value, (list, tuple)):
+        for v in value:
+            _valide_dynamique(cls_name, name, v)
+        return
+    if isinstance(value, dict):
+        for v in value.values():
+            _valide_dynamique(cls_name, name, v)
+        return
+    if isinstance(value, np.ndarray):
+        raise TypeError(
+            f"{cls_name}.{name} : tableau numpy dans un champ dynamique — "
+            f"convertissez avec jnp.asarray(...), ou déclarez le champ "
+            f"StaticField si c'est une constante.")
+    raise TypeError(
+        f"{cls_name}.{name} : un champ dynamique ne peut contenir que des "
+        f"tableaux jax, des sous-modules, ou des list/tuple/dict de ceux-ci — "
+        f"reçu {type(value).__name__} ({value!r:.60}). Si ce n'est pas un "
+        f"paramètre, déclarez-le : {name}: ... = StaticField().")
+
+
+def _valide_statique(cls_name, name, value):
+    """Un champ statique : hachable, ou tableau (enveloppé au vol)."""
+    if _est_tableau(value) or isinstance(value, np.ndarray):
+        return
+    if isinstance(value, list):
+        raise TypeError(
+            f"{cls_name}.{name} : une liste dans un champ statique — un "
+            f"statique fait partie de l'identité du modèle et doit être "
+            f"immuable : utilisez un tuple.")
+    try:
+        hash(value)
+    except TypeError:
+        raise TypeError(
+            f"{cls_name}.{name} : un champ statique doit être hachable "
+            f"(ou être un tableau) — reçu {type(value).__name__}.")
+
+
+class Module:
+    """Classe de base : voir la docstring du module."""
+
+    inference: bool = StaticField(default_value=False)
+
+    _initializing = False   # défaut de classe : les instances unflatten le gardent
+
+    # ------------------------------------------------------------- champs
+    @classmethod
+    def _collecte_champs(cls):
+        statics, dynamics = {}, []
+        for klass in reversed(cls.__mro__):
+            for name in getattr(klass, "__annotations__", {}):
+                default = klass.__dict__.get(name, None)
+                if isinstance(default, StaticField):
+                    statics[name] = default.default_value
+                    if name in dynamics:
+                        dynamics.remove(name)
+                elif name not in statics and name not in dynamics:
+                    if default is not None:
+                        raise TypeError(
+                            f"{cls.__name__}.{name} : valeur par défaut sur un "
+                            f"champ dynamique — non supporté (les dynamiques "
+                            f"s'initialisent dans __init__ ; pour une "
+                            f"métadonnée, utilisez StaticField(default_value=...)).")
+                    dynamics.append(name)
+        cls._yax_statics = statics            # nom -> valeur par défaut
+        cls._yax_dynamics = tuple(dynamics)   # ordre de déclaration
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        cls._collecte_champs()
+
+        init = cls.__init__
+        if not getattr(init, "_yax_wrapped", False):
+            @functools.wraps(init)
+            def init_encadre(self, *args, **kw):
+                object.__setattr__(self, "_initializing", True)
+                init(self, *args, **kw)
+                object.__setattr__(self, "_initializing", False)
+                self._fin_de_construction()
+
+            init_encadre._yax_wrapped = True
+            cls.__init__ = init_encadre
+
+        jtu.register_pytree_with_keys(
+            cls,
+            flatten_with_keys=cls._flatten_with_keys,
+            flatten_func=cls._flatten,
+            unflatten_func=cls._unflatten)
+
+    def _fin_de_construction(self):
+        cls = type(self)
+        for name, default in cls._yax_statics.items():
+            if name not in self.__dict__:
+                object.__setattr__(self, name, default)
+        for name in cls._yax_dynamics:
+            if name not in self.__dict__:
+                raise AttributeError(
+                    f"{cls.__name__} : le champ dynamique '{name}' n'a pas été "
+                    f"initialisé dans __init__.")
+
+    # -------------------------------------------------------- immutabilité
+    def __setattr__(self, name, value):
+        cls = type(self)
+        if not self._initializing:
+            raise AttributeError(
+                f"{cls.__name__} est immuable : impossible d'assigner "
+                f"'{name}'. Pour « modifier » un module, reconstruire une "
+                f"copie avec yax.tree_at.")
+        if name in cls._yax_statics:
+            _valide_statique(cls.__name__, name, value)
+        elif name in cls._yax_dynamics:
+            _valide_dynamique(cls.__name__, name, value)
+        else:
+            raise AttributeError(
+                f"{cls.__name__}.{name} : champ non déclaré — ajoutez "
+                f"l'annotation de classe correspondante.")
+        object.__setattr__(self, name, value)
+
+    # ------------------------------------------------------------- pytree
+    def _aux(self):
+        cls = type(self)
+        out = []
+        for name in cls._yax_statics:
+            v = self.__dict__.get(name)
+            if _est_tableau(v) or isinstance(v, np.ndarray):
+                v = self.__dict__.setdefault("_wrapped_" + name, _HashableArray(v))
+            out.append((name, v))
+        return tuple(out)
+
+    def _flatten(self):
+        cls = type(self)
+        children = tuple(self.__dict__.get(n) for n in cls._yax_dynamics)
+        return children, self._aux()
+
+    def _flatten_with_keys(self):
+        cls = type(self)
+        children = tuple((jtu.GetAttrKey(n), self.__dict__.get(n))
+                         for n in cls._yax_dynamics)
+        return children, self._aux()
+
+    @classmethod
+    def _unflatten(cls, aux, children):
+        # aucun __init__, aucune validation : jax passe ici avec des tracers,
+        # des gradients, ou les sorties de n'importe quel tree.map
+        obj = object.__new__(cls)
+        for name, value in zip(cls._yax_dynamics, children):
+            object.__setattr__(obj, name, value)
+        for name, value in aux:
+            if type(value) is _HashableArray:
+                object.__setattr__(obj, "_wrapped_" + name, value)
+                value = value.value
+            object.__setattr__(obj, name, value)
+        return obj
+
+    # ------------------------------------------------------------ services
+    def set_inference(self, value: bool):
+        """Copie du module où le drapeau statique `inference` vaut `value`,
+        récursivement dans tous les sous-modules. L'original est intact."""
+
+        def reconstruit(x):
+            if isinstance(x, Module):
+                klass = type(x)
+                obj = object.__new__(klass)
+                for name in klass._yax_dynamics:
+                    object.__setattr__(obj, name, reconstruit(x.__dict__.get(name)))
+                for name in klass._yax_statics:
+                    object.__setattr__(obj, name, x.__dict__.get(name))
+                object.__setattr__(obj, "inference", bool(value))
+                return obj
+            if isinstance(x, list):
+                return [reconstruit(v) for v in x]
+            if isinstance(x, tuple):
+                return tuple(reconstruit(v) for v in x)
+            if isinstance(x, dict):
+                return {k: reconstruit(v) for k, v in x.items()}
+            return x
+
+        return reconstruit(self)
+
+    # -------------------------------------------------------------- repr
+    def __repr__(self):
+        return _pformat(self, indent=0)
+
+
+def _pformat(x, indent):
+    pad = "  " * indent
+    if isinstance(x, Module):
+        cls = type(x)
+        noms = list(cls._yax_dynamics) + list(cls._yax_statics)
+        lignes = [f"{pad}{cls.__name__}("]
+        for name in noms:
+            valeur = _pformat(x.__dict__.get(name), indent + 1).lstrip()
+            lignes.append(f"{pad}  {name}={valeur},")
+        lignes.append(f"{pad})")
+        return "\n".join(lignes)
+    if _est_tableau(x) or isinstance(x, np.ndarray):
+        dtype = str(x.dtype).replace("float", "f").replace("int", "i").replace("bool", "b")
+        return pad + f"{dtype}[{','.join(map(str, x.shape))}]"
+    if isinstance(x, (list, tuple)):
+        if not x:
+            return pad + repr(x)
+        ouvre, ferme = ("[", "]") if isinstance(x, list) else ("(", ")")
+        interieur = ",\n".join(_pformat(v, indent + 1) for v in x)
+        return f"{pad}{ouvre}\n{interieur}\n{pad}{ferme}"
+    if isinstance(x, dict):
+        if not x:
+            return pad + "{}"
+        lignes = [pad + "{"]
+        for k, v in x.items():
+            lignes.append(f"{pad}  {k!r}: {_pformat(v, indent + 1).lstrip()},")
+        lignes.append(pad + "}")
+        return "\n".join(lignes)
+    return pad + repr(x)
+
+
+def tree_pprint(x):
+    print(_pformat(x, 0))
+
+
+def tree_at(where, pytree, replace):
+    """Reconstruit `pytree` en remplaçant une ou plusieurs feuilles.
+
+    `where` : fonction qui, appliquée au pytree, rend LA feuille à remplacer
+    (ou un tuple/liste de feuilles). `replace` : la ou les valeurs neuves.
+
+        model2 = tree_at(lambda m: m.layers[-1].bias, model, jnp.ones(3))
+
+    Les feuilles sont repérées par identité (`is`) : `where` doit rendre des
+    objets extraits du pytree lui-même. L'original est intact.
+    """
+    cibles = where(pytree)
+    if not isinstance(cibles, (tuple, list)):
+        cibles, replace = (cibles,), (replace,)
+    if len(cibles) != len(replace):
+        raise ValueError(f"tree_at : {len(cibles)} cible(s) mais "
+                         f"{len(replace)} remplacement(s).")
+
+    leaves, treedef = jtu.tree_flatten(pytree)
+    nouvelles = list(leaves)
+    for cible, valeur in zip(cibles, replace):
+        indices = [i for i, l in enumerate(leaves) if l is cible]
+        if len(indices) != 1:
+            raise ValueError(
+                "tree_at : cible introuvable ou ambigue — `where` doit rendre "
+                "une feuille (un tableau) extraite du pytree lui-même.")
+        nouvelles[indices[0]] = valeur
+    return jtu.tree_unflatten(treedef, nouvelles)
+
+
+if __name__ == "__main__":
+    import jax.numpy as jnp
+    import jax.random as jr
+
+    class Petit(Module):
+        w: jnp.ndarray
+        nom: str = StaticField(default_value="petit")
+
+        def __init__(self, rkey):
+            self.w = jr.normal(rkey, (3,))
+
+        def apply(self, x, rkey=None):
+            return jnp.sum(self.w * x)
+
+    m = Petit(jr.key(0))
+    print(m)
+    print("feuilles :", jax.tree.leaves(m))
+    g = jax.grad(lambda mm: mm.apply(jnp.ones(3)))(m)
+    print("gradient :", type(g).__name__, g.w, "| nom conservé :", g.nom)
+    m2 = m.set_inference(True)
+    print("inference :", m.inference, "->", m2.inference)
+    m3 = tree_at(lambda mm: mm.w, m, jnp.zeros(3))
+    print("tree_at :", m3.w, "| original intact :", m.w)

yax/image/augmentation.py

@@ -0,0 +1,123 @@
+"""Enrichissement des données (augmentation), chapitre 11.2.
+
+Des transformations qui PRÉSERVENT l'étiquette, écrites pour UNE image
+channels-first (canaux, H, W). Tout est différentiable et vmap-able :
+l'augmentation se fait sur le device, dans le pas d'entraînement — pas dans un
+chargeur de données à part. L'aléatoire vient d'une clé, comme partout : une
+transformation différente à chaque époque pour un même exemple.
+
+Convention : f(rkey, img, ...) — la clé d'abord, comme dans jax.random. Ce sont
+des fonctions, pas des modules : elles n'ont aucun paramètre appris.
+
+Piège du cours, à répéter : choisir des transformations qui changent
+l'étiquette. Une symétrie horizontale sur des chiffres manuscrits transforme
+un 2 en quelque chose qui n'est plus un 2. L'invariance qu'on impose doit
+être vraie.
+"""
+
+import jax
+import jax.numpy as jnp
+import jax.random as jr
+
+
+def random_horizontal_flip(rkey, img, prob=0.5):
+    """Symétrie gauche-droite avec probabilité prob."""
+    # jnp.where et non un `if` : le tirage est une valeur tracée (chapitre 5)
+    flip = jr.bernoulli(rkey, prob)
+    return jnp.where(flip, img[:, :, ::-1], img)
+
+
+def random_brightness_contrast(rkey, img, max_shift=0.2, max_factor=0.2):
+    """Luminosité : décalage global ; contraste : dilatation autour de la
+    moyenne de l'image."""
+    k_shift, k_factor = jr.split(rkey)
+    shift = jr.uniform(k_shift, minval=-max_shift, maxval=max_shift)
+    factor = 1.0 + jr.uniform(k_factor, minval=-max_factor, maxval=max_factor)
+    mean = jnp.mean(img)
+    return (img - mean) * factor + mean + shift
+
+
+def random_noise(rkey, img, sigma=0.05):
+    return img + sigma * jr.normal(rkey, img.shape)
+
+
+def random_scale_translate(rkey, img, max_zoom=0.2, max_shift=2.0):
+    """Homothétie (zoom autour du centre) et translation, par
+    jax.image.scale_and_translate. max_shift est en pixels."""
+    C, H, W = img.shape
+    k_zoom, k_shift = jr.split(rkey)
+    zoom = 1.0 + jr.uniform(k_zoom, minval=-max_zoom, maxval=max_zoom)
+    shift = jr.uniform(k_shift, (2,), minval=-max_shift, maxval=max_shift)
+    # scale_and_translate applique sortie(y) = entree((y - translation)/scale) :
+    # cette translation-ci recentre le zoom sur le milieu de l'image
+    scale = jnp.array([zoom, zoom])
+    translation = (1.0 - zoom) * jnp.array([H, W]) / 2.0 + shift
+    return jax.image.scale_and_translate(img, img.shape, (1, 2),
+                                         scale, translation, method="linear")
+
+
+def random_rotation(rkey, img, max_angle_degrees=15.0):
+    """Rotation d'angle quelconque autour du centre, par map_coordinates :
+    on construit la grille de coordonnées TOURNÉE, on interpole (ordre 1).
+    C'est la rotation inverse qu'on applique aux coordonnées — pour remplir
+    chaque pixel de la sortie, on va chercher d'où il vient dans l'entrée."""
+    C, H, W = img.shape
+    angle = jnp.deg2rad(jr.uniform(rkey, minval=-max_angle_degrees,
+                                   maxval=max_angle_degrees))
+    cos, sin = jnp.cos(angle), jnp.sin(angle)
+    rows = jnp.arange(H) - (H - 1) / 2.0
+    cols = jnp.arange(W) - (W - 1) / 2.0
+    r, c = jnp.meshgrid(rows, cols, indexing="ij")
+    src_r = cos * r - sin * c + (H - 1) / 2.0
+    src_c = sin * r + cos * c + (W - 1) / 2.0
+
+    def rotate_channel(channel):
+        return jax.scipy.ndimage.map_coordinates(channel, [src_r, src_c],
+                                                 order=1, mode="constant", cval=0.0)
+
+    return jax.vmap(rotate_channel)(img)
+
+
+def random_augmentation(rkey, img):
+    """Pipeline d'exemple : composer, c'est splitter la clé et enchaîner.
+    À adapter au jeu de données — le flip n'y est volontairement pas, c'est
+    la transformation qu'il faut choisir en connaissance de cause."""
+    k1, k2, k3, k4 = jr.split(rkey, 4)
+    img = random_scale_translate(k1, img)
+    img = random_rotation(k2, img)
+    img = random_brightness_contrast(k3, img)
+    img = random_noise(k4, img)
+    return img
+
+
+if __name__ == "__main__":
+    import matplotlib
+
+    # une image test : un carré excentré avec un coin marqué (l'asymétrie rend
+    # les transformations visibles)
+    img = jnp.zeros((1, 32, 32))
+    img = img.at[0, 8:20, 6:18].set(0.6)
+    img = img.at[0, 8:12, 6:10].set(1.0)
+
+    # formes conservées, jit + vmap passent
+    batch = jnp.stack([img] * 8)
+    keys = jr.split(jr.key(0), 8)
+    augmented = jax.jit(jax.vmap(random_augmentation))(keys, batch)
+    assert augmented.shape == batch.shape
+    # même clé -> même transformation (reproductible). Tolérance large : jit
+    # réordonne les flottants, eager et compilé diffèrent au dernier bit.
+    again = jax.vmap(random_augmentation)(keys, batch)
+    assert jnp.allclose(augmented, again, atol=1e-5)
+    print("jit + vmap + reproductibilité : OK")
+
+    from matplotlib import pyplot as plt
+
+    fig, axs = plt.subplots(1, 9, figsize=(16, 2.2))
+    axs[0].imshow(img[0], cmap="gray", vmin=-0.2, vmax=1.2)
+    axs[0].set_title("originale")
+    for i in range(8):
+        axs[i + 1].imshow(augmented[i, 0], cmap="gray", vmin=-0.2, vmax=1.2)
+    for ax in axs:
+        ax.axis("off")
+    plt.tight_layout()
+    plt.show()

yax/layers/Conv_layer.py

@@ -0,0 +1,70 @@
+"""Couche de convolution 2D.
+
+Implémentée directement sur lax.conv_general_dilated, le primitif XLA — c'est
+exactement ce que ferait une bibliothèque, il n'y a aucune optimisation cachée
+au-dessus. L'API du primitif est générale donc verbeuse (les dimension_numbers
+disent qui est batch, canaux, hauteur, largeur) ; cette couche la fige dans la
+convention du cours.
+
+Convention yax2 : écrit pour UN échantillon, en channels-first (canaux, H, W).
+Le batch vient de vmap. Initialisation de Glorot, avec le fan_in d'une
+convolution : canaux d'entrée x hauteur x largeur du noyau.
+"""
+
+from yax.core import Module, StaticField, tree_at
+import jax
+import jax.numpy as jnp
+import jax.random as jr
+from jax import lax
+
+
+class Conv_layer(Module):
+    weight: jnp.ndarray            # (dim_out, dim_in, k, k)
+    bias: jnp.ndarray              # (dim_out,)
+
+    stride: int = StaticField()
+    padding: str = StaticField()
+
+    def __init__(self, dim_in, dim_out, kernel_size, rkey, *, stride=1,
+                 padding="SAME"):
+        # padding="SAME" : la sortie garde la taille spatiale de l'entrée
+        # (à stride 1) ; "VALID" : pas de remplissage, la taille fond de
+        # kernel_size-1. stride>1 : l'alternative au pooling.
+        assert padding in ("SAME", "VALID")
+        fan_in = dim_in * kernel_size ** 2
+        fan_out = dim_out * kernel_size ** 2
+        lim = jnp.sqrt(6.0 / (fan_in + fan_out))
+        self.weight = jr.uniform(rkey, (dim_out, dim_in, kernel_size, kernel_size),
+                                 minval=-lim, maxval=lim)
+        self.bias = jnp.zeros((dim_out,))
+        self.stride = stride
+        self.padding = padding
+
+    def apply(self, x, rkey=None):
+        # x : (canaux, H, W) — rkey ignoré : couche déterministe.
+        # Le primitif attend un batch : on en fabrique un de taille 1.
+        y = lax.conv_general_dilated(
+            x[None],                                    # (1, dim_in, H, W)
+            self.weight,                                # (dim_out, dim_in, k, k)
+            window_strides=(self.stride, self.stride),
+            padding=self.padding,
+            dimension_numbers=("NCHW", "OIHW", "NCHW"))
+        return y[0] + self.bias[:, None, None]
+
+
+if __name__ == "__main__":
+    import equinox as eqx  # reference de verification uniquement
+    # formes, puis vérification numérique contre eqx.nn.Conv2d — le schéma
+    # constant : réimplémenter, vérifier contre la référence, adopter LA NÔTRE
+    x = jr.normal(jr.key(0), (3, 16, 16))
+    print("SAME  stride 1 :", Conv_layer(3, 8, 3, jr.key(1)).apply(x).shape)
+    print("SAME  stride 2 :", Conv_layer(3, 8, 3, jr.key(1), stride=2).apply(x).shape)
+    print("VALID stride 1 :", Conv_layer(3, 8, 3, jr.key(1), padding="VALID").apply(x).shape)
+
+    for stride, padding in [(1, "SAME"), (2, "SAME"), (1, "VALID")]:
+        ref = eqx.nn.Conv2d(3, 8, 3, stride=stride, padding=padding, key=jr.key(2))
+        ours = Conv_layer(3, 8, 3, jr.key(3), stride=stride, padding=padding)
+        ours = tree_at(lambda m: (m.weight, m.bias), ours,
+                           (ref.weight, ref.bias.reshape(-1)))
+        assert jnp.allclose(ours.apply(x), ref(x), atol=1e-5), (stride, padding)
+    print("Conv_layer == eqx.nn.Conv2d : OK (SAME/VALID, stride 1/2)")

yax/layers/Dropout.py

@@ -0,0 +1,63 @@
+"""Dropout.
+
+Deux choses distinctes, et découplées :
+
+- le MODE : le drapeau statique `inference`, porté par tout yax.Module et
+  basculé par `model = model.set_inference(True/False)` — le Trainer s'en
+  charge (entraînement en False, validation et modèle rendu en True) ;
+- la SOURCE D'ALÉATOIRE : l'argument `rkey` de `apply`.
+
+En mode inférence, le dropout est l'identité, clé ou pas. En mode
+entraînement, la clé est obligatoire — l'oublier est une erreur explicite,
+pas un comportement silencieux.
+
+Le découplage a un bonus : le MC-dropout (échantillonner des prédictions avec
+le dropout actif) s'obtient en repassant simplement le modèle en
+`set_inference(False)` et en fournissant des clés à l'évaluation.
+"""
+
+import jax.numpy as jnp
+import jax.random as jr
+
+from yax.core import Module, StaticField
+
+
+class Dropout(Module):
+    rate: float = StaticField(default_value=0.0)
+
+    def __init__(self, rate):
+        assert 0.0 <= rate < 1.0, f"rate:{rate} doit etre dans [0,1)"
+        self.rate = rate
+
+    def apply(self, x, rkey=None):
+        if self.inference or self.rate == 0.0:
+            return x
+        if rkey is None:
+            raise ValueError(
+                "Dropout en mode entrainement (inference=False) : une cle "
+                "est requise — apply(x, rkey). Pour evaluer, basculer le "
+                "modele avec model.set_inference(True).")
+        keep_prob = 1.0 - self.rate
+        mask = jr.bernoulli(rkey, keep_prob, x.shape)
+        # division par keep_prob : l'esperance de la sortie est inchangee,
+        # rien a corriger a l'evaluation ("inverted dropout").
+        return jnp.where(mask, x / keep_prob, 0.0)
+
+
+if __name__ == "__main__":
+    drop = Dropout(0.4)
+    x = jnp.ones((10_000,))
+
+    y_train = drop.apply(x, jr.key(0))
+    print(f"fraction annulée : {float(jnp.mean(y_train == 0.0)):.3f} (attendu ~0.4)")
+    print(f"moyenne          : {float(jnp.mean(y_train)):.3f} (attendu ~1.0)")
+
+    try:
+        drop.apply(x)   # entrainement sans cle : erreur explicite
+    except ValueError as e:
+        print("ValueError :", str(e)[:60], "...")
+
+    drop_eval = drop.set_inference(True)
+    assert (drop_eval.apply(x) == x).all()
+    assert (drop_eval.apply(x, jr.key(0)) == x).all()   # cle ignoree en inference
+    print("mode inference : identité, avec ou sans clé : OK")

yax/layers/Embedding.py

@@ -0,0 +1,29 @@
+"""Embedding : une table apprise, indexée par des entiers.
+
+Le one-hot suivi d'un Linear EST un embedding, en moins efficace (chapitre 8) :
+ici on lit directement la ligne de la table.
+"""
+
+from yax.core import Module
+import jax.numpy as jnp
+import jax.random as jr
+
+
+class Embedding(Module):
+    weight: jnp.ndarray
+
+    def __init__(self, num_embeddings, dim, rkey):
+        # initialisation normale d'écart-type 0.02, l'usage des transformers
+        self.weight = 0.02 * jr.normal(rkey, (num_embeddings, dim))
+
+    def apply(self, ids, rkey=None):
+        # ids : entiers, forme quelconque -> sortie ids.shape + (dim,).
+        # Indexation par entiers : valide sous jit (contrairement au masque
+        # booléen entre crochets, chapitre 1). rkey ignoré : couche déterministe.
+        return self.weight[ids]
+
+
+if __name__ == "__main__":
+    emb = Embedding(100, 32, jr.key(0))
+    ids = jnp.array([1, 5, 10])
+    print("sortie :", emb.apply(ids).shape)  # (3, 32)