adds pool on/off

Author: Refefer

src/graph.rs

@@ -7,17 +7,23 @@ use crate::vecops::iadd;
 #[derive(Debug)]
 pub struct Graph {
     gradients: HashMap<NodeIdx, Vec<DType>>,
-    freelist: HashMap<usize, Vec<Vec<DType>>>
+    freelist: HashMap<usize, Vec<Vec<DType>>>,
+    nan_check: bool
 }
 
 impl Graph {
     pub fn new() -> Self {
         Graph {
             gradients: HashMap::new(),
-            freelist: HashMap::new()
+            freelist: HashMap::new(),
+            nan_check: false
         }
     }
 
+    pub fn debug_nan(&mut self, check: bool)  {
+        self.nan_check = check;
+    }
+
     pub fn backward(&mut self, end_node: &ANode) {
         let out = Run::new(end_node);
         // dz/dz of course is 1
@@ -86,6 +92,17 @@ impl Graph {
 
                 node.compute_grad(&node_grad, &mut temp_grads);
 
+                if self.nan_check {
+                    for (i, grad) in temp_grads.iter().enumerate() {
+                        for gi in grad.iter() {
+                            if gi.is_nan() {
+                                eprintln!("Nan detected with id {:?}, child {}", node.get_id(), i);
+                                panic!()
+                            }
+                        }
+                    }
+                }
+
                 // Update grads
                 grads.iter_mut().zip(temp_grads.into_iter()).for_each(|(g, tg)| {
                     iadd(g, &tg);

src/lib.rs

@@ -9,7 +9,7 @@ mod pool;
 
 pub use graph::Graph;
 pub use ops::{Variable,Constant};
-pub use pool::clear_pool;
+pub use pool::{clear_pool, use_shared_pool};
 
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;

src/ops.rs

@@ -20,7 +20,7 @@ impl Computation {
     }
 
     fn shared(value: Arc<Vec<DType>>) -> Self {
-        Computation { value: Data::Shared(value) }
+       Computation { value: Data::Shared(value) }
     }
 
     fn pooled(value: MPVec) -> Self {
@@ -427,6 +427,10 @@ impl Node for Power {
         grad.iter().zip(lx.zip(ly)).for_each(|(gi, (xi, yi))| {
             out.add(*gi * *yi * xi.powf(*yi - 1f32));
         });
+        println!("x: {:?}", x);
+        println!("y: {:?}", y);
+        println!("grad: {:?}", grad);
+        println!("dX: {:?}", child_grads[0]);
 
         // df(x,y)/dy = ln(y) * x ^ y
         let (lx, ly) = Broadcast::from_pair(x, y);
@@ -1147,6 +1151,18 @@ mod tests {
         assert_eq!(Some(&vec![0f32, 2f32, 4f32]), x_grad);
     }
 
+    fn euclidean_distance(x: &ANode, y: &ANode) -> ANode {
+        let minus = x - y;
+        println!("{:?}", minus.get_id());
+        let pow = minus.pow(2f32);
+        println!("{:?}", pow.get_id());
+        let sum = pow.sum();
+        println!("{:?}", sum.get_id());
+        let sqrt = sum.pow(0.5);
+        println!("{:?}", sqrt.get_id());
+        sqrt
+    }
+
     #[test]
     fn test_backward_pass_complicated() {
         // (x+2) ^ 2 

src/pool.rs

@@ -3,10 +3,12 @@ use lazy_static::lazy_static;
 use std::convert::{AsRef, AsMut};
 use std::collections::HashMap;
 use std::sync::Mutex;
+use std::sync::atomic::{AtomicBool,Ordering};
 use std::ops::{Drop,Deref,DerefMut};
 
 use crate::DType;
 
+static USE_POOL: AtomicBool = AtomicBool::new(true);
 lazy_static! {
     static ref POOL: Mutex<MemoryPool> = {
         let m = MemoryPool::new();
@@ -43,10 +45,18 @@ impl MemoryPool {
     }
 }
 
+pub fn use_shared_pool(use_pool: bool) {
+    USE_POOL.store(use_pool, Ordering::SeqCst);
+}
+
 pub fn allocate_vec(size: usize) -> MPVec {
-    let mut pool = POOL.lock()
-        .expect("Error accessing memory pool!");
-    pool.get(size)
+    if USE_POOL.load(Ordering::Relaxed) {
+        let mut pool = POOL.lock()
+            .expect("Error accessing memory pool!");
+        pool.get(size)
+    } else {
+        MPVec(vec![0.; size])
+    }
 }
 
 pub fn clear_pool() {
@@ -56,9 +66,11 @@ pub fn clear_pool() {
 }
 
 fn return_vec(v: Vec<DType>) {
-    let mut pool = POOL.lock()
-        .expect("Error accessing memory pool!");
-    pool.ret(v);
+    if USE_POOL.load(Ordering::Relaxed) {
+        let mut pool = POOL.lock()
+            .expect("Error accessing memory pool!");
+        pool.ret(v);
+    }
 }
 
 pub struct MPVec(Vec<DType>);