From 96a136b4e39e275f71cf046ae4b8c8696089599e Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Mon, 6 Jan 2025 23:01:47 -0500
Subject: [PATCH] override sum and max in TableDistribution

---
 gtsam/discrete/TableDistribution.cpp | 15 +++++++++++++++
 gtsam/discrete/TableDistribution.h   | 19 +++++++++++--------
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/gtsam/discrete/TableDistribution.cpp b/gtsam/discrete/TableDistribution.cpp
index 2a9c63d51..3c0605f27 100644
--- a/gtsam/discrete/TableDistribution.cpp
+++ b/gtsam/discrete/TableDistribution.cpp
@@ -91,6 +91,21 @@ bool TableDistribution::equals(const DiscreteFactor& other, double tol) const {
   }
 }
 
+/* ****************************************************************************/
+DiscreteFactor::shared_ptr TableDistribution::sum(size_t nrFrontals) const {
+  return table_.sum(nrFrontals);
+}
+
+/* ****************************************************************************/
+DiscreteFactor::shared_ptr TableDistribution::sum(const Ordering& keys) const {
+  return table_.sum(keys);
+}
+
+/* ****************************************************************************/
+DiscreteFactor::shared_ptr TableDistribution::max(size_t nrFrontals) const {
+  return table_.max(nrFrontals);
+}
+
 /* ****************************************************************************/
 DiscreteFactor::shared_ptr TableDistribution::max(const Ordering& keys) const {
   return table_.max(keys);
diff --git a/gtsam/discrete/TableDistribution.h b/gtsam/discrete/TableDistribution.h
index 3aafdfda7..1c393bb1a 100644
--- a/gtsam/discrete/TableDistribution.h
+++ b/gtsam/discrete/TableDistribution.h
@@ -111,14 +111,17 @@ class GTSAM_EXPORT TableDistribution : public DiscreteConditional {
     return table_.evaluate(values);
   }
 
-  /**
-   * @brief Create new factor by maximizing over all
-   * values with the same separator.
-   *
-   * @param keys The keys to sum over.
-   * @return DiscreteFactor::shared_ptr
-   */
-  virtual DiscreteFactor::shared_ptr max(const Ordering& keys) const override;
+  /// Create new factor by summing all values with the same separator values
+  DiscreteFactor::shared_ptr sum(size_t nrFrontals) const override;
+
+  /// Create new factor by summing all values with the same separator values
+  DiscreteFactor::shared_ptr sum(const Ordering& keys) const override;
+
+  /// Create new factor by maximizing over all values with the same separator.
+  DiscreteFactor::shared_ptr max(size_t nrFrontals) const override;
+
+  /// Create new factor by maximizing over all values with the same separator.
+  DiscreteFactor::shared_ptr max(const Ordering& keys) const override;
 
   /**
    * @brief Return assignment that maximizes value.