grnexus 1.0.2

data/lib/grnexus_layers.rb

@@ -0,0 +1,1103 @@
+require_relative 'grnexus_activations'
+
+module GRNEXUSLayer
+  class Layer
+    def initialize
+      @weights = nil
+      @biases = nil
+      @trainable = true
+      @cache = {}
+    end
+
+    def forward(input)
+      raise NotImplementedError, "Debes implementar el método forward"
+    end
+
+    def backward(gradient, learning_rate)
+      raise NotImplementedError, "Debes implementar el método backward"
+    end
+
+    def trainable?
+      @trainable
+    end
+
+    def parameters
+      [@weights, @biases].compact
+    end
+
+    def zero_gradients!
+      @weight_gradient = nil if @weight_gradient
+      @bias_gradient = nil if @bias_gradient
+    end
+
+    def update_parameters(learning_rate)
+      if @weight_gradient
+        update_matrix!(@weights, @weight_gradient, learning_rate)
+      end
+      if @bias_gradient
+        update_vector!(@biases, @bias_gradient, learning_rate)
+      end
+    end
+
+    private
+
+    def update_matrix!(matrix, gradient, lr)
+      (0...matrix.length).each do |i|
+        (0...matrix[i].length).each do |j|
+          matrix[i][j] -= lr * gradient[i][j]
+        end
+      end
+    end
+
+    def update_vector!(vector, gradient, lr)
+      (0...vector.length).each do |i|
+        vector[i] -= lr * gradient[i]
+      end
+    end
+  end
+
+  class DenseLayer < Layer
+    attr_accessor :units, :input_dim, :activation, :use_bias, :weights, :biases
+    
+    def initialize(units:, input_dim:, activation: nil, use_bias: true, weight_init: :xavier)
+      super()
+      @units = units
+      @input_dim = input_dim
+      @activation = activation.is_a?(Symbol) ? GRNEXUSActivations.const_get(activation).new : activation
+      @use_bias = use_bias
+      
+      @weights = initialize_weights(weight_init, @input_dim, @units)
+      @biases = @use_bias ? Array.new(@units, 0.0) : nil
+    end
+
+    def forward(input)
+      batch_case = input[0].is_a?(Array) && input[0].length > 0 && input[0][0].is_a?(Numeric)
+      
+      if batch_case
+        @cache[:input] = input.dup
+        output_batch = []
+        
+        input.each do |x|
+          # weights is (input_dim x units), x is (input_dim)
+          # result should be (units)
+          z = Array.new(@units, 0.0)
+          @units.times do |i|
+            sum = 0.0
+            @input_dim.times do |j|
+              sum += @weights[j][i] * x[j]
+            end
+            z[i] = sum
+          end
+          
+          z = add_vectors(z, @biases) if @biases
+          
+          if @activation
+            @cache[:pre_activation] ||= []
+            @cache[:pre_activation] << z.dup
+            z = @activation.call(z)
+          end
+          
+          output_batch << z
+        end
+        
+        output_batch
+      else
+        @cache[:input] = input.dup
+        # weights is (input_dim x units), input is (input_dim)
+        # result should be (units)
+        z = Array.new(@units, 0.0)
+        @units.times do |i|
+          sum = 0.0
+          @input_dim.times do |j|
+            sum += @weights[j][i] * input[j]
+          end
+          z[i] = sum
+        end
+        
+        z = add_vectors(z, @biases) if @biases
+        
+        if @activation
+          @cache[:pre_activation] = z.dup
+          z = @activation.call(z)
+        end
+        
+        z
+      end
+    end
+
+    def backward(gradient, learning_rate)
+      batch_case = gradient[0].is_a?(Array) && gradient[0][0].is_a?(Numeric)
+      gradients = batch_case ? gradient : [gradient]
+      inputs = batch_case ? @cache[:input] : [@cache[:input]]
+      
+      input_gradients = []
+      weight_gradients = []
+      bias_gradients = [] if @use_bias
+      
+      gradients.each_with_index do |grad, idx|
+        current_input = inputs[idx]
+        
+        if @activation
+          pre_act = batch_case ? @cache[:pre_activation][idx] : @cache[:pre_activation]
+          act_deriv = @activation.call(pre_act, derivative: true)
+          grad = multiply_elementwise(grad, act_deriv)
+        end
+        
+        # Weight gradient: outer product of input and grad
+        # weights is (input_dim x units), so gradient should be same shape
+        weight_grad = Array.new(@input_dim) { Array.new(@units, 0.0) }
+        @input_dim.times do |i|
+          @units.times do |j|
+            weight_grad[i][j] = current_input[i] * grad[j]
+          end
+        end
+        weight_gradients << weight_grad
+        
+        if @biases
+          bias_gradients << grad.dup
+        end
+        
+        # Input gradient: weights^T * grad
+        input_grad = Array.new(@input_dim, 0.0)
+        @input_dim.times do |i|
+          sum = 0.0
+          @units.times do |j|
+            sum += @weights[i][j] * grad[j]
+          end
+          input_grad[i] = sum
+        end
+        input_gradients << input_grad
+      end
+      
+      # Average gradients and update weights
+      avg_weight_grad = average_matrices(weight_gradients)
+      @input_dim.times do |i|
+        @units.times do |j|
+          @weights[i][j] -= learning_rate * avg_weight_grad[i][j]
+        end
+      end
+      
+      if @biases
+        avg_bias_grad = average_arrays(bias_gradients)
+        @units.times do |i|
+          @biases[i] -= learning_rate * avg_bias_grad[i]
+        end
+      end
+      
+      batch_case ? input_gradients : input_gradients.first
+    end
+
+    private
+
+    def initialize_weights(method, input_dim, units)
+      case method
+      when :xavier
+        std = Math.sqrt(2.0 / (input_dim + units))
+        Array.new(input_dim) { Array.new(units) { rand_normal(0, std) } }
+      when :he
+        std = Math.sqrt(2.0 / input_dim)
+        Array.new(input_dim) { Array.new(units) { rand_normal(0, std) } }
+      when :random
+        Array.new(input_dim) { Array.new(units) { rand * 2 - 1 } }
+      else
+        Array.new(input_dim) { Array.new(units, 0.0) }
+      end
+    end
+
+    def rand_normal(mean, std_dev)
+      u1 = rand
+      u2 = rand
+      z0 = Math.sqrt(-2 * Math.log(u1)) * Math.cos(2 * Math::PI * u2)
+      z0 * std_dev + mean
+    end
+
+    def multiply_matrix_vector(matrix, vector)
+      result = Array.new(matrix.length, 0.0)
+      matrix.each_with_index do |row, i|
+        result[i] = dot_product(row, vector)
+      end
+      result
+    end
+
+    def dot_product(a, b)
+      sum = 0.0
+      (0...a.length).each { |i| sum += a[i] * b[i] }
+      sum
+    end
+
+    def add_vectors(a, b)
+      (0...a.length).map { |i| a[i] + b[i] }
+    end
+
+    def multiply_elementwise(a, b)
+      (0...a.length).map { |i| a[i] * b[i] }
+    end
+
+    def multiply_outer_product(vec1, vec2)
+      Array.new(vec1.length) do |i|
+        Array.new(vec2.length) { |j| vec1[i] * vec2[j] }
+      end
+    end
+
+    def average_matrices(matrices)
+      rows, cols = matrices[0].length, matrices[0][0].length
+      avg = Array.new(rows) { Array.new(cols, 0.0) }
+      
+      matrices.each do |matrix|
+        matrix.each_with_index do |row, i|
+          row.each_with_index { |val, j| avg[i][j] += val }
+        end
+      end
+      
+      avg.map! { |row| row.map! { |val| val / matrices.length } }
+      avg
+    end
+
+    def average_arrays(arrays)
+      length = arrays[0].length
+      avg = Array.new(length, 0.0)
+      
+      arrays.each do |arr|
+        arr.each_with_index { |val, i| avg[i] += val }
+      end
+      
+      avg.map! { |val| val / arrays.length }
+    end
+  end
+
+  class ActivationLayer < Layer
+    attr_accessor :activation
+    
+    def initialize(activation)
+      super()
+      @activation = activation.is_a?(Symbol) ? GRNEXUSActivations.const_get(activation).new : activation
+      @trainable = false
+    end
+
+    def forward(input)
+      @cache[:input] = input.is_a?(Array) && input[0].is_a?(Array) ? input.dup : input.dup
+      if input[0].is_a?(Array)
+        input.map { |x| @activation.call(x) }
+      else
+        @activation.call(input)
+      end
+    end
+
+    def backward(gradient, learning_rate = nil)
+      input_cache = @cache[:input]
+      
+      if gradient[0].is_a?(Array)
+        gradient.zip(input_cache).map do |grad, cached_input|
+          activation_deriv = @activation.call(cached_input, derivative: true)
+          multiply_elementwise(grad, activation_deriv)
+        end
+      else
+        activation_deriv = @activation.call(input_cache, derivative: true)
+        multiply_elementwise(gradient, activation_deriv)
+      end
+    end
+
+    private
+
+    def multiply_elementwise(a, b)
+      (0...a.length).map { |i| a[i] * b[i] }
+    end
+  end
+
+  class DropoutLayer < Layer
+    attr_accessor :rate
+    
+    def initialize(rate: 0.5)
+      super()
+      @rate = rate
+      @mask = nil
+      @trainable = false
+    end
+
+    def forward(input, training: true)
+      @cache[:training] = training
+      if training
+        if input[0].is_a?(Array)
+          @mask = input.map { |sample| sample.map { rand > @rate ? 1.0 / (1.0 - @rate) : 0.0 } }
+          multiply_batch_elementwise(input, @mask)
+        else
+          @mask = input.map { rand > @rate ? 1.0 / (1.0 - @rate) : 0.0 }
+          multiply_elementwise(input, @mask)
+        end
+      else
+        input
+      end
+    end
+
+    def backward(gradient, learning_rate = nil)
+      training = @cache[:training]
+      if training && @mask
+        if gradient[0].is_a?(Array)
+          multiply_batch_elementwise(gradient, @mask)
+        else
+          multiply_elementwise(gradient, @mask)
+        end
+      else
+        gradient
+      end
+    end
+
+    private
+
+    def multiply_elementwise(a, b)
+      (0...a.length).map { |i| a[i] * b[i] }
+    end
+
+    def multiply_batch_elementwise(batch_a, batch_b)
+      batch_a.zip(batch_b).map do |a, b|
+        multiply_elementwise(a, b)
+      end
+    end
+  end
+
+  class BatchNormLayer < Layer
+    attr_accessor :epsilon, :momentum, :gamma, :beta
+    
+    def initialize(epsilon: 1e-5, momentum: 0.1)
+      super()
+      @epsilon = epsilon
+      @momentum = momentum
+      @running_mean = nil
+      @running_var = nil
+      @gamma = 1.0
+      @beta = 0.0
+      @trainable = true
+    end
+
+    def forward(input, training: true)
+      batch_case = input[0].is_a?(Array) && input[0][0].is_a?(Numeric)
+      input_tensor = batch_case ? input : [input]
+      
+      if training
+        batch_mean = calculate_mean(input_tensor)
+        batch_var = calculate_variance(input_tensor, batch_mean)
+        update_running_stats(batch_mean, batch_var)
+        
+        @cache[:mean] = batch_mean
+        @cache[:var] = batch_var
+        @cache[:inv_std] = batch_var.map { |v| 1.0 / Math.sqrt(v + @epsilon) }
+        @cache[:x_norm] = normalize_batch(input_tensor, batch_mean, batch_var)
+        @cache[:x_centered] = center_batch(input_tensor, batch_mean)
+      else
+        # Si no hay running stats, usar las del batch actual
+        if @running_mean.nil? || @running_var.nil?
+          batch_mean = calculate_mean(input_tensor)
+          batch_var = calculate_variance(input_tensor, batch_mean)
+          update_running_stats(batch_mean, batch_var)
+          @cache[:x_norm] = normalize_batch(input_tensor, batch_mean, batch_var)
+        else
+          running_inv_std = @running_var.map { |v| 1.0 / Math.sqrt(v + @epsilon) }
+          @cache[:x_norm] = normalize_batch(input_tensor, @running_mean, @running_var, running_inv_std)
+        end
+      end
+      
+      output_tensor = apply_affine_transform(batch_case ? @cache[:x_norm] : @cache[:x_norm].first)
+      batch_case ? output_tensor : output_tensor.first
+    end
+
+    def backward(gradient, learning_rate = nil)
+      batch_case = gradient[0].is_a?(Array) && gradient[0][0].is_a?(Numeric)
+      grad_tensor = batch_case ? gradient : [gradient]
+      
+      @gamma_gradient = calculate_gamma_gradient(grad_tensor, @cache[:x_norm])
+      @beta_gradient = calculate_beta_gradient(grad_tensor)
+      
+      # dx_norm = grad_tensor * gamma (element-wise)
+      dx_norm = grad_tensor.map { |grad_sample| grad_sample.map { |g| g * @gamma } }
+      
+      n = grad_tensor.length
+      inv_std = @cache[:inv_std]
+      
+      dx = []
+      (0...n).each do |i|
+        sum1 = multiply_elementwise(dx_norm[i], inv_std)
+        mean_dx_norm = dx_norm[i].map { |val| val / n }.reduce(:+)
+        sum2 = @cache[:x_centered][i].zip(inv_std).map { |c, s| c * s**3 }.map { |val| val / n }
+        sum2 = multiply_elementwise(sum2, Array.new(sum2.length, mean_dx_norm))
+        
+        dx << subtract_vectors(sum1, sum2)
+      end
+      
+      batch_case ? dx : dx.first
+    end
+
+    private
+
+    def calculate_mean(batch)
+      features = batch[0].length
+      means = Array.new(features, 0.0)
+      
+      batch.each do |sample|
+        sample.each_with_index { |val, i| means[i] += val }
+      end
+      
+      means.map! { |sum| sum / batch.length }
+    end
+
+    def calculate_variance(batch, means)
+      features = batch[0].length
+      vars = Array.new(features, 0.0)
+      
+      batch.each do |sample|
+        sample.each_with_index { |val, i| vars[i] += (val - means[i])**2 }
+      end
+      
+      vars.map! { |sum| sum / batch.length }
+    end
+
+    def update_running_stats(batch_mean, batch_var)
+      if @running_mean.nil?
+        @running_mean = batch_mean.dup
+        @running_var = batch_var.dup
+      else
+        @running_mean = multiply_scalar_add_vector((1 - @momentum), @running_mean, @momentum, batch_mean)
+        @running_var = multiply_scalar_add_vector((1 - @momentum), @running_var, @momentum, batch_var)
+      end
+    end
+
+    def normalize_batch(batch, means, vars, inv_std = nil)
+      inv_std ||= vars.map { |v| 1.0 / Math.sqrt(v + @epsilon) }
+      
+      batch.map do |sample|
+        sample.zip(means, inv_std).map { |val, mean, std_inv| (val - mean) * std_inv }
+      end
+    end
+
+    def center_batch(batch, means)
+      batch.map do |sample|
+        sample.zip(means).map { |val, mean| val - mean }
+      end
+    end
+
+    def apply_affine_transform(normalized)
+      if normalized[0].is_a?(Array)
+        normalized.map do |sample|
+          sample.map { |val| @gamma * val + @beta }
+        end
+      else
+        normalized.map { |val| @gamma * val + @beta }
+      end
+    end
+
+    def calculate_gamma_gradient(gradient_batch, x_norm_batch)
+      grad_sum = Array.new(x_norm_batch[0].length, 0.0)
+      
+      gradient_batch.each do |grad_sample|
+        grad_sample.each_with_index { |grad_val, i| grad_sum[i] += grad_val * x_norm_batch[gradient_batch.index(grad_sample)][i] }
+      end
+      
+      grad_sum
+    end
+
+    def calculate_beta_gradient(gradient_batch)
+      grad_sum = Array.new(gradient_batch[0].length, 0.0)
+      
+      gradient_batch.each do |grad_sample|
+        grad_sample.each_with_index { |grad_val, i| grad_sum[i] += grad_val }
+      end
+      
+      grad_sum
+    end
+
+    def multiply_batch_elementwise(batch_a, scalar_or_vector)
+      if scalar_or_vector.is_a?(Array)
+        batch_a.map do |sample|
+          sample.zip(scalar_or_vector).map { |val, mult| val * mult }
+        end
+      else
+        batch_a.map do |sample|
+          sample.map { |val| val * scalar_or_vector }
+        end
+      end
+    end
+    
+    def subtract_vectors(a, b)
+      (0...a.length).map { |i| a[i] - b[i] }
+    end
+    
+    def multiply_elementwise(a, b)
+      (0...a.length).map { |i| a[i] * b[i] }
+    end
+    
+    def multiply_scalar_add_vector(scalar1, vec1, scalar2, vec2)
+      (0...vec1.length).map { |i| scalar1 * vec1[i] + scalar2 * vec2[i] }
+    end
+  end
+
+  class Conv2DLayer < Layer
+    def initialize(filters:, kernel_size:, stride: 1, padding: 0)
+      super()
+      @filters = filters
+      @kernel_size = kernel_size.is_a?(Array) ? kernel_size : [kernel_size, kernel_size]
+      @stride = stride
+      @padding = padding
+      
+      kh, kw = @kernel_size
+      @kernels = initialize_conv_kernels(@filters, kh, kw)
+      @biases = Array.new(@filters, 0.0)
+    end
+
+    def forward(input)
+      batch_case = input[0].is_a?(Array) && input[0][0].is_a?(Array) && input[0][0][0].is_a?(Numeric)
+      input_tensor = batch_case ? input : [input]
+      
+      output_batch = input_tensor.map do |single_input|
+        convolve_2d(single_input)
+      end
+      
+      batch_case ? output_batch : output_batch.first
+    end
+
+    def backward(gradient)
+      gradient
+    end
+
+    private
+
+    def initialize_conv_kernels(filters, kh, kw)
+      fan_in = kh * kw
+      std = Math.sqrt(2.0 / fan_in)
+      
+      Array.new(filters) do
+        Array.new(kh) { Array.new(kw) { rand_normal(0, std) } }
+      end
+    end
+
+    def convolve_2d(input_image)
+      h, w = input_image.length, input_image[0].length
+      kh, kw = @kernels[0].length, @kernels[0][0].length
+      
+      out_h = (h + 2 * @padding - kh) / @stride + 1
+      out_w = (w + 2 * @padding - kw) / @stride + 1
+      
+      output = Array.new(out_h) { Array.new(out_w) { Array.new(@filters, 0.0) } }
+      padded = @padding > 0 ? pad_image(input_image, @padding) : input_image
+      
+      (0...out_h).each do |oh|
+        (0...out_w).each do |ow|
+          roi_start_h = oh * @stride
+          roi_start_w = ow * @stride
+          
+          (0...@filters).each do |f|
+            sum = 0.0
+            (0...kh).each do |kh_off|
+              (0...kw).each do |kw_off|
+                ih = roi_start_h + kh_off
+                iw = roi_start_w + kw_off
+                sum += padded[ih][iw] * @kernels[f][kh_off][kw_off]
+              end
+            end
+            output[oh][ow][f] = sum + @biases[f]
+          end
+        end
+      end
+      
+      output
+    end
+
+    def pad_image(image, padding)
+      h, w = image.length, image[0].length
+      padded_h, padded_w = h + 2 * padding, w + 2 * padding
+      
+      padded = Array.new(padded_h) { Array.new(padded_w, 0.0) }
+      
+      (0...h).each do |ih|
+        (0...w).each do |iw|
+          padded[ih + padding][iw + padding] = image[ih][iw]
+        end
+      end
+      
+      padded
+    end
+
+    def rand_normal(mean, std_dev)
+      u1 = rand
+      u2 = rand
+      z0 = Math.sqrt(-2 * Math.log(u1)) * Math.cos(2 * Math::PI * u2)
+      z0 * std_dev + mean
+    end
+  end
+
+  class MaxPoolingLayer < Layer
+    def initialize(pool_size:, stride: nil)
+      super()
+      @pool_size = pool_size.is_a?(Array) ? pool_size : [pool_size, pool_size]
+      @stride = stride || @pool_size
+      @stride = @stride.is_a?(Array) ? @stride : [@stride, @stride]
+      @trainable = false
+    end
+
+    def forward(input)
+      batch_case = input[0].is_a?(Array) && input[0][0].is_a?(Numeric)
+      input_tensor = batch_case ? input : [input]
+      
+      output_batch = input_tensor.map do |single_input|
+        pool_2d(single_input)
+      end
+      
+      batch_case ? output_batch : output_batch.first
+    end
+
+    def backward(gradient)
+      gradient
+    end
+
+    private
+
+    def pool_2d(input_image)
+      h, w = input_image.length, input_image[0].length
+      ph, pw = @pool_size
+      sh, sw = @stride
+      
+      out_h = (h - ph) / sh + 1
+      out_w = (w - pw) / sw + 1
+      
+      output = Array.new(out_h) { Array.new(out_w, 0.0) }
+      @cache[:switch_indices] ||= []
+      switch_indices_map = []
+      
+      (0...out_h).each do |oh|
+        (0...out_w).each do |ow|
+          pool_start_h = oh * sh
+          pool_start_w = ow * sw
+          
+          max_val = -Float::INFINITY
+          max_h, max_w = 0, 0
+          
+          (0...ph).each do |ph_off|
+            (0...pw).each do |pw_off|
+              ih = pool_start_h + ph_off
+              iw = pool_start_w + pw_off
+              if input_image[ih][iw] > max_val
+                max_val = input_image[ih][iw]
+                max_h, max_w = ih, iw
+              end
+            end
+          end
+          
+          output[oh][ow] = max_val
+          switch_indices_map << [max_h, max_w]
+        end
+      end
+      
+      @cache[:switch_indices] << switch_indices_map
+      output
+    end
+  end
+
+  class LSTMLayer < Layer
+    def initialize(units:, input_size:)
+      super()
+      @units = units
+      @input_size = input_size
+      @hidden_size = units
+      
+      @wf = initialize_weights(:xavier, @input_size, @units)
+      @uf = initialize_weights(:xavier, @hidden_size, @units)
+      @bf = Array.new(@units, 0.0)
+      
+      @wi = initialize_weights(:xavier, @input_size, @units)
+      @ui = initialize_weights(:xavier, @hidden_size, @units)
+      @bi = Array.new(@units, 0.0)
+      
+      @wo = initialize_weights(:xavier, @input_size, @units)
+      @uo = initialize_weights(:xavier, @hidden_size, @units)
+      @bo = Array.new(@units, 0.0)
+      
+      @wc = initialize_weights(:xavier, @input_size, @units)
+      @uc = initialize_weights(:xavier, @hidden_size, @units)
+      @bc = Array.new(@units, 0.0)
+    end
+
+    def forward(input_sequence)
+      batch_case = input_sequence[0][0].is_a?(Array) && input_sequence[0][0][0].is_a?(Numeric)
+      sequences = batch_case ? transpose_batch_sequences(input_sequence) : [input_sequence]
+      
+      outputs_batch = sequences.map do |single_sequence|
+        hidden_state = Array.new(@units, 0.0)
+        cell_state = Array.new(@units, 0.0)
+        outputs = []
+        
+        single_sequence.each do |input_t|
+          hidden_state, cell_state = lstm_step(input_t, hidden_state, cell_state)
+          outputs << hidden_state.dup
+        end
+        
+        outputs
+      end
+      
+      if batch_case
+        transpose_batch_sequences(outputs_batch)
+      else
+        outputs_batch.first
+      end
+    end
+
+    def backward(gradient)
+      gradient
+    end
+
+    private
+
+    def lstm_step(input_t, prev_hidden, prev_cell)
+      f_input = add_vectors(multiply_matrix_vector(@wf, input_t), multiply_matrix_vector(@uf, prev_hidden))
+      f_input = add_vectors(f_input, @bf)
+      f_gate = GRNEXUSActivations::Sigmoid.new.call(f_input)
+      
+      i_input = add_vectors(multiply_matrix_vector(@wi, input_t), multiply_matrix_vector(@ui, prev_hidden))
+      i_input = add_vectors(i_input, @bi)
+      i_gate = GRNEXUSActivations::Sigmoid.new.call(i_input)
+      
+      o_input = add_vectors(multiply_matrix_vector(@wo, input_t), multiply_matrix_vector(@uo, prev_hidden))
+      o_input = add_vectors(o_input, @bo)
+      o_gate = GRNEXUSActivations::Sigmoid.new.call(o_input)
+      
+      c_input = add_vectors(multiply_matrix_vector(@wc, input_t), multiply_matrix_vector(@uc, prev_hidden))
+      c_input = add_vectors(c_input, @bc)
+      candidate = GRNEXUSActivations::Tanh.new.call(c_input)
+      
+      new_cell = add_vectors(
+        multiply_elementwise(f_gate, prev_cell),
+        multiply_elementwise(i_gate, candidate)
+      )
+      
+      tanh_cell = GRNEXUSActivations::Tanh.new.call(new_cell)
+      new_hidden = multiply_elementwise(o_gate, tanh_cell)
+      
+      [new_hidden, new_cell]
+    end
+
+    def transpose_batch_sequences(sequences)
+      seq_len = sequences.length
+      batch_size = sequences[0].length
+      input_size = sequences[0][0].length
+      
+      transposed = Array.new(batch_size) { Array.new(seq_len) { Array.new(input_size) } }
+      
+      (0...seq_len).each do |t|
+        (0...batch_size).each do |b|
+          (0...input_size).each do |i|
+            transposed[b][t][i] = sequences[t][b][i]
+          end
+        end
+      end
+      
+      transposed
+    end
+
+    def initialize_weights(method, input_dim, units)
+      case method
+      when :xavier
+        std = Math.sqrt(2.0 / (input_dim + units))
+        Array.new(units) { Array.new(input_dim) { rand_normal(0, std) } }
+      else
+        Array.new(units) { Array.new(input_dim, 0.0) }
+      end
+    end
+
+    def rand_normal(mean, std_dev)
+      u1 = rand
+      u2 = rand
+      z0 = Math.sqrt(-2 * Math.log(u1)) * Math.cos(2 * Math::PI * u2)
+      z0 * std_dev + mean
+    end
+
+    def multiply_matrix_vector(matrix, vector)
+      result = Array.new(matrix.length, 0.0)
+      matrix.each_with_index do |row, i|
+        result[i] = dot_product(row, vector)
+      end
+      result
+    end
+
+    def dot_product(a, b)
+      sum = 0.0
+      (0...a.length).each { |i| sum += a[i] * b[i] }
+      sum
+    end
+
+    def add_vectors(a, b)
+      (0...a.length).map { |i| a[i] + b[i] }
+    end
+
+    def multiply_elementwise(a, b)
+      (0...a.length).map { |i| a[i] * b[i] }
+    end
+  end
+
+  class SoftmaxLayer < Layer
+    def initialize
+      super()
+      @trainable = false
+    end
+
+    def forward(input)
+      batch_case = input[0].is_a?(Array) && input[0][0].is_a?(Numeric)
+      input_tensor = batch_case ? input : [input]
+      
+      output_batch = input_tensor.map { |x| compute_softmax(x) }
+      batch_case ? output_batch : output_batch.first
+    end
+
+    def backward(gradient)
+      gradient
+    end
+
+    private
+
+    def compute_softmax(x)
+      max_val = x.max
+      exps = x.map { |val| Math.exp(val - max_val) }
+      sum_exps = exps.sum
+      exps.map { |exp| exp / sum_exps }
+    end
+  end
+
+  class GRULayer < Layer
+    def initialize(units:, input_size:)
+      super()
+      @units = units
+      @input_size = input_size
+      @hidden_size = units
+      
+      @wr = initialize_weights(:xavier, @input_size, @units)
+      @ur = initialize_weights(:xavier, @hidden_size, @units)
+      @br = Array.new(@units, 0.0)
+      
+      @wz = initialize_weights(:xavier, @input_size, @units)
+      @uz = initialize_weights(:xavier, @hidden_size, @units)
+      @bz = Array.new(@units, 0.0)
+      
+      @wh = initialize_weights(:xavier, @input_size, @units)
+      @uh = initialize_weights(:xavier, @hidden_size, @units)
+      @bh = Array.new(@units, 0.0)
+    end
+
+    def forward(input_sequence)
+      batch_case = input_sequence[0][0].is_a?(Array) && input_sequence[0][0][0].is_a?(Numeric)
+      sequences = batch_case ? transpose_batch_sequences(input_sequence) : [input_sequence]
+      
+      outputs_batch = sequences.map do |single_sequence|
+        hidden_state = Array.new(@units, 0.0)
+        outputs = []
+        
+        single_sequence.each do |input_t|
+          hidden_state = gru_step(input_t, hidden_state)
+          outputs << hidden_state.dup
+        end
+        
+        outputs
+      end
+      
+      if batch_case
+        transpose_batch_sequences(outputs_batch)
+      else
+        outputs_batch.first
+      end
+    end
+
+    def backward(gradient)
+      gradient
+    end
+
+    private
+
+    def gru_step(input_t, prev_hidden)
+      r_input = add_vectors(multiply_matrix_vector(@wr, input_t), multiply_matrix_vector(@ur, prev_hidden))
+      r_input = add_vectors(r_input, @br)
+      r_gate = GRNEXUSActivations::Sigmoid.new.call(r_input)
+      
+      z_input = add_vectors(multiply_matrix_vector(@wz, input_t), multiply_matrix_vector(@uz, prev_hidden))
+      z_input = add_vectors(z_input, @bz)
+      z_gate = GRNEXUSActivations::Sigmoid.new.call(z_input)
+      
+      rh_hidden = multiply_elementwise(r_gate, prev_hidden)
+      h_input = add_vectors(multiply_matrix_vector(@wh, input_t), multiply_matrix_vector(@uh, rh_hidden))
+      h_input = add_vectors(h_input, @bh)
+      h_tilde = GRNEXUSActivations::Tanh.new.call(h_input)
+      
+      one_minus_z = (0...@units).map { |i| 1.0 - z_gate[i] }
+      term1 = multiply_elementwise(one_minus_z, h_tilde)
+      term2 = multiply_elementwise(z_gate, prev_hidden)
+      new_hidden = add_vectors(term1, term2)
+      
+      new_hidden
+    end
+
+    def transpose_batch_sequences(sequences)
+      seq_len = sequences.length
+      batch_size = sequences[0].length
+      input_size = sequences[0][0].length
+      
+      transposed = Array.new(batch_size) { Array.new(seq_len) { Array.new(input_size) } }
+      
+      (0...seq_len).each do |t|
+        (0...batch_size).each do |b|
+          (0...input_size).each do |i|
+            transposed[b][t][i] = sequences[t][b][i]
+          end
+        end
+      end
+      
+      transposed
+    end
+
+    def initialize_weights(method, input_dim, units)
+      case method
+      when :xavier
+        std = Math.sqrt(2.0 / (input_dim + units))
+        Array.new(units) { Array.new(input_dim) { rand_normal(0, std) } }
+      else
+        Array.new(units) { Array.new(input_dim, 0.0) }
+      end
+    end
+
+    def rand_normal(mean, std_dev)
+      u1 = rand
+      u2 = rand
+      z0 = Math.sqrt(-2 * Math.log(u1)) * Math.cos(2 * Math::PI * u2)
+      z0 * std_dev + mean
+    end
+
+    def multiply_matrix_vector(matrix, vector)
+      result = Array.new(matrix.length, 0.0)
+      matrix.each_with_index do |row, i|
+        result[i] = dot_product(row, vector)
+      end
+      result
+    end
+
+    def dot_product(a, b)
+      sum = 0.0
+      (0...a.length).each { |i| sum += a[i] * b[i] }
+      sum
+    end
+
+    def add_vectors(a, b)
+      (0...a.length).map { |i| a[i] + b[i] }
+    end
+
+    def multiply_elementwise(a, b)
+      (0...a.length).map { |i| a[i] * b[i] }
+    end
+  end
+
+  class EmbeddingLayer < Layer
+    attr_accessor :weights
+    
+    def initialize(vocab_size:, embedding_dim:, padding_idx: nil)
+      super()
+      @vocab_size = vocab_size
+      @embedding_dim = embedding_dim
+      @padding_idx = padding_idx
+      @trainable = false  # Set to false for now (no gradient update)
+      
+      # Xavier initialization
+      limit = Math.sqrt(6.0 / (@vocab_size + @embedding_dim))
+      @weights = Array.new(@vocab_size) { Array.new(@embedding_dim) { rand(-limit..limit) } }
+    end
+
+    def forward(input)
+      # Check if input is a batch of sequences
+      batch_case = input[0].is_a?(Array)
+      input_tensor = batch_case ? input : [input]
+      
+      # Convert sequences to embeddings
+      output_batch = input_tensor.map do |sequence|
+        sequence.map do |idx|
+          idx_int = idx.to_i
+          # Clamp index to valid range
+          idx_int = [[idx_int, 0].max, @vocab_size - 1].min
+          @weights[idx_int].dup
+        end
+      end
+      
+      batch_case ? output_batch : output_batch.first
+    end
+
+    def backward(gradient, learning_rate = nil)
+      # For now, pass gradient through
+      # Full backprop through embeddings would require tracking indices
+      gradient
+    end
+  end
+
+  class FlattenLayer < Layer
+    def initialize
+      super()
+      @trainable = false
+      @input_shape = nil
+    end
+
+    def forward(input)
+      # Handle batch of sequences (3D: batch x sequence x features)
+      if input[0].is_a?(Array) && input[0][0].is_a?(Array) && input[0][0][0].is_a?(Numeric)
+        @input_shape = [input.length, input[0].length, input[0][0].length]
+        # Flatten each sample in the batch
+        input.map { |sample| sample.flatten }
+      # Handle batch of vectors (2D: batch x features)
+      elsif input[0].is_a?(Array) && input[0][0].is_a?(Numeric)
+        @input_shape = [input.length, input[0].length]
+        input
+      # Handle single sequence (2D: sequence x features)
+      elsif input[0].is_a?(Array)
+        @input_shape = [input.length, input[0].length]
+        [input.flatten]
+      # Handle single vector (1D: features)
+      else
+        @input_shape = [input.length]
+        [input]
+      end
+    end
+
+    def backward(gradient, learning_rate = nil)
+      gradient
+    end
+  end
+
+  class ReshapeLayer < Layer
+    def initialize(shape)
+      super()
+      @target_shape = shape
+      @trainable = false
+    end
+
+    def forward(input)
+      reshape_tensor(input, @target_shape)
+    end
+
+    def backward(gradient)
+      reshape_tensor(gradient, @original_shape)
+    end
+
+    private
+
+    def reshape_tensor(tensor, new_shape)
+      flattened = tensor.flatten
+      build_tensor(flattened, new_shape)
+    end
+
+    def build_tensor(flat_array, shape)
+      if shape.length == 1
+        flat_array
+      elsif shape.length == 2
+        rows, cols = shape
+        Array.new(rows) { |i| Array.new(cols) { |j| flat_array[i * cols + j] } }
+      else
+        size = shape[0]
+        remaining_shape = shape[1..-1]
+        remaining_size = remaining_shape.reduce(:*)
+        
+        Array.new(size) do |i|
+          sub_array = flat_array[i * remaining_size, remaining_size]
+          build_tensor(sub_array, remaining_shape)
+        end
+      end
+    end
+  end
+end