TP7: Le traceur

IFT3395/IFT6390 — Fondements de l’apprentissage machine

Ce notebook accompagne le Chapitre 7: Réseaux de neurones, section « Le traceur ». Vous y implémenterez un moteur de différentiation automatique en mode inverse en explorant deux mécanismes d’interception: la surcharge d’opérateurs Python et l’approche par espace de noms (à la autograd/JAX).

Objectifs¶

À la fin de ce TP, vous serez en mesure de:

Expliquer le rôle du traceur dans un système de DA en mode inverse
Implémenter un traceur scalaire en surchargeant les opérateurs de float
Construire une fonction grad analogue à jax.grad
Vérifier les gradients par comparaison analytique et par différences finies
Extraire un graphe de calcul à partir de la bande d’enregistrement
Comparer la surcharge d’opérateurs et l’approche par espace de noms pour l’interception des opérations

Prérequis: Ch7, sections « Différentiation automatique » et « Le traceur ».

Partie 0: Configuration¶

import math
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt

plt.rcParams['figure.figsize'] = (8, 5)
plt.rcParams['font.size'] = 12
print("Configuration terminée!")

Configuration terminée!

Partie 1: Pourquoi automatiser la construction du graphe?¶

Le notebook précédent (Différentiation automatique avec NetworkX) montrait comment construire manuellement un graphe de calcul avec NetworkX, puis le parcourir en mode inverse. Cette approche est pédagogique, mais elle a trois défauts:

Redondance: on écrit chaque opération deux fois (une fois dans le graphe, une fois pour évaluer la fonction)
Fragilité: toute modification de la fonction oblige à reconstruire le graphe à la main
Séparation: la description du calcul et son exécution vivent dans deux endroits différents

L’idée du traceur résout ces trois problèmes d’un coup: on encapsule chaque nombre dans un objet qui surcharge les opérateurs Python (+, *, etc.) et enregistre automatiquement les opérations sur une bande (tape). La fonction s’écrit normalement; la bande se construit toute seule.

Bibliothèque de règles VJP¶

Chaque opération primitive a une règle VJP (vector-Jacobian product) qui calcule les cotangents des entrées à partir du cotangent de la sortie. La signature commune est vjp(résidus, cotangent_sortie) → cotangents_entrées.

# ---- Bibliothèque de règles VJP ----
# Signature commune: vjp(résidus, cotangent_sortie) → cotangents_entrées

def add_vjp(res, g):
    return (g, g)                        # ∂(a+b)/∂a = 1, ∂(a+b)/∂b = 1

def mul_vjp(res, g):
    a, b = res
    return (b * g, a * g)                # ∂(a·b)/∂a = b, ∂(a·b)/∂b = a

def sin_vjp(res, g):
    (a,) = res
    return (math.cos(a) * g,)            # ∂sin(a)/∂a = cos(a)

def exp_vjp(res, g):
    (a,) = res
    return (math.exp(a) * g,)            # ∂exp(a)/∂a = exp(a)

def neg_vjp(res, g):
    (a,) = res
    return (-g,)                         # ∂(-a)/∂a = -1

Partie 2: Le traceur — classe `Var`¶

Le traceur est un objet qui se fait passer pour un nombre. Il encapsule:

data: la valeur numérique (un float)
id: un identifiant unique (entier auto-incrémenté)

À chaque opération arithmétique, le traceur:

Calcule le résultat numérique (passe avant)
Enregistre une entrée sur la bande globale _tape: la règle VJP, les résidus, les identifiants des entrées, et l’identifiant de la sortie
Retourne un nouveau Var contenant le résultat

La méthode _record (donnée ci-dessous) effectue les étapes 2 et 3. L’opérateur __add__ est donné comme exemple.

_tape = []      # bande globale: [(vjp_fn, résidus, ids_entrées, id_sortie)]
_n_vars = 0     # compteur d'identifiants

