Merge pull request #39 from Ebanflo42/graph/backprop

src/core/graph/constant.rs

@@ -58,6 +58,7 @@ impl Context {
             operation: Operation::Constant(ConstantBinding { value }),
             dtype,
         });
+        self.constants.push(node_id);
         Ok(node_id)
     }
 
@@ -73,6 +74,7 @@ impl Context {
             operation: Operation::Constant(ConstantBinding { value }),
             dtype: T::TY,
         });
+        self.constants.push(node_id);
         Ok(node_id)
     }
 
@@ -94,6 +96,7 @@ impl Context {
             operation: Operation::Constant(ConstantBinding { value: reshaped }),
             dtype: T::TY,
         });
+        self.constants.push(node_id);
         Ok(node_id)
     }
 
@@ -108,6 +111,7 @@ impl Context {
             operation: Operation::Constant(ConstantBinding { value: l }),
             dtype: t,
         });
+        self.constants.push(node_id);
         Ok(node_id)
     }
 
@@ -134,6 +138,7 @@ impl Context {
             operation: Operation::Constant(ConstantBinding { value: new_value }),
             dtype: self.nodes[const_id].dtype,
         });
+        self.constants.push(node_id);
         Ok(node_id)
     }
 
@@ -153,6 +158,7 @@ impl Context {
             operation: Operation::Constant(ConstantBinding { value: new_value }),
             dtype: new_type,
         });
+        self.constants.push(node_id);
         Ok(node_id)
     }
 }

src/core/graph/context.rs

@@ -1,11 +1,15 @@
+use std::collections::HashMap;
+
 use super::*;
 use slotmap::SlotMap;
 
 /// XLA computation graph context.
 // TODO: rename this to something meaningful
 pub struct Context {
     pub(crate) nodes: SlotMap<NodeIdentifier, Node>,
-    pub(crate) param_indices: Vec<NodeIdentifier>,
+    pub(crate) constants: Vec<NodeIdentifier>,
+    pub(crate) parameters: Vec<NodeIdentifier>,
+    pub(crate) dependent_nodes: HashMap<NodeIdentifier, Vec<NodeIdentifier>>,
 }
 
 impl Default for Context {
@@ -44,7 +48,10 @@ pub enum ContextError {
     MultipleReturns(),
 
     #[error("Operation is not differentiable, to use it as a constant in a differentiable computation, wrap it with Context::stop_gradient.")]
-    NonDifferentiableError(Callsite),
+    NonDifferentiableOpError(Callsite),
+
+    #[error("Type is not differentiable, differentiable types are F16, Bf16, F32, F64, C64, C128")]
+    NonDifferentiableTypeError(Callsite),
 }
 
 pub type Result<T> = std::result::Result<T, ContextError>;
@@ -53,7 +60,9 @@ impl Context {
     pub fn new() -> Self {
         Self {
             nodes: SlotMap::with_key(),
-            param_indices: Vec::new(),
+            constants: Vec::new(),
+            parameters: Vec::new(),
+            dependent_nodes: HashMap::new(),
         }
     }
 
@@ -64,14 +73,18 @@ impl Context {
             Operation::Constant(a) => format!("Constant {} {}", input_node.shape, a),
             Operation::Parameter(a) => format!("Parameter {} {}", input_node.shape, a),
             Operation::StopGradient(a) => {
-                format!("StopGradient {} {}", input_node.shape, self.to_string(a))
+                format!("StopGradient ({})", self.to_string(a))
             }
-            Operation::Diff(a, b) => format!("Diff ({}) {}", self.to_string(a), self.to_string(b)),
             Operation::Add(a, b) => format!("Add ({}) ({})", self.to_string(a), self.to_string(b)),
+            Operation::Sub(a, b) => format!("Sub ({}) ({})", self.to_string(a), self.to_string(b)),
             Operation::Mul(a, b) => format!("Mul ({}) ({})", self.to_string(a), self.to_string(b)),
+            Operation::Neg(a) => format!("Neg ({})", self.to_string(a)),
             Operation::Equal(a, b) => {
                 format!("LessThan ({}) ({})", self.to_string(a), self.to_string(b))
             }
+            Operation::NotEqual(a, b) => {
+                format!("NotEqual ({}) ({})", self.to_string(a), self.to_string(b))
+            }
             Operation::LessThan(a, b) => {
                 format!("LessThan ({}) ({})", self.to_string(a), self.to_string(b))
             }
@@ -98,18 +111,26 @@ impl Context {
                 self.to_string(on_true),
                 self.to_string(on_false)
             ),
-            Operation::TypeCast(a, ty) => format!("TypeCast {} {}", self.to_string(a), ty),
+            Operation::TypeCast(a, ty) => format!("TypeCast ({}) {}", self.to_string(a), ty),
             Operation::SliceInDim {
                 node,
                 start,
                 stop,
                 stride,
                 dim,
             } => format!(
-                "SliceInDim {} {} {} {} {}",
+                "SliceInDim ({}) {} {} {} {}",
                 self.to_string(node), start, stop, stride, dim
             ),
-            Operation::ZerosLike(node) => format!("ZerosLike {}", self.to_string(node))
+            Operation::ZerosLike(node) => format!("ZerosLike {}", self.to_string(node)),
+            Operation::ReduceMax {
+                node,
+                dim,
+                keepdims,
+            } => format!(
+                "SliceInDim {} {} {}",
+                self.to_string(node), dim, keepdims
+            ),
         }
     }
 }

