sgtlearn 0.2.0__tar.gz → 0.3.0__tar.gz

{sgtlearn-0.2.0 → sgtlearn-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: sgtlearn
-Version: 0.2.0
+Version: 0.3.0
 Summary: Shape Generalized Trees learning library
 Author-Email: Nakul Upadhya <nakulupadhya1@gmail.com>, Joshua Lee <joshua.lee.9880@gmail.com>, Eldan Cohen <eldan.cohen@utoronto.ca>
 License: MIT
@@ -31,7 +31,7 @@ Description-Content-Type: text/markdown
 
 `sgtlearn` is a Python package for learning [Shape Generalized Trees (SGTs)](https://neurips.cc/virtual/2025/loc/san-diego/poster/115950).
 
-- 🌳 **Shape Generalized Trees (SGTs):** A class of decision trees where each node applies a learnable, axis-aligned shape function to a feature for non-linear and interpretable splits.
+- 🌳 **Shape Generalized Trees (SGTs):** A class of decision trees where each node applies a learnable, axis-aligned shape function to one or two logical features for non-linear and interpretable splits.
 - 👁 **Interpretability:** Each node's shape function can be visualized directly.
 - ⚡ **ShapeCART Algorithm:** An efficient induction method for learning SGTs from data.
 - 🔀 **Extensions:**
@@ -39,10 +39,9 @@ Description-Content-Type: text/markdown
   - **SGT<sub>K</sub>:** Multi-way branching generalization.
   - **Shape²CART & ShapeCART<sub>K</sub>:** Algorithms for learning S²GTs and SGT<sub>K</sub>s.
 
+
 > [!NOTE]
-> This codebase is an efficient, but working implementation of the algorithms in the paper "Empowering Decision Trees via Shape Function Branching". Please refer to the [ROADMAP](ROADMAP.md) for a detailed list of features that are currently implemented and those that are planned for future releases. For the canonical code base for the paper, please refer to https://github.com/optimal-uoft/Empowering-DTs-via-Shape-Functions. Features in the paper that are not yet implemented in this codebase include: 
-> * Bivariate shape functions (Shape$^2$CART) + Higher branching factors for bivariate splits (Shape$^2$SGT$_K$)
-> * Visualization for bivariate splits (ex. contour plots)
+> This codebase is an efficient implementation of the algorithms in "Empowering Decision Trees via Shape Function Branching." See the [ROADMAP](ROADMAP.md) for implementation status and the [canonical research code](https://github.com/optimal-uoft/Empowering-DTs-via-Shape-Functions) for the paper's original implementation.
 
 ## Installation
 

{sgtlearn-0.2.0 → sgtlearn-0.3.0}/README.md

@@ -4,7 +4,7 @@
 
 `sgtlearn` is a Python package for learning [Shape Generalized Trees (SGTs)](https://neurips.cc/virtual/2025/loc/san-diego/poster/115950).
 
-- 🌳 **Shape Generalized Trees (SGTs):** A class of decision trees where each node applies a learnable, axis-aligned shape function to a feature for non-linear and interpretable splits.
+- 🌳 **Shape Generalized Trees (SGTs):** A class of decision trees where each node applies a learnable, axis-aligned shape function to one or two logical features for non-linear and interpretable splits.
 - 👁 **Interpretability:** Each node's shape function can be visualized directly.
 - ⚡ **ShapeCART Algorithm:** An efficient induction method for learning SGTs from data.
 - 🔀 **Extensions:**
@@ -12,10 +12,9 @@
   - **SGT<sub>K</sub>:** Multi-way branching generalization.
   - **Shape²CART & ShapeCART<sub>K</sub>:** Algorithms for learning S²GTs and SGT<sub>K</sub>s.
 
+
 > [!NOTE]
-> This codebase is an efficient, but working implementation of the algorithms in the paper "Empowering Decision Trees via Shape Function Branching". Please refer to the [ROADMAP](ROADMAP.md) for a detailed list of features that are currently implemented and those that are planned for future releases. For the canonical code base for the paper, please refer to https://github.com/optimal-uoft/Empowering-DTs-via-Shape-Functions. Features in the paper that are not yet implemented in this codebase include: 
-> * Bivariate shape functions (Shape$^2$CART) + Higher branching factors for bivariate splits (Shape$^2$SGT$_K$)
-> * Visualization for bivariate splits (ex. contour plots)
+> This codebase is an efficient implementation of the algorithms in "Empowering Decision Trees via Shape Function Branching." See the [ROADMAP](ROADMAP.md) for implementation status and the [canonical research code](https://github.com/optimal-uoft/Empowering-DTs-via-Shape-Functions) for the paper's original implementation.
 
 ## Installation
 

{sgtlearn-0.2.0 → sgtlearn-0.3.0}/ROADMAP.md

@@ -17,6 +17,16 @@
 - [x] **NaN routing at predict**: if training saw missing at that split, follow the stored direction; otherwise route to the majority child.
 
 ## v0.3.0
-- [ ] multioutput support
-- [ ] Shape$^2$CART
-- [ ] Shape$^2$CART Random Forest Ensembling  
+- [x] multioutput support
+- [x] Opt-in Shape$^2$CART for SGT estimators, including continuous/categorical
+  pairs, joint missing routing, and multiway branching ([tutorial](https://sgtlearn.readthedocs.io/en/latest/tutorials/bivariate-branching.html))
+- [x] Shape$^2$CART Random Forest Ensembling
+- [x] Pair-aware TAO refinement ([#48](https://github.com/optimal-uoft/sgtlearn/issues/48))
+- [x] Shape$^2$CART routing heatmap visualization ([#28](https://github.com/optimal-uoft/sgtlearn/issues/28))
+
+See the implementation specification in [#42](https://github.com/optimal-uoft/sgtlearn/issues/42)
+and the umbrella issue [#27](https://github.com/optimal-uoft/sgtlearn/issues/27).
+
+
+## v1.0.0
+- [ ] Boosting

{sgtlearn-0.2.0 → sgtlearn-0.3.0}/cpp/CMakeLists.txt

@@ -77,6 +77,10 @@ add_library(sgtlearn_core STATIC
         src/Splitters/categorical/CategoricalRegressionSplitter.h
         src/Splitters/categorical/CategoricalRegressionSplitter.cpp
         src/Discretizers/univariate/UnivariateClassificationDiscretizer.cpp
+        src/Discretizers/pair/PairClassificationDiscretizer.h
+        src/Discretizers/pair/PairClassificationDiscretizer.cpp
+        src/Discretizers/pair/PairRegressionDiscretizer.h
+        src/Discretizers/pair/PairRegressionDiscretizer.cpp
         src/Splitters/univariate/SquaredErrorSplitter.h
         src/Splitters/univariate/SquaredErrorSplitter.cpp
         src/Discretizers/univariate/UnivariateRegressionDiscretizer.cpp
@@ -84,6 +88,8 @@ add_library(sgtlearn_core STATIC
         src/Discretizers/univariate/GainHessianUnivariateDiscretizer.h
         src/algorithms/WaveletTreeMAE.h
         src/algorithms/WaveletTreeMAE.cpp
+        src/algorithms/WeightedMAETree.h
+        src/algorithms/WeightedMAETree.cpp
         src/Discretizers/univariate/GainHessianUnivariateDiscretizer.cpp
         src/Splitters/univariate/AbsoluteErrorSplitter.h
         src/Splitters/univariate/AbsoluteErrorSplitter.cpp
@@ -119,8 +125,14 @@ add_library(sgtlearn_core STATIC
         src/Estimators/RegressionShapeGeneralizedTree.h
         src/Estimators/RegressionShapeGeneralizedTree.cpp
         src/BranchAssignmentObjectives/BranchAssignment.h
+        src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentCommon.h
         src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignment.h
         src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignment.cpp
+        src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentBst.h
+        src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentBst.cpp
+        src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentSort.h
+        src/BranchAssignmentObjectives/AbsoluteErrorBranchAssignmentSort.cpp
+        src/BranchAssignmentObjectives/MaeBranchConfig.h
         src/BranchAssignmentObjectives/LeafAggregateProcessor.h
         src/BranchAssignmentObjectives/LeafAggregationBranchAssignment.h
         src/BranchAssignmentObjectives/LeafAggregationBranchAssignment.cpp
@@ -223,8 +235,10 @@ if (SGTLEARN_BUILD_TESTS)
 
     add_executable(cpp_tests
             tests/test_wavelet_tree_mae.cpp
+            tests/test_weighted_mae_tree.cpp
             tests/test_splitters.cpp
             tests/test_branch_assignment.cpp
+            tests/bench_mae_branch_assignment.cpp
     )
 
     target_link_libraries(cpp_tests PRIVATE
@@ -240,4 +254,4 @@ if (SGTLEARN_BUILD_TESTS)
 
 endif ()
 
-# endregion
+# endregion

{sgtlearn-0.2.0 → sgtlearn-0.3.0}/cpp/bindings/Discretizers.cpp

@@ -15,6 +15,9 @@
 #include <string>
 #include <variant>
 
+#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
+#include <numpy/arrayobject.h>
+
 #include "Splitters/univariate/AbsoluteErrorSplitter.h"
 #include "Discretizers/categorical/CategoricalClassificationDiscretizer.h"
 #include "Discretizers/categorical/CategoricalRegressionDiscretizer.h"
@@ -58,14 +61,6 @@ py::array_t<size_t> col_to_numpy_1d(const arma::Col<size_t> &col) {
   return out;
 }
 
-py::array_t<float> vector_float_to_numpy_1d(const std::vector<float> &v) {
-  py::array_t<float> out({static_cast<py::ssize_t>(v.size())});
-  auto buf = out.mutable_unchecked<1>();
-  for (size_t i = 0; i < v.size(); ++i)
-    buf(static_cast<py::ssize_t>(i)) = v[i];
-  return out;
-}
-
 std::string normalize_regression_criterion(std::string s) {
   s = normalize_criterion(std::move(s));
   if (s == "mse")
@@ -115,6 +110,119 @@ VectorOfVectorsToNumpyList(const std::vector<std::vector<size_t>> &bins) {
   return out;
 }
 
+/**
+ * Build a ``(n_outputs, n_samples)`` label matrix from a NumPy array that is
+ * either 1-D ``(n_samples,)`` or 2-D ``(n_samples, n_outputs)``.
+ */
+arma::Mat<size_t> classification_y_to_arma(const py::array_t<size_t> &y,
+                                           arma::uword n_samples) {
+  py::array_t<size_t> ycopy = py::array_t<size_t>::ensure(y);
+  arma::Mat<size_t> yArma;
+  if (ycopy.ndim() == 1) {
+    yArma = carma::arr_to_row<size_t>(ycopy, true);
+  } else if (ycopy.ndim() == 2) {
+    yArma = arma::Mat<size_t>(carma::arr_to_mat<size_t>(ycopy, true).t());
+  } else {
+    throw std::invalid_argument(
+        "y must be a 1D (n_samples,) or 2D (n_samples, n_outputs) numpy array");
+  }
+  if (yArma.n_cols != n_samples)
+    throw std::invalid_argument("y length must match X.shape[0]");
+  return yArma;
+}
+
+/**
+ * Build a ``(n_outputs, n_samples)`` target matrix from a NumPy array that is
+ * either 1-D ``(n_samples,)`` or 2-D ``(n_samples, n_outputs)``.
+ */
+arma::Mat<float> regression_y_to_arma(const py::array_t<float> &y,
+                                      arma::uword n_samples) {
+  py::array_t<float> ycopy = py::array_t<float>::ensure(y);
+  arma::Mat<float> yArma;
+  if (ycopy.ndim() == 1) {
+    yArma = carma::arr_to_row<float>(ycopy, true);
+  } else if (ycopy.ndim() == 2) {
+    yArma = arma::Mat<float>(carma::arr_to_mat<float>(ycopy, true).t());
+  } else {
+    throw std::invalid_argument(
+        "y must be a 1D (n_samples,) or 2D (n_samples, n_outputs) numpy array");
+  }
+  if (yArma.n_cols != n_samples)
+    throw std::invalid_argument("y length must match X.shape[0]");
+  return yArma;
+}
+
+/**
+ * Resolve a Python ``numClasses`` argument (int or sequence of ints) into a
+ * per-output class-count vector of length ``n_outputs``. An int is broadcast
+ * to every output; a sequence must match ``n_outputs``.
+ */
+std::vector<size_t> num_classes_per_output_from_py(const py::object &numClasses,
+                                                   arma::uword n_outputs) {
+  py::object numbers = py::module_::import("numbers");
+  if (py::isinstance<py::bool_>(numClasses))
+    throw std::invalid_argument("numClasses cannot be bool");
+  if (py::isinstance(numClasses, numbers.attr("Integral")))
+    return std::vector<size_t>(static_cast<size_t>(n_outputs),
+                               py::cast<size_t>(numClasses));
+  if (py::isinstance<py::iterable>(numClasses)) {
+    std::vector<size_t> out;
+    for (const py::handle item : numClasses)
+      out.push_back(py::cast<size_t>(py::reinterpret_borrow<py::object>(item)));
+    if (out.size() != static_cast<size_t>(n_outputs))
+      throw std::invalid_argument(
+          "numClasses sequence length must match the number of outputs "
+          "(y.shape[1])");
+    return out;
+  }
+  throw std::invalid_argument(
+      "numClasses must be an int or a sequence of ints");
+}
+
+/**
+ * Convert per-bin classification predictions (one vector of length
+ * ``n_outputs`` per bin) into a NumPy array: 1-D ``(n_bins,)`` for a single
+ * output, else 2-D ``(n_bins, n_outputs)``.
+ */
+template <typename T>
+py::array bin_predictions_to_numpy_impl(const std::vector<std::vector<T>> &preds) {
+  const size_t nBins = preds.size();
+  // Bins with no samples (e.g. the trailing NaN routing bin) carry an empty
+  // prediction vector, so derive the output width from the widest bin.
+  size_t nOut = 0;
+  for (const auto &p : preds)
+    nOut = std::max(nOut, p.size());
+  if (nOut == 0)
+    nOut = 1;
+  if (nOut == 1) {
+    py::array_t<T> out({static_cast<py::ssize_t>(nBins)});
+    auto buf = out.template mutable_unchecked<1>();
+    for (size_t i = 0; i < nBins; ++i)
+      buf(static_cast<py::ssize_t>(i)) =
+          preds[i].empty() ? static_cast<T>(0) : preds[i][0];
+    return out;
+  }
+  py::array_t<T> out(
+      {static_cast<py::ssize_t>(nBins), static_cast<py::ssize_t>(nOut)});
+  auto buf = out.template mutable_unchecked<2>();
+  for (size_t i = 0; i < nBins; ++i)
+    for (size_t o = 0; o < nOut; ++o)
+      buf(static_cast<py::ssize_t>(i), static_cast<py::ssize_t>(o)) =
+          o < preds[i].size() ? preds[i][o] : static_cast<T>(0);
+  return out;
+}
+
+py::array
+bin_predictions_to_numpy(const std::vector<std::vector<size_t>> &preds) {
+  return bin_predictions_to_numpy_impl(preds);
+}
+
+/** Regression counterpart of :func:`bin_predictions_to_numpy`. */
+py::array
+bin_predictions_to_numpy(const std::vector<std::vector<float>> &preds) {
+  return bin_predictions_to_numpy_impl(preds);
+}
+
 class UnivariateClassificationDiscretizerPy {
   std::variant<GiniDisc, EntropyDisc> impl_;
 
@@ -132,7 +240,7 @@ public:
   }
 
   void Train(const py::array_t<float> &X, const py::array_t<size_t> &features,
-             const py::array_t<size_t> &y, size_t numClasses,
+             const py::array_t<size_t> &y, py::object numClasses,
              size_t minLeafSize, double minGainSplit, size_t maxDepth,
              size_t maxLeafNodes, py::object sample_weights = py::none()) {
     py::array_t<float> Xcopy = py::array_t<float>::ensure(X);
@@ -148,12 +256,9 @@ public:
     const arma::uvec featuresArma =
         arma::conv_to<arma::uvec>::from(carma::arr_to_col<size_t>(fcopy, true));
 
-    py::array_t<size_t> ycopy = py::array_t<size_t>::ensure(y);
-    if (ycopy.ndim() != 1)
-      throw std::invalid_argument("y must be a 1D numpy array");
-    arma::Mat<size_t> yArma = carma::arr_to_row<size_t>(ycopy, true);
-    if (yArma.n_elem != xArma.n_cols)
-      throw std::invalid_argument("y length must match X.shape[0]");
+    arma::Mat<size_t> yArma = classification_y_to_arma(y, xArma.n_cols);
+    const std::vector<size_t> nClassesPerOutput =
+        num_classes_per_output_from_py(numClasses, yArma.n_rows);
 
     const arma::Row<float> wRow =
         sampleWeightRowFromPy(sample_weights, xArma.n_cols);
@@ -161,7 +266,7 @@ public:
     arma::uvec featuresMut = featuresArma;
     std::visit(
         [&](auto &d) {
-          d.Train(xArma, featuresMut, yArma, numClasses, minLeafSize,
+          d.Train(xArma, featuresMut, yArma, nClassesPerOutput, minLeafSize,
                   minGainSplit, maxDepth, maxLeafNodes, wRow);
         },
         impl_);
@@ -198,15 +303,9 @@ public:
         impl_);
   }
 
-  py::array_t<size_t> getBinPredictions() {
+  py::array getBinPredictions() {
     return std::visit(
-        [](auto &d) {
-          const auto &preds = d.getBinPredictions();
-          arma::Col<size_t> col(preds.size());
-          for (size_t i = 0; i < preds.size(); ++i)
-            col(i) = preds[i];
-          return col_to_numpy_1d(col);
-        },
+        [](auto &d) { return bin_predictions_to_numpy(d.getBinPredictions()); },
         impl_);
   }
 
@@ -252,12 +351,7 @@ public:
     const arma::uvec featuresArma =
         arma::conv_to<arma::uvec>::from(carma::arr_to_col<size_t>(fcopy, true));
 
-    py::array_t<float> ycopy = py::array_t<float>::ensure(y);
-    if (ycopy.ndim() != 1)
-      throw std::invalid_argument("y must be a 1D numpy array");
-    arma::Mat<float> yArma = carma::arr_to_row<float>(ycopy, true);
-    if (yArma.n_elem != xArma.n_cols)
-      throw std::invalid_argument("y length must match X.shape[0]");
+    arma::Mat<float> yArma = regression_y_to_arma(y, xArma.n_cols);
 
     const arma::Row<float> wRow =
         sampleWeightRowFromPy(sample_weights, xArma.n_cols);
@@ -302,11 +396,9 @@ public:
         impl_);
   }
 
-  py::array_t<float> getBinPredictions() {
+  py::array getBinPredictions() {
     return std::visit(
-        [](auto &d) {
-          return vector_float_to_numpy_1d(d.getBinPredictions());
-        },
+        [](auto &d) { return bin_predictions_to_numpy(d.getBinPredictions()); },
         impl_);
   }
 
@@ -336,7 +428,7 @@ public:
   }
 
   void Train(const py::array_t<float> &X, const py::array_t<size_t> &features,
-             const py::array_t<size_t> &y, size_t numClasses,
+             const py::array_t<size_t> &y, py::object numClasses,
              size_t minLeafSize, double minGainSplit, size_t maxDepth,
              size_t maxLeafNodes, py::object sample_weights = py::none()) {
     py::array_t<float> Xcopy = py::array_t<float>::ensure(X);
@@ -352,19 +444,16 @@ public:
     const arma::uvec featuresArma =
         arma::conv_to<arma::uvec>::from(carma::arr_to_col<size_t>(fcopy, true));
 
-    py::array_t<size_t> ycopy = py::array_t<size_t>::ensure(y);
-    if (ycopy.ndim() != 1)
-      throw std::invalid_argument("y must be a 1D numpy array");
-    arma::Mat<size_t> yArma = carma::arr_to_row<size_t>(ycopy, true);
-    if (yArma.n_elem != xArma.n_cols)
-      throw std::invalid_argument("y length must match X.shape[0]");
+    arma::Mat<size_t> yArma = classification_y_to_arma(y, xArma.n_cols);
+    const std::vector<size_t> nClassesPerOutput =
+        num_classes_per_output_from_py(numClasses, yArma.n_rows);
 
     const arma::Row<float> wRow =
         sampleWeightRowFromPy(sample_weights, xArma.n_cols);
 
     arma::uvec featuresMut = featuresArma;
-    impl_.Train(xArma, featuresMut, yArma, numClasses, minLeafSize, minGainSplit,
-                maxDepth, maxLeafNodes, wRow);
+    impl_.Train(xArma, featuresMut, yArma, nClassesPerOutput, minLeafSize,
+                minGainSplit, maxDepth, maxLeafNodes, wRow);
   }
 
   py::array_t<size_t> transform(const py::array_t<float> &X) {
@@ -390,12 +479,8 @@ public:
     return VectorOfVectorsToNumpyList(impl_.inSampleDiscretizations());
   }
 
-  py::array_t<size_t> getBinPredictions() {
-    const auto &preds = impl_.getBinPredictions();
-    arma::Col<size_t> col(preds.size());
-    for (size_t i = 0; i < preds.size(); ++i)
-      col(i) = preds[i];
-    return col_to_numpy_1d(col);
+  py::array getBinPredictions() {
+    return bin_predictions_to_numpy(impl_.getBinPredictions());
   }
 
   size_t getNumLeaves() const { return impl_.numLeaves(); }
@@ -436,12 +521,7 @@ public:
     const arma::uvec featuresArma =
         arma::conv_to<arma::uvec>::from(carma::arr_to_col<size_t>(fcopy, true));
 
-    py::array_t<float> ycopy = py::array_t<float>::ensure(y);
-    if (ycopy.ndim() != 1)
-      throw std::invalid_argument("y must be a 1D numpy array");
-    arma::Mat<float> yArma = carma::arr_to_row<float>(ycopy, true);
-    if (yArma.n_elem != xArma.n_cols)
-      throw std::invalid_argument("y length must match X.shape[0]");
+    arma::Mat<float> yArma = regression_y_to_arma(y, xArma.n_cols);
 
     const arma::Row<float> wRow =
         sampleWeightRowFromPy(sample_weights, xArma.n_cols);
@@ -474,8 +554,8 @@ public:
     return VectorOfVectorsToNumpyList(impl_.inSampleDiscretizations());
   }
 
-  py::array_t<float> getBinPredictions() {
-    return vector_float_to_numpy_1d(impl_.getBinPredictions());
+  py::array getBinPredictions() {
+    return bin_predictions_to_numpy(impl_.getBinPredictions());
   }
 
   size_t getNumLeaves() const { return impl_.numLeaves(); }
@@ -486,6 +566,16 @@ public:
 } // namespace
 
 PYBIND11_MODULE(Discretizers, m) {
+  // CARMA's allocator may lazily call _import_array() on free/alloc; prime it
+  // here (GIL held) so destructions during ShapeGeneralizedTrees.fit (which can
+  // resolve to this module's statically-linked core symbols) stay safe.
+  if (_import_array() < 0) {
+    PyErr_Clear();
+    throw std::runtime_error(
+        "Discretizers: numpy.core.multiarray failed to import; "
+        "ensure numpy is installed and importable before importing this module");
+  }
+
   py::class_<UnivariateClassificationDiscretizerPy>(
       m, "UnivariateClassificationDiscretizer")
       .def(py::init<std::string>(), py::arg("criterion") = "gini")

{sgtlearn-0.2.0 → sgtlearn-0.3.0}/cpp/bindings/ShapeGeneralizedTrees.cpp

@@ -43,7 +43,7 @@ PYBIND11_MODULE(ShapeGeneralizedTrees, m) {
 
   py::class_<ClassificationShapeGeneralizedTreePy>(
       m, "ClassificationShapeGeneralizedTree")
-      .def(py::init([](std::string criterion, size_t num_classes,
+      .def(py::init([](std::string criterion, py::object num_classes,
                        size_t num_partitions, size_t outer_min_leaf_size,
                        double outer_min_gain_split, size_t outer_max_depth,
                        size_t outer_max_leaf_nodes, size_t inner_min_leaf_size,
@@ -52,15 +52,17 @@ PYBIND11_MODULE(ShapeGeneralizedTrees, m) {
                        size_t coordinate_descent_max_iters,
                        size_t coordinate_descent_patience,
                        bool coordinate_descent_smart_init, uint64_t random_state,
-                       py::object max_features) {
+                       py::object max_features, size_t pairwise_candidates,
+                       double pairwise_penalty) {
              return ClassificationShapeGeneralizedTreePy(
-                 std::move(criterion), num_classes, num_partitions,
+                 std::move(criterion), std::move(num_classes), num_partitions,
                  outer_min_leaf_size, outer_min_gain_split, outer_max_depth,
                  outer_max_leaf_nodes, inner_min_leaf_size, inner_min_gain_split,
                  inner_max_depth, inner_max_leaf_nodes,
                  coordinate_descent_max_iters, coordinate_descent_patience,
                  coordinate_descent_smart_init, random_state,
-                 std::move(max_features));
+                 std::move(max_features), pairwise_candidates,
+                 pairwise_penalty);
            }),
            py::arg("criterion") = "gini", py::arg("num_classes"),
            py::arg("num_partitions") = 2,
@@ -76,25 +78,38 @@ PYBIND11_MODULE(ShapeGeneralizedTrees, m) {
            py::arg("coordinate_descent_patience") = 5,
            py::arg("coordinate_descent_smart_init") = true,
            py::arg("random_state") = 42,
-           py::arg("max_features") = py::none())
+           py::arg("max_features") = py::none(),
+           py::arg("pairwise_candidates") = 0,
+           py::arg("pairwise_penalty") = 0.0)
       .def("fit", &ClassificationShapeGeneralizedTreePy::fit, py::arg("X"),
            py::arg("y"), py::arg("sample_weight") = py::none(),
            py::arg("features"),
            "Fit the routing tree. X is (n_samples, n_features) float32; y is "
-           "1-D uint class labels. Optional sample_weight is 1-D float32.")
+           "uint class labels, 1-D (n_samples,) or 2-D (n_samples, n_outputs). "
+           "Optional sample_weight is 1-D float32.")
       .def("predict", &ClassificationShapeGeneralizedTreePy::predict,
            py::arg("X"),
-           "Predict class labels for X (shape (n_samples, n_features)).")
+           "Predict class labels for X (shape (n_samples, n_features)). Returns "
+           "(n_samples,) for a single output or (n_samples, n_outputs) for "
+           "multi-output.")
       .def("predict_proba", &ClassificationShapeGeneralizedTreePy::predictProba,
            py::arg("X"),
-           "Predict class probabilities for X; output shape "
-           "(n_samples, n_classes).")
+           "Predict class probabilities for X. Single output: a "
+           "(n_samples, n_classes) array. Multi-output: a list of such arrays, "
+           "one per output.")
       .def_property_readonly(
           "num_leaves", &ClassificationShapeGeneralizedTreePy::numLeaves)
       .def_property_readonly(
           "num_nodes", &ClassificationShapeGeneralizedTreePy::numNodes)
+      .def_property_readonly(
+          "n_outputs", &ClassificationShapeGeneralizedTreePy::nOutputs)
+      .def_property_readonly(
+          "classes_per_output",
+          &ClassificationShapeGeneralizedTreePy::classesPerOutput)
       .def_property_readonly(
           "is_fitted", &ClassificationShapeGeneralizedTreePy::isFitted)
+      .def_property_readonly(
+          "has_pair_nodes", &ClassificationShapeGeneralizedTreePy::hasPairNodes)
       .def_property_readonly(
           "feature_importance",
           &ClassificationShapeGeneralizedTreePy::featureImportance,
@@ -113,14 +128,16 @@ PYBIND11_MODULE(ShapeGeneralizedTrees, m) {
                        size_t coordinate_descent_max_iters,
                        size_t coordinate_descent_patience,
                        bool coordinate_descent_smart_init, uint64_t random_state,
-                       py::object max_features) {
+                       py::object max_features, size_t pairwise_candidates,
+                       double pairwise_penalty) {
              return RegressionShapeGeneralizedTreePy(
                  std::move(criterion), num_partitions, outer_min_leaf_size,
                  outer_min_gain_split, outer_max_depth, outer_max_leaf_nodes,
                  inner_min_leaf_size, inner_min_gain_split, inner_max_depth,
                  inner_max_leaf_nodes, coordinate_descent_max_iters,
                  coordinate_descent_patience, coordinate_descent_smart_init,
-                 random_state, std::move(max_features));
+                 random_state, std::move(max_features), pairwise_candidates,
+                 pairwise_penalty);
            }),
            py::arg("criterion") = "squared_error",
            py::arg("num_partitions") = 2, py::arg("outer_min_leaf_size") = 1,
@@ -133,6 +150,8 @@ PYBIND11_MODULE(ShapeGeneralizedTrees, m) {
            py::arg("coordinate_descent_patience") = 5,
            py::arg("coordinate_descent_smart_init") = true,
            py::arg("random_state") = 42, py::arg("max_features") = py::none(),
+           py::arg("pairwise_candidates") = 0,
+           py::arg("pairwise_penalty") = 0.0,
            R"(Regression tree: inner bins are round-robin seeded. ``squared_error`` runs
 coordinate descent and keeps the map only if branch MSE improves clearly vs the seed;
 otherwise the snapshot is restored and the branch objective is rebuilt.
@@ -142,15 +161,21 @@ is accepted for API parity with ClassificationShapeGeneralizedTree but ignored.)
            py::arg("y"), py::arg("sample_weight") = py::none(),
            py::arg("features"),
            "Fit the routing tree. X is (n_samples, n_features) float32; y is "
-           "1-D float32 targets. Optional sample_weight is 1-D float32.")
+           "float32 targets, 1-D (n_samples,) or 2-D (n_samples, n_outputs). "
+           "Optional sample_weight is 1-D float32.")
       .def("predict", &RegressionShapeGeneralizedTreePy::predict, py::arg("X"),
-           "Predict scalar targets for X (shape (n_samples,)).")
+           "Predict targets for X. Returns (n_samples,) for a single output or "
+           "(n_samples, n_outputs) for multi-output.")
       .def_property_readonly("num_leaves",
                              &RegressionShapeGeneralizedTreePy::numLeaves)
       .def_property_readonly("num_nodes",
                              &RegressionShapeGeneralizedTreePy::numNodes)
+      .def_property_readonly("n_outputs",
+                             &RegressionShapeGeneralizedTreePy::nOutputs)
       .def_property_readonly("is_fitted",
                              &RegressionShapeGeneralizedTreePy::isFitted)
+      .def_property_readonly(
+          "has_pair_nodes", &RegressionShapeGeneralizedTreePy::hasPairNodes)
       .def_property_readonly(
           "feature_importance",
           &RegressionShapeGeneralizedTreePy::featureImportance,