class Var:
    """Traceur scalaire: encapsule un flottant et un identifiant unique."""

    def __init__(self, data):
        global _n_vars
        self.data = float(data)
        self.id = _n_vars
        _n_vars += 1

    def _record(self, vjp_fn, res, inputs, out_data):
        """Enregistre une opération sur la bande et retourne un nouveau Var."""
        out = Var(out_data)
        _tape.append((vjp_fn, res, [v.id for v in inputs], out.id))
        return out

    # --- Exemple: addition ---
    def __add__(self, other):
        other = other if isinstance(other, Var) else Var(other)
        return self._record(add_vjp, (self.data, other.data),
                            [self, other], self.data + other.data)

    def __radd__(self, other): return self.__add__(other)

    # ============================================
    # Exercice 1 ★★: Implémenter les opérateurs manquants
    #
    # Sur le modèle de __add__ ci-dessus, implémentez:
    #   __mul__(self, other)  — multiplication
    #   __rmul__(self, other) — multiplication (opérande gauche non-Var)
    #   sin(self)             — sinus
    #   exp(self)             — exponentielle
    #   __neg__(self)         — négation unaire
    #
    # Chaque méthode appelle self._record avec la bonne
    # règle VJP, les bons résidus, et le bon résultat numérique.
    # ============================================

    def __mul__(self, other):
        pass  # <- Complétez

    def __rmul__(self, other):
        pass  # <- Complétez

    def sin(self):
        pass  # <- Complétez

    def exp(self):
        pass  # <- Complétez

    def __neg__(self):
        pass  # <- Complétez

def __mul__(self, other):
    other = other if isinstance(other, Var) else Var(other)
    return self._record(mul_vjp, (self.data, other.data),
                        [self, other], self.data * other.data)

def __rmul__(self, other): return self.__mul__(other)

def sin(self):
    return self._record(sin_vjp, (self.data,),
                        [self], math.sin(self.data))

def exp(self):
    return self._record(exp_vjp, (self.data,),
                        [self], math.exp(self.data))

def __neg__(self):
    return self._record(neg_vjp, (self.data,),
                        [self], -self.data)

# Vérification: tracer f(x,y) = sin(xy) + exp(x)
_tape, _n_vars = [], 0
x = Var(0.5)
y = Var(3.0)
result = (x * y).sin() + x.exp()

print(f"f(0.5, 3.0) = {result.data:.6f}")
print(f"Vérification: sin(0.5*3.0) + exp(0.5) = {math.sin(1.5) + math.exp(0.5):.6f}")
print(f"\nContenu de la bande ({len(_tape)} entrées):")
for vjp_fn, res, in_ids, out_id in _tape:
    op = vjp_fn.__name__.replace('_vjp', '')
    print(f"  {op}: entrées={in_ids} → sortie={out_id}, résidus={res}")

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[4], line 5
      3 x = Var(0.5)
      4 y = Var(3.0)
----> 5 result = (x * y).sin() + x.exp()
      7 print(f"f(0.5, 3.0) = {result.data:.6f}")
      8 print(f"Vérification: sin(0.5*3.0) + exp(0.5) = {math.sin(1.5) + math.exp(0.5):.6f}")

AttributeError: 'NoneType' object has no attribute 'sin'

Partie 3: La fonction `grad`¶

La fonction grad est un opérateur d’ordre supérieur: elle prend une fonction f et retourne une nouvelle fonction qui calcule le gradient de f. C’est l’analogue de jax.grad.

Son fonctionnement en trois étapes:

Réinitialiser la bande et le compteur
Passe avant: encapsuler les entrées dans des Var, appeler f — la bande se remplit automatiquement
Passe arrière: parcourir la bande en ordre inverse, propager les cotangents par accumulation

