wip

Author: acharneski

Cargo.lock

@@ -4,7 +4,8 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
-luminal = "0.2"
+luminal = {path="/home/andrew/code/luminal"}
+luminal_training = {path="/home/andrew/code/luminal/crates/luminal_training"}
 serde = { version = "1.0", features = ["derive"] }
 thiserror = "2.0.12"
 rand = { version = "0.9.1", features = ["small_rng", "std"] }
@@ -28,6 +29,8 @@ rand_chacha = "0.9.0"
 rayon = "1.10.0"
 html-escape = "0.2.13"
 itertools = "0.13.0"
+dfdx = { version = "0.13", features = ["f16"] }
+luminal_nn = {path="/home/andrew/code/luminal/crates/luminal_nn"}
 
 [features]
 default = ["plotting"]

Cargo.toml

@@ -3,9 +3,12 @@
 
 use crate::benchmarks::functions::OptimizationProblem;
 use crate::optimizers::optimizer::Optimizer;
+use dfdx::prelude::Shape;
 use log::{debug, info, warn};
+use luminal::prelude::*;
 use rand::prelude::StdRng;
 use rand::{Rng, SeedableRng};
+use rand_distr::num_traits::ToPrimitive;
 use serde::{Deserialize, Serialize};
 use statrs::statistics::Statistics;
 use std::cmp::max;
@@ -14,8 +17,6 @@ use std::sync::atomic::AtomicBool;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 use std::time::{Duration, Instant};
-use luminal::prelude::*;
-use rand_distr::num_traits::ToPrimitive;
 use tokio::time::timeout;
 /// Global flag to disable optimal value thresholds for all problems
 static NO_THRESHOLD_MODE: AtomicBool = AtomicBool::new(false);
@@ -41,7 +42,6 @@ pub fn create_1d_tensor(data: &[f32], _device: &Device) -> Result<Tensor, String
     Ok(Tensor::new(data.to_vec()))
 }
 
