Implement Const Folding for Mul, Add and Sub

src/core/graph/autodiff.rs

@@ -61,7 +61,6 @@ impl Context {
                 self.nodes[with_respect_to].callsite.clone(),
             ));
         }
-
         if node == with_respect_to {
             return self.scalar(1, wrt_dtype);
         }

src/core/graph/compile.rs

@@ -20,10 +20,10 @@ pub enum CompileError {
 }
 
 impl Context {
-    pub fn compile<A: Into<NodeIdentifier> + Copy, const N: usize>(
+    pub fn compile<const N: usize>(
         &mut self,
         name: &str,
-        returns: [A; N],
+        returns: [NodeIdentifier; N],
         client: &xla::PjRtClient,
     ) -> Result<xla::PjRtLoadedExecutable> {
         // TODO: gate debug mode behind a feature flag
@@ -33,7 +33,7 @@ impl Context {
         }
 
         for a in returns.iter() {
-            self.foldconsts(*a, usize::MAX)?;
+            self.fold_consts(*a, usize::MAX)?;
         }
         //while self.foldconsts(a, 1)? {
         //    println!("{}", self.to_string(a));