def grad(f):
    """Retourne une fonction qui calcule le gradient de f."""
    def grad_fn(*args):
        global _tape, _n_vars
        _tape, _n_vars = [], 0               # réinitialiser la bande

        # Passe avant: tracer l'exécution
        traced = [Var(a) for a in args]
        result = f(*traced)

        # Passe arrière: propager les cotangents
        adjoints = [0.0] * _n_vars
        adjoints[result.id] = 1.0            # ∂f/∂f = 1

        # ============================================
        # Exercice 2 ★★: Compléter la passe arrière
        #
        # Parcourez la bande en ordre inverse avec reversed(_tape).
        # Pour chaque entrée (vjp_fn, res, in_ids, out_id):
        #   1. Récupérez le cotangent de la sortie: adjoints[out_id]
        #   2. Appelez vjp_fn(res, cotangent) pour obtenir les cotangents
        #   3. Accumulez: adjoints[idx] += ct pour chaque (idx, ct)
        # ============================================

        pass  # <- Complétez

        return tuple(adjoints[v.id] for v in traced)
    return grad_fn

Partie 4: Vérification¶

Testons le traceur sur les mêmes fonctions que le notebook NetworkX, avec les mêmes valeurs numériques, pour comparer directement.

Fonction $f(x,y) = \sin(xy) + e^x$ ¶

Gradients analytiques:

\frac{\partial f}{\partial x} = y\cos(xy) + e^x, \qquad \frac{\partial f}{\partial y} = x\cos(xy)

(1)

def f(x, y):
    return (x * y).sin() + x.exp()

x_val, y_val = 0.5, 3.0
df_dx, df_dy = grad(f)(x_val, y_val)

# Gradients analytiques
df_dx_exact = y_val * math.cos(x_val * y_val) + math.exp(x_val)
df_dy_exact = x_val * math.cos(x_val * y_val)

# Différences finies
def f_np(x, y):
    return np.sin(x * y) + np.exp(x)

eps = 1e-7
df_dx_num = (f_np(x_val + eps, y_val) - f_np(x_val - eps, y_val)) / (2 * eps)
df_dy_num = (f_np(x_val, y_val + eps) - f_np(x_val, y_val - eps)) / (2 * eps)

print("--- f(x,y) = sin(xy) + exp(x) ---")
print(f"f({x_val}, {y_val}) = {f_np(x_val, y_val):.6f}")
print(f"\n{'Méthode':<25} {'∂f/∂x':>12} {'∂f/∂y':>12}")
print(f"{'-'*25} {'-'*12} {'-'*12}")
print(f"{'Traceur':<25} {df_dx:>12.6f} {df_dy:>12.6f}")
print(f"{'Analytique':<25} {df_dx_exact:>12.6f} {df_dy_exact:>12.6f}")
print(f"{'Différences finies':<25} {df_dx_num:>12.6f} {df_dy_num:>12.6f}")

assert np.allclose(df_dx, df_dx_exact), f"Erreur sur ∂f/∂x: {df_dx} ≠ {df_dx_exact}"
assert np.allclose(df_dy, df_dy_exact), f"Erreur sur ∂f/∂y: {df_dy} ≠ {df_dy_exact}"
print("\nTous les gradients correspondent.")

Fonction $g(x,y) = -(\sin(x) \cdot \sin(y))$ ¶

Gradients analytiques:

\frac{\partial g}{\partial x} = -\cos(x)\sin(y), \qquad \frac{\partial g}{\partial y} = -\sin(x)\cos(y)

(2)

def g(x, y):
    return -(x.sin() * y.sin())

x_val2, y_val2 = 1.0, 2.0
dg_dx, dg_dy = grad(g)(x_val2, y_val2)

# Gradients analytiques
dg_dx_exact = -math.cos(x_val2) * math.sin(y_val2)
dg_dy_exact = -math.sin(x_val2) * math.cos(y_val2)

# Différences finies
def g_np(x, y):
    return -(np.sin(x) * np.sin(y))

dg_dx_num = (g_np(x_val2 + eps, y_val2) - g_np(x_val2 - eps, y_val2)) / (2 * eps)
dg_dy_num = (g_np(x_val2, y_val2 + eps) - g_np(x_val2, y_val2 - eps)) / (2 * eps)

