diff --git a/gtsam_unstable/nonlinear/Expression-inl.h b/gtsam_unstable/nonlinear/Expression-inl.h index 11fafd686..ef5d62185 100644 --- a/gtsam_unstable/nonlinear/Expression-inl.h +++ b/gtsam_unstable/nonlinear/Expression-inl.h @@ -135,6 +135,7 @@ struct JacobianTrace { T value() const { return t; } + virtual void reverseAD(JacobianMap& jacobians) const = 0; virtual void reverseAD(const Matrix& H, JacobianMap& jacobians) const = 0; }; @@ -212,9 +213,11 @@ public: /// Trace structure for reverse AD struct Trace: public JacobianTrace { - /// Return value and derivatives + /// If the expression is just a constant, we do nothing + virtual void reverseAD(JacobianMap& jacobians) const { + } + /// Base case: we simply ignore the given df/dT virtual void reverseAD(const Matrix& H, JacobianMap& jacobians) const { - // Base case: don't touch jacobians } }; @@ -269,9 +272,13 @@ public: /// Trace structure for reverse AD struct Trace: public JacobianTrace { Key key; - /// Return value and derivatives + /// If the expression is just a leaf, we just insert an identity matrix + virtual void reverseAD(JacobianMap& jacobians) const { + size_t n = T::Dim(); + jacobians.add(key, Eigen::MatrixXd::Identity(n, n)); + } + /// Base case: given df/dT, add it jacobians with correct key and we are done virtual void reverseAD(const Matrix& H, JacobianMap& jacobians) const { - // Base case: just insert a new H in the JacobianMap with correct key jacobians.add(key, H); } }; @@ -338,11 +345,12 @@ public: struct Trace: public JacobianTrace { boost::shared_ptr > trace1; Matrix H1; - /// Return value and derivatives + /// Start the reverse AD process + virtual void reverseAD(JacobianMap& jacobians) const { + trace1->reverseAD(H1, jacobians); + } + /// Given df/dT, multiply in dT/dA and continue reverse AD process virtual void reverseAD(const Matrix& H, JacobianMap& jacobians) const { - // This is a top-down calculation - // The end-result needs Jacobians to all leaf nodes. - // Since this is not a leaf node, we compute what is needed for leaf nodes here trace1->reverseAD(H * H1, jacobians); } }; @@ -421,12 +429,13 @@ public: boost::shared_ptr > trace1; boost::shared_ptr > trace2; Matrix H1, H2; - /// Return value and derivatives + /// Start the reverse AD process + virtual void reverseAD(JacobianMap& jacobians) const { + trace1->reverseAD(H1, jacobians); + trace2->reverseAD(H2, jacobians); + } + /// Given df/dT, multiply in dT/dA and continue reverse AD process virtual void reverseAD(const Matrix& H, JacobianMap& jacobians) const { - // This is a top-down calculation - // The end-result needs Jacobians to all leaf nodes. - // Since this is not a leaf node, we compute what is needed for leaf nodes here - // The binary node represents a fork in the tree, and hence we will get two Augmented maps trace1->reverseAD(H * H1, jacobians); trace2->reverseAD(H * H2, jacobians); } diff --git a/gtsam_unstable/nonlinear/Expression.h b/gtsam_unstable/nonlinear/Expression.h index f3653abdf..02d7aa620 100644 --- a/gtsam_unstable/nonlinear/Expression.h +++ b/gtsam_unstable/nonlinear/Expression.h @@ -107,8 +107,7 @@ public: #ifdef REVERSE_AD boost::shared_ptr > trace = root_->traceExecution(values); Augmented augmented(trace->value()); - size_t n = T::Dim(); - trace->reverseAD(Eigen::MatrixXd::Identity(n, n), augmented.jacobians()); + trace->reverseAD(augmented.jacobians()); return augmented; #else return root_->forward(values);