Bugfixes, improvements, and new features, formatted files for better readability, changed the way transformers work

Author: Klus3kk

examples/advanced/neural_ode_example.py

@@ -27,9 +27,9 @@ def generate_spiral_data(n_points=100, noise=0.05):
     Returns:
         Tuple of (t, trajectory) where t are time points and trajectory is the path
     """
-    t = np.linspace(0, 2*np.pi, n_points)
-    x0 = np.cos(t) * np.exp(-0.1*t) + noise * np.random.randn(n_points)
-    x1 = np.sin(t) * np.exp(-0.1*t) + noise * np.random.randn(n_points)
+    t = np.linspace(0, 2 * np.pi, n_points)
+    x0 = np.cos(t) * np.exp(-0.1 * t) + noise * np.random.randn(n_points)
+    x1 = np.sin(t) * np.exp(-0.1 * t) + noise * np.random.randn(n_points)
 
     trajectory = np.stack([x0, x1], axis=1)
     return t, trajectory
@@ -38,67 +38,68 @@ def generate_spiral_data(n_points=100, noise=0.05):
 class SimpleNeuralODE:
     """
     A simplified Neural ODE implementation using Euler's method.
-    
+
     This demonstrates the core concept: learning the derivative function
     that governs the dynamics of a system.
     """
-    
+
     def __init__(self, ode_func, dt=0.01):
         """
         Initialize the Neural ODE.
-        
+
         Args:
             ode_func: Neural network that learns dx/dt = f(x, t)
             dt: Time step for numerical integration
         """
         self.ode_func = ode_func
         self.dt = dt
-    
+
     def __call__(self, x0, n_steps):
         """
         Integrate the ODE forward in time.
-        
+
         Args:
             x0: Initial condition (batch_size, state_dim)
             n_steps: Number of integration steps
-            
+
         Returns:
             Trajectory of states
         """
         trajectory = [x0]
         x = x0
-        
+
         for _ in range(n_steps):
             # Compute derivative: dx/dt = f(x)
             dx_dt = self.ode_func(x)
-            
+
             # Euler step: x_{t+1} = x_t + dt * dx/dt
             x_next_data = x.data + self.dt * dx_dt.data
             x_next = Tensor(x_next_data, requires_grad=x.requires_grad)
-            
+
             # Set up backward pass
             if x.requires_grad:
+
                 def _backward(x_curr=x, dx_dt_curr=dx_dt, x_next_curr=x_next):
                     if x_next_curr.grad is not None:
                         # Gradient flows back through Euler step
                         if x_curr.grad is None:
                             x_curr.grad = x_next_curr.grad.copy()
                         else:
                             x_curr.grad += x_next_curr.grad
-                        
+
                         # Gradient w.r.t. derivative
                         dx_dt_grad = self.dt * x_next_curr.grad
                         if dx_dt_curr.grad is None:
                             dx_dt_curr.grad = dx_dt_grad
                         else:
                             dx_dt_curr.grad += dx_dt_grad
-                
+
                 x_next._backward = _backward
                 x_next._prev = {x, dx_dt}
-            
+
             trajectory.append(x_next)
             x = x_next
-        
+
         return trajectory
 
 
@@ -108,89 +109,83 @@ def run_neural_ode_example():
     """
     print("Neural ODE Example: Learning Spiral Dynamics")
     print("=" * 50)
-    
+
     # Generate spiral data
     t, true_trajectory = generate_spiral_data(n_points=50, noise=0.02)
-    
+
     print(f"Generated spiral data with {len(true_trajectory)} points")
-    
+
     # Create Neural ODE function
     # This network learns dx/dt = f(x)
-    ode_func = Sequential(
-        Linear(2, 16),
-        Tanh(),
-        Linear(16, 16), 
-        Tanh(),
-        Linear(16, 2)
-    )
-    
+    ode_func = Sequential(Linear(2, 16), Tanh(), Linear(16, 16), Tanh(), Linear(16, 2))
+
     # Create Neural ODE solver
     neural_ode = SimpleNeuralODE(ode_func, dt=0.05)
-    
+
     # Prepare training data
     # Use consecutive points: x[i] -> x[i+1]
     X_data = true_trajectory[:-1]  # Current states