print("--- g(x,y) = -(sin(x)·sin(y)) ---")
print(f"g({x_val2}, {y_val2}) = {g_np(x_val2, y_val2):.6f}")
print(f"\n{'Méthode':<25} {'∂g/∂x':>12} {'∂g/∂y':>12}")
print(f"{'-'*25} {'-'*12} {'-'*12}")
print(f"{'Traceur':<25} {dg_dx:>12.6f} {dg_dy:>12.6f}")
print(f"{'Analytique':<25} {dg_dx_exact:>12.6f} {dg_dy_exact:>12.6f}")
print(f"{'Différences finies':<25} {dg_dx_num:>12.6f} {dg_dy_num:>12.6f}")

assert np.allclose(dg_dx, dg_dx_exact), f"Erreur sur ∂g/∂x: {dg_dx} ≠ {dg_dx_exact}"
assert np.allclose(dg_dy, dg_dy_exact), f"Erreur sur ∂g/∂y: {dg_dy} ≠ {dg_dy_exact}"
print("\nTous les gradients correspondent.")

Partie 5: Visualisation du graphe et récapitulatif¶

La bande contient toute l’information nécessaire pour reconstruire le graphe de calcul. Contrairement au notebook précédent où le graphe était construit à la main, ici il est extrait automatiquement de la trace.

# ============================================
# Exercice 3 ★: Extraire un DiGraph de la bande
#
# Complétez tape_to_graph pour construire un nx.DiGraph à partir
# de la bande _tape remplie lors du dernier appel à grad.
#
# Pour chaque entrée (vjp_fn, res, in_ids, out_id) de la bande:
#   - Récupérez le nom de l'opération: vjp_fn.__name__.replace('_vjp', '')
#   - Ajoutez le noeud out_id avec un attribut label (ex. "v2 (mul)")
#   - Ajoutez une arête de chaque in_id vers out_id
# ============================================

def tape_to_graph(tape, input_names):
    G = nx.DiGraph()
    # Ajouter les noeuds d'entrée
    for i, name in enumerate(input_names):
        G.add_node(i, label=name)

    pass  # <- Complétez: parcourir la bande et ajouter noeuds + arêtes

    return G

def tape_to_graph(tape, input_names):
    G = nx.DiGraph()
    for i, name in enumerate(input_names):
        G.add_node(i, label=name)
    for vjp_fn, res, in_ids, out_id in tape:
        op_name = vjp_fn.__name__.replace('_vjp', '')
        G.add_node(out_id, label=f"v{out_id} ({op_name})")
        for in_id in in_ids:
            G.add_edge(in_id, out_id)
    return G

# Tracer f pour remplir la bande
grad(f)(0.5, 3.0)
G_f = tape_to_graph(_tape, ["x", "y"])
print(f"Graphe de f: {G_f.number_of_nodes()} noeuds, {G_f.number_of_edges()} arêtes")

# Tracer g pour remplir la bande
grad(g)(1.0, 2.0)
G_g = tape_to_graph(_tape, ["x", "y"])
print(f"Graphe de g: {G_g.number_of_nodes()} noeuds, {G_g.number_of_edges()} arêtes")

# Visualisation côte à côte
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

for ax, G, title in [(axes[0], G_f, r"$f(x,y) = \sin(xy) + e^x$"),
                       (axes[1], G_g, r"$g(x,y) = -(\sin(x) \cdot \sin(y))$")]:
    labels = nx.get_node_attributes(G, 'label')
    # Disposition en couches par profondeur topologique
    for layer, nodes in enumerate(nx.topological_generations(G)):
        for node in nodes:
            G.nodes[node]["layer"] = layer
    pos = nx.multipartite_layout(G, subset_key="layer")
    nx.draw(G, pos, ax=ax, with_labels=True, labels=labels,
            node_color="lightblue", node_size=2500,
            font_size=9, font_weight="bold",
            arrowsize=20, edge_color="gray")
    ax.set_title(title, fontsize=13)

