grx-tensor 0.1.0 → 0.2.0

data/lib/grx/c_api.rb

@@ -7,49 +7,53 @@ module GRX
   module CAPI
     extend Fiddle::Importer
 
-    LIB_NAME = case RUBY_PLATFORM
-               when /mingw|mswin|windows/i then "grx_core.dll"
-               when /darwin/i              then "libgrx_core.dylib"
-               else                             "libgrx_core.so"
-               end
+    CANDIDATE_NAMES = case RUBY_PLATFORM
+                      when /mingw|mswin|windows/i
+                        ["grx_core.dll", "libgrx_core.dll", "libgrx_core.so", "grx_core.so"]
+                      when /darwin/i
+                        ["libgrx_core.dylib", "grx_core.bundle", "libgrx_core.so", "grx_core.so"]
+                      else
+                        ["libgrx_core.so", "grx_core.so", "libgrx_core.dylib", "grx_core.dll"]
+                      end
 
-    # rake-compiler siempre genera el archivo como "grx_core.so" / "grx_core.bundle" / "grx_core.dll"
-    # (sin el prefijo "lib"), y lo pone un nivel arriba de lib/grx/
-    RAKE_COMPILER_NAME = case RUBY_PLATFORM
-                         when /mingw|mswin|windows/i then "grx_core.dll"
-                         when /darwin/i              then "grx_core.bundle"
-                         else                             "grx_core.so"
-                         end
-
-    LIB_PATHS = [
-      # 1. make -C ext/unix  →  lib/grx/libgrx_core.so
-      File.expand_path(LIB_NAME, __dir__),
-      # 2. gem install (rake-compiler)  →  lib/grx_core.so  (un nivel arriba)
-      File.expand_path("../#{RAKE_COMPILER_NAME}", __dir__),
-      # 3. gem install en Ruby versioned path  →  lib/ruby/X.X.X/grx_core.so
-      File.expand_path("../../#{RAKE_COMPILER_NAME}", __dir__),
-      # 4. desarrollo local sin instalar
-      File.expand_path("../../ext/grx/#{LIB_NAME}", __dir__),
+    SEARCH_DIRS = [
+      File.expand_path(__dir__),                     # lib/grx/
+      File.expand_path("..", __dir__),              # lib/
+      File.expand_path("../../ext/grx", __dir__),    # ext/grx/
+      File.expand_path("../../ext/unix", __dir__),   # ext/unix/
+      File.expand_path("../../ext/windows", __dir__) # ext/windows/
     ].freeze
 
+    LIB_PATHS = SEARCH_DIRS.flat_map do |dir|
+      CANDIDATE_NAMES.flat_map do |name|
+        [
+          File.join(dir, name),
+          File.join(dir, name).tr("/", "\\")
+        ]
+      end
+    end.uniq.freeze
+
     LOADED = begin
-      path = LIB_PATHS.find { |p| File.exist?(p) }
-      raise Fiddle::DLError, "No se encontró #{LIB_NAME} en #{LIB_PATHS.inspect}" unless path
-      dlload path
-      true
+      path = LIB_PATHS.find { |p| File.file?(p) && File.exist?(p) }
+      if path
+        dlload path
+        true
+      else
+        raise Fiddle::DLError, "Binary library not found (#{CANDIDATE_NAMES.join(', ')}) in #{SEARCH_DIRS.inspect}"
+      end
     rescue Fiddle::DLError => e
-      warn "[GRX] Extensión C no disponible: #{e.message}\n" \
-           "      → Ejecuta: make -C ext/unix install\n" \
-           "      → Corriendo en modo Ruby puro (sin SIMD)."
+      warn "[GRX] C extension unavailable: #{e.message}\n" \
+           "      → Run: make -C ext/unix all (Linux/macOS) or make -C ext/windows -f Makefile.mingw all (Windows)\n" \
+           "      → Running in pure Ruby fallback mode (without SIMD)."
       false
     end
 
     if LOADED
-      # Memoria
+      # Memory management
       extern "double* grx_alloc(size_t)"
       extern "void    grx_free(double*)"
 
-      # Aritmética element-wise
+      # Element-wise arithmetic
       extern "void grx_add       (double*, double*, double*, size_t)"
       extern "void grx_sub       (double*, double*, double*, size_t)"
       extern "void grx_mul       (double*, double*, double*, size_t)"
@@ -58,7 +62,7 @@ module GRX
       extern "void grx_add_scalar(double*, double,  double*, size_t)"
       extern "void grx_negate    (double*, double*, size_t)"
 
-      # Matemáticas element-wise
+      # Element-wise math
       extern "void grx_abs   (double*, double*, size_t)"
       extern "void grx_sqrt  (double*, double*, size_t)"
       extern "void grx_square(double*, double*, size_t)"
@@ -67,28 +71,28 @@ module GRX
       extern "void grx_pow   (double*, double,  double*, size_t)"
       extern "void grx_clip  (double*, double, double, double*, size_t)"
 
-      # Reducciones
+      # Reductions
       extern "double grx_sum (double*, size_t)"
       extern "double grx_mean(double*, size_t)"
       extern "double grx_max (double*, size_t)"
       extern "double grx_min (double*, size_t)"
 
-      # Álgebra lineal
+      # Linear algebra
       extern "double grx_dot    (double*, double*, size_t)"
       extern "void   grx_matmul (double*, double*, double*, size_t, size_t, size_t)"
 
-      # Activaciones
+      # Activations
       extern "void grx_relu       (double*, double*, size_t)"
       extern "void grx_leaky_relu (double*, double, double*, size_t)"
       extern "void grx_tanh_act   (double*, double*, size_t)"
       extern "void grx_sigmoid    (double*, double*, size_t)"
       extern "void grx_softmax    (double*, double*, size_t)"
 
-      # Optimizadores
+      # Optimizers
       extern "void grx_sgd_step (double*, double*, double, size_t)"
       extern "void grx_adam_step(double*, double*, double*, double*, double, double, double, double, double, double, size_t)"
 
-      # Inicialización de pesos
+      # Weight initialization
       extern "void grx_init_xavier_uniform(double*, size_t, size_t, size_t)"
       extern "void grx_init_he_normal     (double*, size_t, size_t)"
     end

data/lib/grx/data.rb

@@ -0,0 +1,87 @@
+# frozen_string_literal: true
+
+module GRX
+  module Data
+    # ================================================================
+    # Dataset — Base class for data collections
+    # ================================================================
+    class Dataset
+      def size
+        raise NotImplementedError, "#{self.class}#size must be implemented"
+      end
+
+      def [](index)
+        raise NotImplementedError, "#{self.class}#[] must be implemented"
+      end
+    end
+
+    # ================================================================
+    # TensorDataset — Dataset wrapping parallel tensors (e.g. X and Y)
+    # ================================================================
+    class TensorDataset < Dataset
+      attr_reader :tensors, :size
+
+      def initialize(*tensors)
+        raise ArgumentError, "Must provide at least one tensor" if tensors.empty?
+        first_dim = tensors.first.shape[0]
+        unless tensors.all? { |t| t.shape[0] == first_dim }
+          raise ArgumentError, "All tensors must have the same size in batch dimension (dimension 0)"
+        end
+        @tensors = tensors
+        @size = first_dim
+      end
+
+      def [](index)
+        @tensors.map do |t|
+          cols = t.numel / @size
+          offset = index * cols
+          data = t.to_a.slice(offset, cols)
+          new_shape = t.shape.size == 1 ? [1] : [1] + t.shape[1..]
+          Tensor.create(data, new_shape)
+        end
+      end
+    end
+
+    # ================================================================
+    # DataLoader — Mini-batch iterator with optional shuffling
+    # ================================================================
+    class DataLoader
+      include Enumerable
+
+      attr_reader :dataset, :batch_size, :shuffle
+
+      def initialize(dataset, batch_size: 32, shuffle: true)
+        @dataset    = dataset
+        @batch_size = batch_size
+        @shuffle    = shuffle
+      end
+
+      def each
+        return to_enum(:each) unless block_given?
+
+        indices = (0...@dataset.size).to_a
+        indices.shuffle! if @shuffle
+
+        indices.each_slice(@batch_size) do |batch_indices|
+          batch_samples = batch_indices.map { |i| @dataset[i] }
+          num_tensors = batch_samples.first.size
+
+          batched = (0...num_tensors).map do |t_idx|
+            slices = batch_samples.map { |sample| sample[t_idx].to_a }
+            flat_data = slices.flatten
+            sample_shape = batch_samples.first[t_idx].shape
+            batch_dim = batch_indices.size
+            target_shape = [batch_dim] + sample_shape[1..]
+            Tensor.create(flat_data, target_shape)
+          end
+
+          yield(*batched)
+        end
+      end
+
+      def size
+        (@dataset.size.to_f / @batch_size).ceil
+      end
+    end
+  end
+end

data/lib/grx/loss.rb

@@ -4,11 +4,11 @@ module GRX
   module Loss
     # ================================================================
     # MSELoss — Mean Squared Error
-    # L = mean((pred - target)^2)   → retorna Float
+    # L = mean((pred - target)^2)   -> returns differentiable scalar Tensor
     # ================================================================
     class MSELoss
       def call(pred, target)
-        raise ShapeError, "Shapes incompatibles" if pred.shape != target.shape
+        raise ShapeError, "Incompatible shapes: #{pred.shape} vs #{target.shape}" if pred.shape != target.shape
         (pred - target).square.mean
       end
     end
@@ -19,7 +19,7 @@ module GRX
     # ================================================================
     class MAELoss
       def call(pred, target)
-        raise ShapeError, "Shapes incompatibles" if pred.shape != target.shape
+        raise ShapeError, "Incompatible shapes: #{pred.shape} vs #{target.shape}" if pred.shape != target.shape
         (pred - target).abs.mean
       end
     end
@@ -27,54 +27,84 @@ module GRX
     # ================================================================
     # BCELoss — Binary Cross-Entropy
     # L = -mean(t*log(p) + (1-t)*log(1-p))
-    # pred debe estar en (0,1) — aplica sigmoid antes si usas logits.
+    # pred must be in (0,1) — apply Sigmoid before if using logits.
     # ================================================================
     class BCELoss
       EPS = 1e-7
 
       def call(pred, target)
-        raise ShapeError, "Shapes incompatibles" if pred.shape != target.shape
-        p_data = pred.to_a.map { |v| v < EPS ? EPS : (v > 1-EPS ? 1-EPS : v) }
-        t_data = target.to_a
-        total  = p_data.size.to_f
-        loss   = p_data.each_with_index.sum do |p, i|
-          t = t_data[i]
-          -(t * Math.log(p) + (1 - t) * Math.log(1 - p))
-        end
-        loss / total
+        raise ShapeError, "Incompatible shapes: #{pred.shape} vs #{target.shape}" if pred.shape != target.shape
+        p_clamped = pred.clip(EPS, 1.0 - EPS)
+        ones = Tensor.ones_like(target)
+        term1 = target * p_clamped.log
+        term2 = (ones - target) * (ones - p_clamped).log
+        (-(term1 + term2)).mean
       end
     end
 
     # ================================================================
-    # CrossEntropyLoss — Softmax + NLL
-    # L = -mean(sum(target * log(softmax(logits))))
+    # CrossEntropyLoss — Softmax + NLL (multi-class)
+    # L = -sum(target * log(softmax(logits))) / batch_size
     # ================================================================
     class CrossEntropyLoss
       EPS = 1e-7
 
       def call(logits, target)
-        raise ShapeError, "Shapes incompatibles" if logits.shape != target.shape
-        probs  = logits.softmax.to_a.map { |v| v < EPS ? EPS : v }
+        raise ShapeError, "Incompatible shapes: #{logits.shape} vs #{target.shape}" if logits.shape != target.shape
+        probs = logits.softmax
+        p_data = probs.to_a
         t_data = target.to_a
-        loss   = probs.each_with_index.sum { |p, i| -t_data[i] * Math.log(p) }
-        loss / probs.size.to_f
+        batch_size = logits.shape[0].to_f
+
+        loss_val = t_data.each_with_index.sum do |t, i|
+          next 0.0 if t == 0.0
+          p = [p_data[i], EPS].max
+          -t * Math.log(p)
+        end / batch_size
+
+        out = Tensor.create([loss_val], [1], requires_grad: logits.requires_grad || target.requires_grad)
+        if logits.requires_grad || target.requires_grad
+          out._grafo_hijos.push(logits, target)
+          out.backward_fn = ->(g) {
+            scale = g.item / batch_size
+            grad_logits = p_data.zip(t_data).map { |p, t| (p - t) * scale }
+            logits.agregar_gradiente(Tensor.create(grad_logits, logits.shape)) if logits.requires_grad
+          }
+        end
+        out
       end
     end
 
     # ================================================================
-    # HuberLoss — Smooth L1
+    # HuberLoss — Smooth L1 (robust against outliers)
     # ================================================================
     class HuberLoss
+      attr_reader :delta
+
       def initialize(delta: 1.0)
-        @delta = delta
+        @delta = delta.to_f
       end
 
       def call(pred, target)
-        raise ShapeError, "Shapes incompatibles" if pred.shape != target.shape
+        raise ShapeError, "Shapes incompatibles: #{pred.shape} vs #{target.shape}" if pred.shape != target.shape
+        diff_data = (pred - target).abs.to_a
         d = @delta
-        diffs = (pred - target).abs.to_a
-        loss  = diffs.sum { |v| v <= d ? 0.5 * v * v : d * (v - 0.5 * d) }
-        loss / diffs.size.to_f
+        loss_val = diff_data.sum { |v| v <= d ? 0.5 * v * v : d * (v - 0.5 * d) } / diff_data.size.to_f
+
+        out = Tensor.create([loss_val], [1], requires_grad: pred.requires_grad || target.requires_grad)
+        if pred.requires_grad || target.requires_grad
+          out._grafo_hijos.push(pred, target)
+          n = diff_data.size.to_f
+          out.backward_fn = ->(g) {
+            grad_pred = (pred - target).to_a.map do |err|
+              abs_err = err.abs
+              (abs_err <= d ? err : d * (err > 0 ? 1.0 : -1.0)) * (g.item / n)
+            end
+            pred.agregar_gradiente(Tensor.create(grad_pred, pred.shape)) if pred.requires_grad
+            target.agregar_gradiente(Tensor.create(grad_pred.map(&:-@), target.shape)) if target.requires_grad
+          }
+        end
+        out
       end
     end
   end

data/lib/grx/nn.rb

@@ -3,10 +3,10 @@
 module GRX
   module NN
     # ================================================================
-    # Module — clase base para todas las capas
+    # Module — Base class for all neural network layers
     # ================================================================
     class Module
-      # Retorna todos los parámetros entrenables (para pasarlos al optimizador)
+      # Returns all trainable parameters for optimizer registration
       def parameters
         instance_variables.flat_map do |var|
           val = instance_variable_get(var)
@@ -29,14 +29,22 @@ module GRX
         parameters.each(&:zero_grad!)
       end
 
-      # Subclases implementan forward
+      def save_weights(path)
+        GRX::Serialization.save(self, path)
+      end
+
+      def load_weights(path)
+        GRX::Serialization.load(self, path)
+      end
+
+      # Subclasses implement forward computation
       def call(*args)
         forward(*args)
       end
     end
 
     # ================================================================
-    # Linear — Capa densa (fully connected)
+    # Linear — Dense fully connected layer
     # y = x @ W^T + b
     # ================================================================
     class Linear < Module
@@ -47,25 +55,20 @@ module GRX
         @out_features = out_features
         @use_bias     = bias
 
-        # Pesos: Xavier uniform (bueno para tanh/sigmoid)
+        # Weights: Xavier uniform initialization
         @weight = Tensor.xavier_uniform([out_features, in_features], requires_grad: true)
 
-        # Bias: ceros
+        # Bias: initialized to zeros
         @bias = bias ? Tensor.zeros([out_features], requires_grad: true) : nil
       end
 
       def forward(x)
-        # x: [batch, in_features]  →  out: [batch, out_features]
+        # x: [batch, in_features]  ->  out: [batch, out_features]
         # out = x @ W^T
         out = x.matmul(@weight.transpose)
 
         if @use_bias
-          # Sumamos bias fila por fila.
-          # Repetimos @bias batch_size veces para crear un tensor [batch, out_features]
-          # que comparte el grafo con @bias original.
           batch_size = x.shape[0]
-          # Tile del bias: concatenamos el mismo tensor bias_size veces
-          # usando operaciones que mantienen el grafo conectado
           bias_tiled = _tile_bias(@bias, batch_size, @out_features)
           out + bias_tiled
         else
@@ -75,17 +78,11 @@ module GRX
 
       private
 
-      # Crea un tensor [batch, out_features] repitiendo @bias batch veces.
-      # Usa add_scalar(0) para crear un nuevo nodo conectado al bias en el grafo.
       def _tile_bias(bias, batch_size, out_features)
-        # Construimos el tensor tileado sumando el bias a un tensor de ceros
-        # de la forma correcta. Esto conecta el grafo al bias original.
         data = Array.new(batch_size) { bias.to_a }.flatten
         tiled = GRX::Tensor.create(data, [batch_size, out_features])
-        # Conectamos al bias original via suma con ceros — mantiene el grafo
         zero_row = GRX::Tensor.zeros([batch_size, out_features])
         result = zero_row + tiled
-        # Registramos manualmente la conexión al bias para backprop
         if bias.requires_grad
           result.requires_grad = true
           result._grafo_hijos << bias
@@ -93,8 +90,7 @@ module GRX
           bf = result.backward_fn
           result.backward_fn = ->(g) {
             bf&.call(g)
-            # Acumulamos gradiente en bias: suma sobre el batch
-            grad_data = g.to_a.each_slice(out_features).reduce([0.0]*out_features) { |acc, row|
+            grad_data = g.to_a.each_slice(out_features).reduce([0.0] * out_features) { |acc, row|
               acc.zip(row).map { |a, r| a + r }
             }
             b.agregar_gradiente(GRX::Tensor.create(grad_data, [out_features]))
@@ -106,12 +102,12 @@ module GRX
       public
 
       def to_s
-        "Linear(#{@in_features} → #{@out_features}, bias: #{@use_bias})"
+        "Linear(#{@in_features} -> #{@out_features}, bias: #{@use_bias})"
       end
     end
 
     # ================================================================
-    # Sequential — Contenedor de capas en secuencia
+    # Sequential — Chains layers sequentially in order
     # ================================================================
     class Sequential < Module
       def initialize(*layers)
@@ -126,6 +122,16 @@ module GRX
         @layers.flat_map(&:parameters)
       end
 
+      def train!
+        @layers.each { |l| l.train! if l.respond_to?(:train!) }
+        self
+      end
+
+      def eval!
+        @layers.each { |l| l.eval! if l.respond_to?(:eval!) }
+        self
+      end
+
       def to_s
         layers_str = @layers.each_with_index.map { |l, i| "  (#{i}): #{l}" }.join("\n")
         "Sequential(\n#{layers_str}\n)"
@@ -133,7 +139,7 @@ module GRX
     end
 
     # ================================================================
-    # Activaciones como capas (para usar en Sequential)
+    # Activation layers (for use inside Sequential pipelines)
     # ================================================================
     class ReLU < Module
       def forward(x) = x.relu
@@ -164,11 +170,11 @@ module GRX
     end
 
     # ================================================================
-    # Dropout — regularización durante entrenamiento
+    # Dropout — Random feature dropout during training mode
     # ================================================================
     class Dropout < Module
       def initialize(p = 0.5)
-        @p       = p
+        @p        = p
         @training = true
       end
 
@@ -178,8 +184,6 @@ module GRX
       def forward(x)
         return x unless @training && @p > 0
 
-        # Máscara binaria: 1 con prob (1-p), 0 con prob p
-        # Escalamos por 1/(1-p) para mantener la esperanza (inverted dropout)
         scale = 1.0 / (1.0 - @p)
         mask_data = x.to_a.map { rand > @p ? scale : 0.0 }
         mask = Tensor.create(mask_data, x.shape)
@@ -190,8 +194,128 @@ module GRX
     end
 
     # ================================================================
-    # BatchNorm1d — Normalización por batch
-    # Estabiliza el entrenamiento de redes profundas.
+    # Embedding — Dense vector lookup table for token indices
+    # ================================================================
+    class Embedding < Module
+      attr_reader :weight, :num_embeddings, :embedding_dim
+
+      def initialize(num_embeddings, embedding_dim)
+        @num_embeddings = num_embeddings
+        @embedding_dim  = embedding_dim
+        @weight         = Tensor.he_normal([num_embeddings, embedding_dim], requires_grad: true)
+      end
+
+      def forward(indices)
+        ids = indices.is_a?(Tensor) ? indices.to_a.map(&:to_i) : indices.map(&:to_i)
+        batch_size = ids.size
+        out_data = ids.flat_map do |id|
+          raise IndexError, "Token index #{id} out of range [0, #{@num_embeddings})" if id < 0 || id >= @num_embeddings
+          @weight.to_a.slice(id * @embedding_dim, @embedding_dim)
+        end
+
+        out = Tensor.create(out_data, [batch_size, @embedding_dim])
+        if @weight.requires_grad
+          out.requires_grad = true
+          out._grafo_hijos << @weight
+          w = @weight; dim = @embedding_dim; num_emb = @num_embeddings
+          out.backward_fn = ->(g) {
+            grad_w = Array.new(num_emb * dim, 0.0)
+            g_data = g.to_a
+            ids.each_with_index do |id, i|
+              slice = g_data.slice(i * dim, dim)
+              dim.times { |d| grad_w[id * dim + d] += slice[d] }
+            end
+            w.agregar_gradiente(Tensor.create(grad_w, w.shape))
+          }
+        end
+        out
+      end
+
+      def to_s
+        "Embedding(#{@num_embeddings}, #{@embedding_dim})"
+      end
+    end
+
+    # ================================================================
+    # LayerNorm — Layer normalization across channel dimensions
+    # ================================================================
+    class LayerNorm < Module
+      attr_reader :gamma, :beta, :normalized_shape, :epsilon
+
+      def initialize(normalized_shape, epsilon: 1e-5)
+        @normalized_shape = normalized_shape.is_a?(Array) ? normalized_shape : [normalized_shape]
+        @dim              = @normalized_shape.reduce(1, :*)
+        @epsilon          = epsilon
+
+        @gamma = Tensor.ones(@normalized_shape, requires_grad: true)
+        @beta  = Tensor.zeros(@normalized_shape, requires_grad: true)
+      end
+
+      def forward(x)
+        batch_size = x.shape[0]
+        x_data = x.to_a
+
+        means = Array.new(batch_size) do |b|
+          x_data.slice(b * @dim, @dim).sum / @dim.to_f
+        end
+        vars = Array.new(batch_size) do |b|
+          m = means[b]
+          x_data.slice(b * @dim, @dim).sum { |v| (v - m)**2 } / @dim.to_f
+        end
+
+        gamma_data = @gamma.to_a
+        beta_data  = @beta.to_a
+        norm_data  = Array.new(batch_size * @dim)
+
+        batch_size.times do |b|
+          m = means[b]
+          inv_std = 1.0 / Math.sqrt(vars[b] + @epsilon)
+          @dim.times do |j|
+            norm_data[b * @dim + j] = gamma_data[j] * (x_data[b * @dim + j] - m) * inv_std + beta_data[j]
+          end
+        end
+
+        out = Tensor.create(norm_data, x.shape)
+        if x.requires_grad || @gamma.requires_grad || @beta.requires_grad
+          out.requires_grad = true
+          out._grafo_hijos.push(x, @gamma, @beta)
+          g_param = @gamma; b_param = @beta; d = @dim; eps = @epsilon
+          out.backward_fn = ->(g) {
+            g_data = g.to_a
+            grad_gamma = Array.new(d, 0.0)
+            grad_beta  = Array.new(d, 0.0)
+            grad_x     = Array.new(batch_size * d, 0.0)
+
+            batch_size.times do |b|
+              m = means[b]; v = vars[b]
+              inv_std = 1.0 / Math.sqrt(v + eps)
+              x_hat = Array.new(d) { |j| (x_data[b * d + j] - m) * inv_std }
+              dl_dxhat = Array.new(d) { |j| g_data[b * d + j] * gamma_data[j] }
+              sum_dl = dl_dxhat.sum
+              sum_dl_x = dl_dxhat.zip(x_hat).sum { |a, c| a * c }
+
+              d.times do |j|
+                grad_gamma[j] += g_data[b * d + j] * x_hat[j]
+                grad_beta[j]  += g_data[b * d + j]
+                grad_x[b * d + j] = (inv_std / d.to_f) * (d.to_f * dl_dxhat[j] - sum_dl - x_hat[j] * sum_dl_x)
+              end
+            end
+
+            g_param.agregar_gradiente(Tensor.create(grad_gamma, g_param.shape)) if g_param.requires_grad
+            b_param.agregar_gradiente(Tensor.create(grad_beta, b_param.shape))   if b_param.requires_grad
+            x.agregar_gradiente(Tensor.create(grad_x, x.shape))                 if x.requires_grad
+          }
+        end
+        out
+      end
+
+      def to_s
+        "LayerNorm(#{@normalized_shape})"
+      end
+    end
+
+    # ================================================================
+    # BatchNorm1d — Normalizacion por batch
     # ================================================================
     class BatchNorm1d < Module
       def initialize(num_features, epsilon: 1e-5, momentum: 0.1)
@@ -200,11 +324,9 @@ module GRX
         @momentum     = momentum
         @training     = true
 
-        # Parámetros entrenables
         @gamma = Tensor.ones([num_features],  requires_grad: true)
         @beta  = Tensor.zeros([num_features], requires_grad: true)
 
-        # Estadísticas corrientes (no entrenables, para inferencia)
         @running_mean = Tensor.zeros([num_features])
         @running_var  = Tensor.ones([num_features])
       end
@@ -213,11 +335,9 @@ module GRX
       def eval!;  @training = false; self; end
 
       def forward(x)
-        # x: [batch, num_features]
         batch_size = x.shape[0]
 
         if @training
-          # Calculamos media y varianza del batch
           batch_data = x.to_a
           means = Array.new(@num_features) do |j|
             batch_data.each_slice(@num_features).map { |row| row[j] }.sum / batch_size
@@ -227,7 +347,6 @@ module GRX
             col.sum { |v| (v - means[j]) ** 2 } / batch_size
           end
 
-          # Actualizamos estadísticas corrientes
           means.each_with_index do |m, j|
             rm = @running_mean.to_a; rm[j] = (1 - @momentum) * rm[j] + @momentum * m
             @running_mean = Tensor.create(rm, [@num_features])
@@ -244,8 +363,6 @@ module GRX
           var_t  = @running_var
         end
 
-        # Normalizamos: x_hat = (x - mean) / sqrt(var + eps)
-        # Luego escalamos: y = gamma * x_hat + beta
         norm_data = x.to_a.each_slice(@num_features).flat_map do |row|
           row.each_with_index.map do |v, j|
             x_hat = (v - mean_t.to_a[j]) / Math.sqrt(var_t.to_a[j] + @epsilon)