-
 /// Wrapper for Duration that implements bincode traits
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct DurationWrapper {
@@ -135,9 +135,14 @@ impl OptimizationTrace {
         if self.iterations.is_empty() {
             None
         } else {
-            Some(Statistics::min(
-                self.iterations.iter().map(|data| data.function_value as f64),
-            ).to_f32()?)
+            Some(
+                Statistics::min(
+                    self.iterations
+                        .iter()
+                        .map(|data| data.function_value as f64),
+                )
+                .to_f32()?,
+            )
         }
     }
 
@@ -310,10 +315,10 @@ impl BenchmarkRunner {
     }
 
     /// Run benchmarks for all combinations of problems and optimizers
-    pub async fn run_benchmarks<S: Shape>(
+    pub async fn run_benchmarks(
         &self,
         problems: Vec<Box<ProblemSpec>>,
-        mut optimizers: Vec<Box<dyn Optimizer<S>>>,
+        mut optimizers: Vec<Box<dyn Optimizer>>,
     ) -> Result<BenchmarkResults, BenchmarkError> {
         let mut results = BenchmarkResults::new(self.config.clone());
         info!(
@@ -353,10 +358,10 @@ impl BenchmarkRunner {
     }
 
     /// Run a single benchmark with one problem and one optimizer
-    pub async fn run_single_benchmark<S: Shape>(
+    pub async fn run_single_benchmark(
         &self,
         problem: &ProblemSpec,
-        optimizer: &mut Box<dyn Optimizer<S>>,
+        optimizer: &mut Box<dyn Optimizer>,
         run_id: usize,
         opt_name: &str,
         initial_point: Result<Vec<f32>, Result<SingleResult, BenchmarkError>>,
@@ -504,10 +509,10 @@ impl BenchmarkRunner {
         }
     }
 
-    async fn optimization_loop<S: Shape>(
+    async fn optimization_loop(
         &self,
         problem: &ProblemSpec,
-        optimizer: &mut dyn Optimizer<S>,
+        optimizer: &mut dyn Optimizer,
         input_floats: &mut [f32],
         iteration: &mut usize,
         function_evaluations: &mut usize,
@@ -683,7 +688,7 @@ impl BenchmarkRunner {
             // This benchmark runner needs to be rewritten to support graph-based optimizers.
             // For now, we return an error to allow compilation.
             return Err(BenchmarkError::OptimizerError("Optimizer::step signature mismatch: BenchmarkRunner requires update for graph-based optimizers".to_string()));
-            
+
             #[allow(unreachable_code)]
             let step_result = crate::optimizers::optimizer::StepResult {
                 step_size: 0.0,
@@ -1096,4 +1101,4 @@ pub fn new_initial_point(
         *xi += (random_delta * 2.0 - 1.0) * noise; // Random perturbation
     }
     Ok(x)
-}
+}

src/benchmarks/evaluation.rs

@@ -132,9 +132,9 @@ impl ProblemTestResults {
             && self.finite_values_maintained
             && self.clone_behavior_correct
             && self
-            .derivative_validation_results
-            .numerical_gradient_accuracy
-            > 0.7
+                .derivative_validation_results
+                .numerical_gradient_accuracy
+                > 0.7
             && (self.derivative_validation_results.robustness_score > 0.5 ||
             // For ML problems, allow lower robustness scores if other metrics are good
             ((self.problem_name.contains("Regression") || self.problem_name.contains("SVM") || self.problem_name.contains("NeuralNetwork"))
@@ -445,8 +445,8 @@ impl UnifiedProblemTester {
                     // Only flag as error if the values are extremely large
                     if f_val.abs() > self.config.finite_check_tolerance
                         || grad
-                        .iter()
-                        .any(|&g| g.abs() > self.config.finite_check_tolerance)
+                            .iter()
+                            .any(|&g| g.abs() > self.config.finite_check_tolerance)
                     {
                         all_finite = false;
                         break;
@@ -1685,7 +1685,7 @@ mod tests {
                         .collect(),
                     0.01,
                 )
-                    .unwrap(),
+                .unwrap(),
             ),
             Box::new(SupportVectorMachine::new(svm_x, svm_y, 1.0).unwrap()),
             Box::new(NeuralNetworkTraining::mlp_classification(vec![3, 5, 2], &mut rng).unwrap()),
@@ -1733,7 +1733,7 @@ mod tests {
                     &mut rng,
                     None,
                 )
-                    .unwrap(),
+                .unwrap(),
             ),
             #[cfg(feature = "onednn")]
             Box::new(
@@ -1789,7 +1789,7 @@ mod tests {
                         .collect(),
                     0.01,
                 )
-                    .unwrap(),
+                .unwrap(),
             ),
         ];
         let results = test_multiple_problems(problems, None);
@@ -1956,4 +1956,4 @@ mod tests {
         // Should still pass with stricter config
         assert!(results.is_valid() || !results.errors.is_empty());
     }
-}
+}

src/benchmarks/unified_tests.rs

@@ -4,6 +4,7 @@ use crate::benchmarks::evaluation::{
     DurationWrapper, ProblemSpec,
 };
 use crate::Optimizer;
+use dfdx::prelude::Shape;
 use itertools::Itertools;
 use log::{debug, info, trace, warn};
 use rand::prelude::*;
@@ -17,7 +18,6 @@ use std::iter::Take;
 use std::slice::Iter;
 use std::sync::Arc;
 use std::time::Duration;
-use luminal::prelude::Shape;
 use tokio::sync::Semaphore;
 
 /// Detailed tracking of evolutionary events
@@ -158,11 +158,11 @@ impl AdaptiveExperimentRunner {
     }
 
     /// Run adaptive parameter evolution to find best optimizer configurations for each problem
-    pub async fn run_adaptive_evolution<S: Shape>(
+    pub async fn run_adaptive_evolution(
         &mut self,
         problems: Vec<ProblemSpec>,
         optimizer_types: Vec<super::parameter_evolution::OptimizerType>,
-    ) -> anyhow::Result<HashMap<String, Vec<(String, Arc<dyn Optimizer<S>>)>>> {
+    ) -> anyhow::Result<HashMap<String, Vec<(String, Arc<dyn Optimizer>)>>> {
         info!("Starting adaptive parameter evolution");
         info!(
             "Population size: {}, Generations: {}, Evaluation runs: {}",
@@ -221,7 +221,7 @@ impl AdaptiveExperimentRunner {
                 );
 
                 let family_best = self
-                    .evolve_optimizer_for_problem_family::<S>(
+                    .evolve_optimizer_for_problem_family(
                         family_problems.clone(),
                         optimizer_type.clone(),
                         &evolution_dir,
@@ -259,7 +259,7 @@ impl AdaptiveExperimentRunner {
                   family_name, all_best_genomes.len());
 
             // Convert best genomes to optimizers
-            let mut optimizers: Vec<(String, Arc<dyn Optimizer<S>>)> = Vec::new();
+            let mut optimizers: Vec<(String, Arc<dyn Optimizer>)> = Vec::new();
             let best: Vec<OptimizerGenome> = all_best_genomes
                 .iter()
                 .into_group_map_by(|x| x.optimizer_type.to_string())
@@ -289,7 +289,7 @@ impl AdaptiveExperimentRunner {
                     "Created evolved optimizer '{}' with fitness: {:?}",
                     name, genome.fitness
                 );
-                optimizers.push((name, genome.to_optimizer::<S>()));
+                optimizers.push((name, genome.to_optimizer()));
             }
 
             family_best_optimizers.insert(family_name.clone(), optimizers);
@@ -336,7 +336,7 @@ impl AdaptiveExperimentRunner {
         }
     }
 
-    async fn evolve_optimizer_for_problem_family<S: Shape>(
+    async fn evolve_optimizer_for_problem_family(
         &self,
         family_problems: Vec<ProblemSpec>,
         optimizer_type: super::parameter_evolution::OptimizerType,
@@ -394,7 +394,7 @@ impl AdaptiveExperimentRunner {
 
             // Evaluate fitness of population
             debug!("Starting fitness evaluation for generation {}", generation);
-            self.evaluate_population_on_family::<S>(
+            self.evaluate_population_on_family(
                 &mut population,
                 &family_problems,
                 &mut tracker,
@@ -492,7 +492,7 @@ impl AdaptiveExperimentRunner {
                 if i < sorted_population.len() {
                     let genome = &mut sorted_population[i];
                     match Self::evaluate_genome_with_metrics(
-                        genome.to_optimizer::<S>(),
+                        genome.to_optimizer(),
                         (&family_problems[0]).clone().clone(),
                         self.config.clone(),
                         1, // Just one run for emergency evaluation
@@ -573,7 +573,7 @@ impl AdaptiveExperimentRunner {
                     OptimizerGenome::new_random(optimizer_type.clone(), &mut rng);
                 // Try to evaluate the fallback genome
                 match Self::evaluate_genome(
-                    fallback_genome.to_optimizer::<S>(),
+                    fallback_genome.to_optimizer(),
                     (&family_problems[0]).clone(),
                     self.config.clone(),
                     1,
@@ -659,7 +659,7 @@ impl AdaptiveExperimentRunner {
         Ok(best_genomes)
     }
 
-    async fn evaluate_population_on_family<S: Shape>(
+    async fn evaluate_population_on_family(
         &self,
         population: &mut [OptimizerGenome],
         family_problems: &[ProblemSpec],
@@ -705,7 +705,7 @@ impl AdaptiveExperimentRunner {
             });
 
             let semaphore = semaphore.clone();
-            let optimizer = genome.to_optimizer::<S>();
+            let optimizer = genome.to_optimizer();
             let problems = family_problems.to_vec();
             let config = self.config.clone();
             let evaluation_runs = self.evaluation_runs;
@@ -854,7 +854,7 @@ impl AdaptiveExperimentRunner {
         Ok(())
     }
 
-    async fn evaluate_population<S: Shape>(
+    async fn evaluate_population(
         &self,
         population: &mut [OptimizerGenome],
         problem: &ProblemSpec,
@@ -895,7 +895,7 @@ impl AdaptiveExperimentRunner {
             });
 
             let semaphore = semaphore.clone();
-            let optimizer = genome.to_optimizer::<S>();
+            let optimizer = genome.to_optimizer();
             let problem = problem.clone();
             let config = self.config.clone();
             let evaluation_runs = self.evaluation_runs;
@@ -1297,8 +1297,8 @@ impl AdaptiveExperimentRunner {
         new_population
     }
 
-    async fn evaluate_genome_with_metrics<S: Shape>(
-        optimizer: Arc<dyn Optimizer<S>>,
+    async fn evaluate_genome_with_metrics(
+        optimizer: Arc<dyn Optimizer>,
         problem: ProblemSpec,
         config: BenchmarkConfig,
         num_runs: usize,
@@ -1376,8 +1376,8 @@ impl AdaptiveExperimentRunner {
         ))
     }
 
-    async fn evaluate_genome<S: Shape>(
-        optimizer: Arc<dyn Optimizer<S>>,
+    async fn evaluate_genome(
+        optimizer: Arc<dyn Optimizer>,
         problem: ProblemSpec,
         config: BenchmarkConfig,
         num_runs: usize,
@@ -1543,10 +1543,10 @@ impl AdaptiveExperimentRunner {
     }
 
     /// Run final championship with evolved optimizers
-    pub async fn run_evolved_championship<S: Shape>(
+    pub async fn run_evolved_championship(
         &self,
         problems: Vec<ProblemSpec>,
-        evolved_optimizers: HashMap<String, Vec<(String, Arc<dyn Optimizer<S>>)>>,
+        evolved_optimizers: HashMap<String, Vec<(String, Arc<dyn Optimizer>)>>,
     ) -> anyhow::Result<()> {
         info!("Running championship with evolved optimizers");
         info!("Championship includes {} problems", problems.len());
@@ -1927,7 +1927,7 @@ impl FamilyRepresentation {
 }
 
 /// Convenience function to run adaptive evolution experiments
-pub async fn run_adaptive_benchmark<S: Shape>(
+pub async fn run_adaptive_benchmark(
     report_path_prefix: &str,
     max_evals: usize,
     num_runs: usize,
@@ -1984,7 +1984,7 @@ pub async fn run_adaptive_benchmark<S: Shape>(
     // First, evolve optimizer parameters for each problem
     info!("Starting parameter evolution phase");
     let evolved_optimizers = runner
-        .run_adaptive_evolution::<S>(problems.clone(), optimizer_types)
+        .run_adaptive_evolution(problems.clone(), optimizer_types)
         .await?;
     info!("Parameter evolution phase completed");
 
@@ -1999,4 +1999,4 @@ pub async fn run_adaptive_benchmark<S: Shape>(
     info!("Results saved to: {}", output_dir.display());
 
     Ok(())
-}
+}