plt.tight_layout()
plt.show()

Comparaison: construction manuelle vs traceur¶

	Construction manuelle (NetworkX)	Traceur (`Var` + `grad`)
Graphe	Construit à la main (`add_node`, `add_edge`)	Construit automatiquement par la trace
Fonction	Écrite séparément du graphe	Écrite une seule fois, normalement
Modification	Refaire le graphe à chaque changement	Rien à changer: la trace s’adapte
Contrôle de flux	Impossible (le graphe est statique)	Boucles et conditions gérées naturellement

L’astuce d’importation¶

Le mécanisme du traceur explique pourquoi, dans JAX, on écrit import jax.numpy as jnp plutôt que import numpy as np. Les fonctions jnp.sin, jnp.exp, etc. savent manipuler les traceurs de JAX, tandis que les fonctions NumPy ordinaires ne les reconnaissent pas. Dans notre implémentation, les méthodes sin() et exp() jouent le même rôle que jnp.sin et jnp.exp: elles interceptent l’appel, enregistrent l’opération sur la bande, et retournent un nouveau traceur.

Partie 6: L’approche par espace de noms¶

Les parties précédentes utilisaient la surcharge d’opérateurs: sin et exp étaient des méthodes de Var, et la fonction à dériver s’écrivait (x * y).sin() + x.exp(). Cette syntaxe est explicite, mais elle ne ressemble pas au code NumPy habituel.

Les bibliothèques de DA réelles (autograd, JAX) adoptent une stratégie différente: elles fournissent un espace de noms qui réimplémente les fonctions NumPy. Quand l’argument est un traceur, la fonction enregistre l’opération sur la bande; quand c’est un nombre ordinaire, elle délègue à la bibliothèque standard. L’utilisateur écrit alors du code NumPy normal, en remplaçant np par l’espace de noms différentiable:

import autograd.numpy as anp   # autograd
import jax.numpy as jnp        # JAX

Construisons un tel espace de noms pour notre traceur.

class tracer_numpy:
    """Espace de noms qui réimplémente les fonctions math pour les traceurs.

    Analogue simplifié de autograd.numpy ou jax.numpy.
    Si l'argument est un Var, l'opération est enregistrée sur la bande.
    Sinon, la fonction standard est appelée.
    """

    @staticmethod
    def sin(x):
        if isinstance(x, Var):
            return x._record(sin_vjp, (x.data,), [x], math.sin(x.data))
        return math.sin(x)

    @staticmethod
    def exp(x):
        if isinstance(x, Var):
            return x._record(exp_vjp, (x.data,), [x], math.exp(x.data))
        return math.exp(x)

    @staticmethod
    def cos(x):
        # Utile pour les vérifications analytiques, pas besoin de VJP ici
        if isinstance(x, Var):
            raise NotImplementedError("cos n'est pas encore tracé")
        return math.cos(x)

# Alias court, comme jnp pour jax.numpy
tnp = tracer_numpy()

La classe Var garde ses opérateurs arithmétiques (+, *, -) car ceux-ci sont interceptés par Python via les méthodes spéciales (__add__, __mul__, etc.). En revanche, les fonctions comme sin et exp passent par l’espace de noms. La fonction à dériver s’écrit maintenant comme du code NumPy ordinaire:

# Avec la surcharge d'opérateurs (parties 2-4):
#   def f(x, y): return (x * y).sin() + x.exp()
#
# Avec l'espace de noms:
def f_ns(x, y):
    return tnp.sin(x * y) + tnp.exp(x)

def g_ns(x, y):
    return -(tnp.sin(x) * tnp.sin(y))

# Vérification: les gradients sont identiques
df_dx, df_dy = grad(f_ns)(0.5, 3.0)
df_dx_exact = 3.0 * math.cos(0.5 * 3.0) + math.exp(0.5)
df_dy_exact = 0.5 * math.cos(0.5 * 3.0)