-    y_data = true_trajectory[1:]   # Next states
-    
+    y_data = true_trajectory[1:]  # Next states
+
     X_tensor = Tensor(X_data, requires_grad=True)
     y_tensor = Tensor(y_data, requires_grad=False)
-    
+
     # Loss function and optimizer
     loss_fn = MSELoss()
     optimizer = Adam(ode_func.parameters(), lr=0.001)
-    
+
     # Training loop
     epochs = 500
     losses = []
-    
+
     print("\nTraining Neural ODE...")
-    
+
     for epoch in range(epochs):
         # Zero gradients
         optimizer.zero_grad()
-        
+
         # Forward pass: predict next states
         predicted_trajectory = []
         for i in range(len(X_data)):
-            x_current = Tensor(X_data[i:i+1], requires_grad=True)
+            x_current = Tensor(X_data[i : i + 1], requires_grad=True)
             trajectory = neural_ode(x_current, n_steps=1)
             predicted_trajectory.append(trajectory[1])  # Next state
-        
+
         # Stack predictions
         predicted_next = Tensor(
             np.array([pred.data[0] for pred in predicted_trajectory]),
-            requires_grad=True
+            requires_grad=True,
         )
-        
+
         # Compute loss
         loss = loss_fn(predicted_next, y_tensor)
-        
+
         # Backward pass
         loss.backward()
-        
+
         # Update parameters
         optimizer.step()
-        
+
         losses.append(loss.data)
-        
+
         if epoch % 50 == 0:
             print(f"Epoch {epoch}: Loss = {loss.data:.6f}")
-    
+
     print("Training completed!")
-    
+
     # Generate predictions for longer trajectory
     print("\nGenerating predictions...")
-    
+
     # Start from initial condition
     x0 = Tensor(true_trajectory[0:1], requires_grad=False)
-    predicted_trajectory = neural_ode(x0, n_steps=len(true_trajectory)-1)
-    
+    predicted_trajectory = neural_ode(x0, n_steps=len(true_trajectory) - 1)
+
     # Extract data for plotting
     predicted_data = np.array([state.data[0] for state in predicted_trajectory])
-    
+
     # Plot results
     plot_results(true_trajectory, predicted_data, losses)
-    
+
     return True
 
 