src/core/graph/logic.rs

@@ -1,14 +1,15 @@
 use super::*;
 
 impl Context {
-
-    pub fn select<A: Into<NodeIdentifier> + Copy, B: Into<NodeIdentifier> + Copy, C: Into<NodeIdentifier> + Copy>(
+    /// TODO: Typecheck pred
+    pub fn select(
         &mut self,
-        pred: A,
-        on_true: B,
-        on_false: C
+        pred: NodeIdentifier,
+        on_true: NodeIdentifier,
+        on_false: NodeIdentifier,
     ) -> Result<NodeIdentifier> {
         let pred = pred.into();
+        let pred = self.stop_gradient(pred);
         let on_true = on_true.into();
         let on_false = on_false.into();
         let node_pred = &self.nodes[pred];
@@ -28,25 +29,42 @@ impl Context {
                     node_false.shape.clone(),
                     callsite!(1),
                 )),
-                Some(s) => {
-                    match s.broadcast(&node_pred.shape) {
-                        None => Err(ContextError::IncompatibleOperandShapes(
-                            s,
-                            node_pred.shape.clone(),
-                            callsite!(1),
-                        )),
-                        Some(sh) => {
-                            let node = Node {
-                                callsite: callsite!(1),
-                                shape: sh,
-                                operation: Operation::Select{ pred: pred, on_true: on_true, on_false: on_false },
-                                dtype: node_true.dtype,
-                            };
-                            Ok(self.nodes.insert(node))
+                Some(s) => match s.broadcast(&node_pred.shape) {
+                    None => Err(ContextError::IncompatibleOperandShapes(
+                        s,
+                        node_pred.shape.clone(),
+                        callsite!(1),
+                    )),
+                    Some(sh) => {
+                        let node = Node {
+                            callsite: callsite!(1),
+                            shape: sh,
+                            operation: Operation::Select {
+                                pred: pred,
+                                on_true: on_true,
+                                on_false: on_false,
+                            },
+                            dtype: node_true.dtype,
+                        };
+                        let node_id = self.nodes.insert(node);
+                        self.dependent_nodes
+                            .entry(pred)
+                            .or_insert(Vec::new())
+                            .push(node_id);
+                        self.dependent_nodes
+                            .entry(on_true)
+                            .or_insert(Vec::new())
+                            .push(node_id);
+                        if on_true != on_false {
+                            self.dependent_nodes
+                                .entry(on_false)
+                                .or_insert(Vec::new())
+                                .push(node_id);
                         }
+                        Ok(node_id)
                     }
-                }
+                },
             }
         }
     }
-}
+}