L'Écho du Silicium

import torch
import torch.nn as nn
import torch.nn.functional as F

# L'essence de l'attention (Self-Attention Head)
class Head(nn.Module):
    def __init__(self, head_size, n_embd, block_size):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        # Masque pour ne pas tricher en regardant le futur
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))

    def forward(self, x):
        B, T, C = x.shape
        k = self.key(x) 
        q = self.query(x) 
        
        # Calcul des scores d'attention ("qui regarde qui ?")
        wei = q @ k.transpose(-2, -1) * C**-0.5
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
        wei = F.softmax(wei, dim=-1)
        
        v = self.value(x)
        out = wei @ v
        return out

# Le bloc Transformer standard
class Block(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout):
        super().__init__()
        head_size = n_embd // n_head
        self.sa = MultiHeadAttention(n_head, head_size, n_embd, block_size, dropout)
        self.ffwd = FeedFoward(n_embd, dropout)
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)

    def forward(self, x):
        # Communication (Attention) + Réflexion (FeedForward)
        x = x + self.sa(self.ln1(x))
        x = x + self.ffwd(self.ln2(x))
        return x

class JEPA(nn.Module):
    def __init__(self, input_dim, hidden_dim, latent_dim):
        super().__init__()
        
        # 1. Context Encoder (L'Observateur)
        # Transforme l'entrée brute en concept abstrait
        self.context_encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, latent_dim)
        )
        
        # 2. Target Encoder (La Vérité)
        # Copie l'encodeur principal mais évolue plus lentement (EMA)
        self.target_encoder = copy.deepcopy(self.context_encoder)
        
        # 3. Predictor (Le Cerveau)
        # Devine la représentation future à partir du contexte + action
        self.predictor = nn.Sequential(
            nn.Linear(latent_dim + input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, latent_dim)
        )

    def forward(self, x_context, z_transformation):
        # On encode le contexte
        context_rep = self.context_encoder(x_context)
        
        # On prédit la représentation de la cible (pas les pixels !)
        pred_input = torch.cat([context_rep, z_transformation], dim=1)
        predicted_rep = self.predictor(pred_input)
        
        return predicted_rep

# 1. PYTORCH (Style Classe & Objet)
class PyTorchNeuron(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 1)
    def forward(self, x):
        return self.linear(x)

# 2. JAX (Style Fonctionnel & Mathématique)
def jax_neuron(params, x):
    return jnp.dot(x, params['w']) + params['b']
# Compilation JIT pour la vitesse
fast_neuron = jit(jax_neuron)

# 3. MOJO (Le Futur ?)
# struct MyModel:
#    fn forward(self, x: Tensor) -> Tensor:
#        return x @ self.w + self.b

class SpikingNeuron:
    def __init__(self, seuil=1.0, decay=0.9):
        self.potential = 0.0
        self.seuil = seuil
        self.decay = decay

    def stimulate(self, courant_entrant, temps_ms):
        # Simulation physique : Le potentiel monte, mais fuit avec le temps
        self.potential = self.potential * self.decay + courant_entrant
        
        if self.potential >= self.seuil:
            # Décharge électrique (Spike) et Reset
            self.potential = 0.0 
            print(f"[{temps_ms}ms] *** SPIKE *** -> Action déclenchée")
            return True
        else:
            # Pas de condition arbitraire, juste de la physique
            print(f"[{temps_ms}ms] Voltage: {self.potential:.2f}v")
            return False