@@ -200,55 +195,73 @@ def plot_results(true_trajectory, predicted_trajectory, losses):
     """
     try:
         plt.figure(figsize=(15, 5))
-        
+
         # Plot trajectories
         plt.subplot(1, 3, 1)
-        plt.plot(true_trajectory[:, 0], true_trajectory[:, 1], 'b-', label='True', linewidth=2)
-        plt.plot(predicted_trajectory[:, 0], predicted_trajectory[:, 1], 'r--', label='Predicted', linewidth=2)
-        plt.scatter(true_trajectory[0, 0], true_trajectory[0, 1], c='green', s=100, label='Start')
-        plt.title('Trajectory Comparison')
-        plt.xlabel('X0')
-        plt.ylabel('X1')
+        plt.plot(
+            true_trajectory[:, 0],
+            true_trajectory[:, 1],
+            "b-",
+            label="True",
+            linewidth=2,
+        )
+        plt.plot(
+            predicted_trajectory[:, 0],
+            predicted_trajectory[:, 1],
+            "r--",
+            label="Predicted",
+            linewidth=2,
+        )
+        plt.scatter(
+            true_trajectory[0, 0],
+            true_trajectory[0, 1],
+            c="green",
+            s=100,
+            label="Start",
+        )
+        plt.title("Trajectory Comparison")
+        plt.xlabel("X0")
+        plt.ylabel("X1")
         plt.legend()
         plt.grid(True)
-        plt.axis('equal')
-        
+        plt.axis("equal")
+
         # Plot loss curve
         plt.subplot(1, 3, 2)
         plt.plot(losses)
-        plt.title('Training Loss')
-        plt.xlabel('Epoch')
-        plt.ylabel('Loss')
-        plt.yscale('log')
+        plt.title("Training Loss")
+        plt.xlabel("Epoch")
+        plt.ylabel("Loss")
+        plt.yscale("log")
         plt.grid(True)
-        
+
         # Plot error over time
         plt.subplot(1, 3, 3)
         errors = np.linalg.norm(true_trajectory - predicted_trajectory, axis=1)
         plt.plot(errors)
-        plt.title('Prediction Error')
-        plt.xlabel('Time Step')
-        plt.ylabel('L2 Error')
+        plt.title("Prediction Error")
+        plt.xlabel("Time Step")
+        plt.ylabel("L2 Error")
         plt.grid(True)
-        
+
         plt.tight_layout()
-        plt.savefig('neural_ode_results.png', dpi=150, bbox_inches='tight')
+        plt.savefig("neural_ode_results.png", dpi=150, bbox_inches="tight")
         plt.show()
-        
+
         print("Results saved as 'neural_ode_results.png'")
-        
+
         # Print final statistics
         final_error = np.mean(errors)
         print(f"\nFinal average prediction error: {final_error:.4f}")
-        
+
     except Exception as e:
         print(f"Could not generate plot: {e}")
 
 
 if __name__ == "__main__":
     print("Simplified Neural ODE Example")
     print("=" * 30)
-    
+
     try:
         success = run_neural_ode_example()
         if success:
@@ -258,4 +271,5 @@ def plot_results(true_trajectory, predicted_trajectory, losses):
     except Exception as e:
         print(f"\n❌ Error running Neural ODE example: {e}")
         import traceback
-        traceback.print_exc()
+
+        traceback.print_exc()

examples/advanced/optimizer_comparison.py

@@ -73,7 +73,9 @@ def train_xor_with_optimizer(optimizer_name, epochs=2000, verbose=True):
     model.layers[0].bias.data = first_bias
 
     # Initialize second layer with smaller random weights
-    model.layers[2].weight.data = np.random.normal(0, 0.1, model.layers[2].weight.data.shape)
+    model.layers[2].weight.data = np.random.normal(
+        0, 0.1, model.layers[2].weight.data.shape
+    )
     model.layers[2].bias.data = np.zeros_like(model.layers[2].bias.data)
 
     # Set up optimizer
@@ -107,15 +109,15 @@ def train_xor_with_optimizer(optimizer_name, epochs=2000, verbose=True):
             outputs = model(X_tensor)
             loss = loss_fn(outputs, y_tensor)
             loss.backward()
-            
+
             # First step: perturb weights
             optimizer.first_step(zero_grad=True)
-            
+
             # Second forward pass
             outputs = model(X_tensor)
             loss = loss_fn(outputs, y_tensor)
             loss.backward()
-            
+
             # Second step: update weights
             optimizer.second_step(zero_grad=True)
         else:
@@ -135,8 +137,10 @@ def train_xor_with_optimizer(optimizer_name, epochs=2000, verbose=True):
         accuracies.append(accuracy)
 
         # Print progress occasionally
-        if verbose and (epoch % 200 == 0 or epoch == 1 or epoch == epochs-1):
-            print(f"{optimizer_name} - Epoch {epoch}/{epochs}, Loss: {losses[-1]:.4f}, Accuracy: {accuracy:.1f}%")
+        if verbose and (epoch % 200 == 0 or epoch == 1 or epoch == epochs - 1):
+            print(
+                f"{optimizer_name} - Epoch {epoch}/{epochs}, Loss: {losses[-1]:.4f}, Accuracy: {accuracy:.1f}%"
+            )
 
         # Early stopping if converged
         if accuracy == 100.0 and losses[-1] < 0.01 and epoch > 100:
@@ -287,8 +291,10 @@ def compare_optimizers():
     # Print summary
     print("\nOptimizer Performance Summary:")
     for opt in optimizer_names:
-        print(f"- {opt}: {results[opt]['final_accuracy']:.1f}% accuracy, final loss = {results[opt]['losses'][-1]:.6f}")
+        print(
+            f"- {opt}: {results[opt]['final_accuracy']:.1f}% accuracy, final loss = {results[opt]['losses'][-1]:.6f}"
+        )
 
 
 if __name__ == "__main__":
-    compare_optimizers()
+    compare_optimizers()