grx-tensor 0.1.0 → 0.2.1

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
@@ -141,8 +147,10 @@ module GRX
     end
 
     class LeakyReLU < Module
-      def initialize(alpha = 0.01)
-        @alpha = alpha
+      attr_reader :alpha
+
+      def initialize(alpha_arg = nil, alpha: nil)
+        @alpha = (alpha || alpha_arg || 0.01).to_f
       end
       def forward(x) = x.leaky_relu(@alpha)
       def to_s = "LeakyReLU(alpha=#{@alpha})"
@@ -164,11 +172,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 +186,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,21 +196,139 @@ 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, eps: nil, epsilon: 1e-5)
+        @normalized_shape = normalized_shape.is_a?(Array) ? normalized_shape : [normalized_shape]
+        @dim              = @normalized_shape.reduce(1, :*)
+        @epsilon          = (eps || epsilon).to_f
+
+        @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)
+      def initialize(num_features, eps: nil, epsilon: 1e-5, momentum: 0.1)
         @num_features = num_features
-        @epsilon      = epsilon
-        @momentum     = momentum
+        @epsilon      = (eps || epsilon).to_f
+        @momentum     = momentum.to_f
         @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 +337,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 +349,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 +365,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)

data/lib/grx/optim.rb

@@ -3,7 +3,7 @@
 module GRX
   module Optim
     # ================================================================
-    # SGD — Stochastic Gradient Descent (con momentum opcional)
+    # SGD — Stochastic Gradient Descent (with optional momentum)
     # ================================================================
     class SGD
       def initialize(params, lr: 0.01, momentum: 0.0, weight_decay: 0.0)
@@ -11,7 +11,7 @@ module GRX
         @lr           = lr
         @momentum     = momentum
         @weight_decay = weight_decay
-        # Buffer de velocidad para momentum
+        # Velocity buffer for momentum
         @velocity = params.map { |p| Tensor.zeros_like(p) }
       end
 
@@ -21,7 +21,7 @@ module GRX
 
           grad = param.grad
 
-          # L2 regularización (weight decay)
+          # L2 regularization (weight decay)
           if @weight_decay > 0
             grad = grad + param.scale(@weight_decay)
           end
@@ -35,7 +35,7 @@ module GRX
           if CAPI::LOADED
             CAPI.grx_sgd_step(param.storage.ptr, grad.storage.ptr, @lr, param.numel)
           else
-            # Fallback Ruby
+            # Ruby fallback
             param_data = param.to_a
             grad_data  = grad.to_a
             param_data.each_with_index { |v, j| param_data[j] = v - @lr * grad_data[j] }
@@ -51,27 +51,32 @@ module GRX
 
     # ================================================================
     # Adam — Adaptive Moment Estimation (Kingma & Ba, 2015)
-    # El optimizador estándar para redes neuronales profundas.
+    # The standard optimizer for deep neural networks.
     # ================================================================
     class Adam
-      def initialize(params, lr: 0.001, beta1: 0.9, beta2: 0.999,
-                     epsilon: 1e-8, weight_decay: 0.0)
+      def initialize(params, lr: 0.001, betas: nil, beta1: 0.9, beta2: 0.999,
+                     eps: nil, epsilon: 1e-8, weight_decay: 0.0)
         @params       = params
-        @lr           = lr
-        @beta1        = beta1
-        @beta2        = beta2
-        @epsilon      = epsilon
-        @weight_decay = weight_decay
-        @t            = 0  # paso actual
+        @lr           = lr.to_f
+        if betas
+          @beta1 = betas[0].to_f
+          @beta2 = betas[1].to_f
+        else
+          @beta1 = beta1.to_f
+          @beta2 = beta2.to_f
+        end
+        @epsilon      = eps ? eps.to_f : epsilon.to_f
+        @weight_decay = weight_decay.to_f
+        @t            = 0  # current step
 
-        # Momentos de primer y segundo orden (inicializados en cero)
+        # First and second order moment vectors (zero-initialized)
         @m = params.map { |p| Tensor.zeros_like(p) }
         @v = params.map { |p| Tensor.zeros_like(p) }
       end
 
       def step
         @t += 1
-        beta1t = @beta1 ** @t  # beta1^t para corrección de bias
+        beta1t = @beta1 ** @t  # beta1^t for bias correction
         beta2t = @beta2 ** @t
 
         @params.each_with_index do |param, i|
@@ -94,7 +99,7 @@ module GRX
               param.numel
             )
           else
-            # Fallback Ruby puro
+            # Pure Ruby fallback
             p_data = param.to_a
             m_data = @m[i].to_a
             v_data = @v[i].to_a

data/lib/grx/serialization.rb

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+module GRX
+  # ===================================================================
+  # Serialization — Native binary .grx format
+  #
+  # Binary layout:
+  # - 8 bytes:  Magic header "GRX1\0\0\0\0"
+  # - 4 bytes:  Number of parameter tensors (unsigned 32-bit big-endian)
+  # For each tensor:
+  #   - 2 bytes: Rank (number of dimensions)
+  #   - 4 bytes * rank: Dimensions of the shape (uint32 big-endian)
+  #   - 8 bytes: Total numel (uint64 big-endian)
+  #   - numel * 8 bytes: Direct binary packed IEEE 754 doubles
+  # ===================================================================
+  module Serialization
+    MAGIC = "GRX1\x00\x00\x00\x00".b
+
+    def self.save(model, path)
+      params = model.parameters
+      File.open(path, "wb") do |f|
+        f.write(MAGIC)
+        f.write([params.size].pack("N"))
+        params.each do |p|
+          shape = p.shape
+          f.write([shape.size].pack("n"))
+          f.write(shape.pack("N*"))
+          f.write([p.numel].pack("Q>"))
+          # Direct binary copy from native C buffer
+          bytes = if p.storage.ptr
+            p.storage.ptr[0, p.numel * 8]
+          else
+            p.to_a.pack("d*")
+          end
+          f.write(bytes)
+        end
+      end
+      path
+    end
+
+    def self.load(model, path)
+      params = model.parameters
+      File.open(path, "rb") do |f|
+        magic = f.read(8)
+        raise StorageError, "Invalid format: not a valid .grx binary file" unless magic == MAGIC
+        count = f.read(4).unpack1("N")
+        raise StorageError, "Parameter count mismatch: model has #{params.size}, file has #{count}" unless count == params.size
+
+        params.each_with_index do |p, idx|
+          rank = f.read(2).unpack1("n")
+          shape = f.read(rank * 4).unpack("N*")
+          numel = f.read(8).unpack1("Q>")
+          raise ShapeError, "Shape mismatch for parameter #{idx}: expected #{p.shape}, got #{shape}" unless shape == p.shape
+
+          bytes = f.read(numel * 8)
+          if p.storage.ptr
+            p.storage.ptr[0, bytes.bytesize] = bytes
+          else
+            p.storage.instance_variable_set(:@data, bytes.unpack("d*"))
+          end
+        end
+      end
+      model
+    end
+  end
+end