print("--- f(x,y) = sin(xy) + exp(x), via espace de noms ---")
print(f"∂f/∂x = {df_dx:.6f}  (exact: {df_dx_exact:.6f})")
print(f"∂f/∂y = {df_dy:.6f}  (exact: {df_dy_exact:.6f})")
assert np.allclose(df_dx, df_dx_exact) and np.allclose(df_dy, df_dy_exact)

dg_dx, dg_dy = grad(g_ns)(1.0, 2.0)
dg_dx_exact = -math.cos(1.0) * math.sin(2.0)
dg_dy_exact = -math.sin(1.0) * math.cos(2.0)

print(f"\n--- g(x,y) = -(sin(x)·sin(y)), via espace de noms ---")
print(f"∂g/∂x = {dg_dx:.6f}  (exact: {dg_dx_exact:.6f})")
print(f"∂g/∂y = {dg_dy:.6f}  (exact: {dg_dy_exact:.6f})")
assert np.allclose(dg_dx, dg_dx_exact) and np.allclose(dg_dy, dg_dy_exact)
print("\nTous les gradients correspondent.")

La bande, la fonction grad, et les règles VJP sont exactement les mêmes dans les deux approches. Seul le mécanisme d’interception change:

	Surcharge d’opérateurs	Espace de noms
Interception	Méthodes spéciales Python (`__add__`, `sin`)	Fonctions dans un module (`tnp.sin`)
Syntaxe	`(x * y).sin()`	`tnp.sin(x * y)`
Ajout d’une primitive	Ajouter une méthode à `Var`	Ajouter une fonction à `tracer_numpy`
Ressemble à NumPy	Non	Oui
Bibliothèques	PyTorch (en partie)	autograd, JAX

En pratique, les deux mécanismes coexistent. PyTorch utilise la surcharge d’opérateurs pour +, *, etc., et des fonctions de l’espace de noms torch pour torch.sin, torch.exp, etc. JAX fait de même: les opérateurs arithmétiques passent par __add__/__mul__ sur les tableaux JAX, tandis que les fonctions comme sin et exp passent par jnp.

Bonus¶

Exercice 4 ★: Soustraction par composition¶

Ajoutez __sub__ et __rsub__ à la classe Var sans définir de nouvelle règle VJP. Réutilisez __add__ et __neg__ (puisque $a - b = a + (-b)$ ).

Testez sur $h(x,y) = x - y$ avec les valeurs $x = 3$ , $y = 1$ .

# ============================================
# Exercice 4 ★ (bonus): Ajouter __sub__ et __rsub__
#
# Rappel: a - b = a + (-b) et la négation est déjà implémentée.
# Pensez à gérer le cas où other n'est pas un Var.
# ============================================

# Décommentez et complétez:
# Var.__sub__ = lambda self, other: ...
# Var.__rsub__ = lambda self, other: ...

# Test (décommentez après avoir complété):
# dh_dx, dh_dy = grad(lambda x, y: x - y)(3.0, 1.0)
# print(f"∂h/∂x = {dh_dx} (attendu: 1.0)")
# print(f"∂h/∂y = {dh_dy} (attendu: -1.0)")

def _sub(self, other):
    other = other if isinstance(other, Var) else Var(other)
    return self + (-other)

def _rsub(self, other):
    other = other if isinstance(other, Var) else Var(other)
    return other + (-self)

Var.__sub__ = _sub
Var.__rsub__ = _rsub

dh_dx, dh_dy = grad(lambda x, y: x - y)(3.0, 1.0)
print(f"∂h/∂x = {dh_dx} (attendu: 1.0)")
print(f"∂h/∂y = {dh_dy} (attendu: -1.0)")

La soustraction ne nécessite pas de nouvelle règle VJP: elle se décompose en une négation suivie d’une addition, et la rétropropagation traverse les deux opérations correctement.