diff --git a/gtsam/discrete/DiscreteFactorGraph.cpp b/gtsam/discrete/DiscreteFactorGraph.cpp
index 4ededbb8b..f81c6085e 100644
--- a/gtsam/discrete/DiscreteFactorGraph.cpp
+++ b/gtsam/discrete/DiscreteFactorGraph.cpp
@@ -112,7 +112,7 @@ namespace gtsam {
 
   /* ************************************************************************ */
   // Alternate eliminate function for MPE
-  std::pair<DiscreteConditional::shared_ptr, DecisionTreeFactor::shared_ptr>  //
+  std::pair<DiscreteConditional::shared_ptr, DiscreteFactor::shared_ptr>
   EliminateForMPE(const DiscreteFactorGraph& factors,
                   const Ordering& frontalKeys) {
     // PRODUCT: multiply all factors
@@ -201,7 +201,7 @@ namespace gtsam {
   }
 
   /* ************************************************************************ */
-  std::pair<DiscreteConditional::shared_ptr, DecisionTreeFactor::shared_ptr>  //
+  std::pair<DiscreteConditional::shared_ptr, DiscreteFactor::shared_ptr>
   EliminateDiscrete(const DiscreteFactorGraph& factors,
                     const Ordering& frontalKeys) {
     // PRODUCT: multiply all factors
diff --git a/gtsam/discrete/DiscreteFactorGraph.h b/gtsam/discrete/DiscreteFactorGraph.h
index d0dc282b4..a5324811c 100644
--- a/gtsam/discrete/DiscreteFactorGraph.h
+++ b/gtsam/discrete/DiscreteFactorGraph.h
@@ -48,7 +48,7 @@ class DiscreteJunctionTree;
  * @ingroup discrete
  */
 GTSAM_EXPORT
-std::pair<DiscreteConditional::shared_ptr, DecisionTreeFactor::shared_ptr>
+std::pair<DiscreteConditional::shared_ptr, DiscreteFactor::shared_ptr>
 EliminateDiscrete(const DiscreteFactorGraph& factors,
                   const Ordering& frontalKeys);
 
@@ -61,7 +61,7 @@ EliminateDiscrete(const DiscreteFactorGraph& factors,
  * @ingroup discrete
  */
 GTSAM_EXPORT
-std::pair<DiscreteConditional::shared_ptr, DecisionTreeFactor::shared_ptr>
+std::pair<DiscreteConditional::shared_ptr, DiscreteFactor::shared_ptr>
 EliminateForMPE(const DiscreteFactorGraph& factors,
                 const Ordering& frontalKeys);
 
@@ -133,6 +133,7 @@ class GTSAM_EXPORT DiscreteFactorGraph
 
   /// @}
 
+  //TODO(Varun): Make compatible with TableFactor
   /** Add a decision-tree factor */
   template <typename... Args>
   void add(Args&&... args) {
@@ -146,7 +147,7 @@ class GTSAM_EXPORT DiscreteFactorGraph
   DiscreteKeys discreteKeys() const;
 
   /** return product of all factors as a single factor */
-  DecisionTreeFactor product() const;
+  DiscreteFactor::shared_ptr product() const;
 
   /** 
    * Evaluates the factor graph given values, returns the joint probability of