sleipnirgroup-jormungandr 0.6.7.dev128__tar.gz → 0.6.7.dev130__tar.gz

{sleipnirgroup_jormungandr-0.6.7.dev128 → sleipnirgroup_jormungandr-0.6.7.dev130}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sleipnirgroup-jormungandr
-Version: 0.6.7.dev128
+Version: 0.6.7.dev130
 Summary: Reverse mode autodiff library and NLP solver DSL
 License-Expression: BSD-3-Clause
 License-File: LICENSE.txt

{sleipnirgroup_jormungandr-0.6.7.dev128 → sleipnirgroup_jormungandr-0.6.7.dev130}/include/sleipnir/autodiff/variable_matrix.hpp

@@ -9,6 +9,7 @@
 #include <initializer_list>
 #include <iterator>
 #include <span>
+#include <type_traits>
 #include <utility>
 #include <vector>
 
@@ -566,24 +567,6 @@ class VariableMatrix : public SleipnirBase {
     return result;
   }
 
-  /// Matrix-scalar multiplication operator.
-  ///
-  /// @param lhs Operator left-hand side.
-  /// @param rhs Operator right-hand side.
-  template <EigenMatrixLike LHS>
-  friend VariableMatrix<Scalar> operator*(const LHS& lhs,
-                                          const Variable<Scalar>& rhs) {
-    VariableMatrix<Scalar> result(detail::empty, lhs.rows(), lhs.cols());
-
-    for (int row = 0; row < result.rows(); ++row) {
-      for (int col = 0; col < result.cols(); ++col) {
-        result[row, col] = lhs[row, col] * rhs;
-      }
-    }
-
-    return result;
-  }
-
   /// Matrix-scalar multiplication operator.
   ///
   /// @param lhs Operator left-hand side.
@@ -601,24 +584,6 @@ class VariableMatrix : public SleipnirBase {
     return result;
   }
 
-  /// Scalar-matrix multiplication operator.
-  ///
-  /// @param lhs Operator left-hand side.
-  /// @param rhs Operator right-hand side.
-  template <EigenMatrixLike RHS>
-  friend VariableMatrix<Scalar> operator*(const Variable<Scalar>& lhs,
-                                          const RHS& rhs) {
-    VariableMatrix<Scalar> result(detail::empty, rhs.rows(), rhs.cols());
-
-    for (int row = 0; row < result.rows(); ++row) {
-      for (int col = 0; col < result.cols(); ++col) {
-        result[row, col] = rhs[row, col] * lhs;
-      }
-    }
-
-    return result;
-  }
-
   /// Scalar-matrix multiplication operator.
   ///
   /// @param lhs Operator left-hand side.
@@ -671,25 +636,6 @@ class VariableMatrix : public SleipnirBase {
     return *this;
   }
 
-  /// Binary division operator.
-  ///
-  /// @param lhs Operator left-hand side.
-  /// @param rhs Operator right-hand side.
-  /// @return Result of division.
-  template <EigenMatrixLike LHS>
-  friend VariableMatrix<Scalar> operator/(const LHS& lhs,
-                                          const Variable<Scalar>& rhs) {
-    VariableMatrix<Scalar> result(detail::empty, lhs.rows(), lhs.cols());
-
-    for (int row = 0; row < result.rows(); ++row) {
-      for (int col = 0; col < result.cols(); ++col) {
-        result[row, col] = lhs[row, col] / rhs;
-      }
-    }
-
-    return result;
-  }
-
   /// Binary division operator.
   ///
   /// @param lhs Operator left-hand side.
@@ -1373,6 +1319,64 @@ template <typename Derived>
 VariableMatrix(const Eigen::DiagonalBase<Derived>&)
     -> VariableMatrix<typename Derived::Scalar>;
 
+/// Matrix-scalar multiplication operator.
+///
+/// @tparam Scalar Scalar type.
+/// @param lhs Operator left-hand side.
+/// @param rhs Operator right-hand side.
+template <typename Scalar, EigenMatrixLike LHS>
+  requires std::same_as<typename std::decay_t<LHS>::Scalar, Scalar>
+VariableMatrix<Scalar> operator*(const LHS& lhs, const Variable<Scalar>& rhs) {
+  VariableMatrix<Scalar> result(detail::empty, lhs.rows(), lhs.cols());
+
+  for (int row = 0; row < result.rows(); ++row) {
+    for (int col = 0; col < result.cols(); ++col) {
+      result[row, col] = lhs[row, col] * rhs;
+    }
+  }
+
+  return result;
+}
+
+/// Scalar-matrix multiplication operator.
+///
+/// @tparam Scalar Scalar type.
+/// @param lhs Operator left-hand side.
+/// @param rhs Operator right-hand side.
+template <typename Scalar, EigenMatrixLike RHS>
+  requires std::same_as<typename std::decay_t<RHS>::Scalar, Scalar>
+VariableMatrix<Scalar> operator*(const Variable<Scalar>& lhs, const RHS& rhs) {
+  VariableMatrix<Scalar> result(detail::empty, rhs.rows(), rhs.cols());
+
+  for (int row = 0; row < result.rows(); ++row) {
+    for (int col = 0; col < result.cols(); ++col) {
+      result[row, col] = rhs[row, col] * lhs;
+    }
+  }
+
+  return result;
+}
+
+/// Binary division operator.
+///
+/// @tparam Scalar Scalar type.
+/// @param lhs Operator left-hand side.
+/// @param rhs Operator right-hand side.
+/// @return Result of division.
+template <typename Scalar, EigenMatrixLike LHS>
+  requires std::same_as<typename std::decay_t<LHS>::Scalar, Scalar>
+VariableMatrix<Scalar> operator/(const LHS& lhs, const Variable<Scalar>& rhs) {
+  VariableMatrix<Scalar> result(detail::empty, lhs.rows(), lhs.cols());
+
+  for (int row = 0; row < result.rows(); ++row) {
+    for (int col = 0; col < result.cols(); ++col) {
+      result[row, col] = lhs[row, col] / rhs;
+    }
+  }
+
+  return result;
+}
+
 /// Applies a coefficient-wise reduce operation to two matrices.
 ///
 /// @tparam Scalar Scalar type.

{sleipnirgroup_jormungandr-0.6.7.dev128 → sleipnirgroup_jormungandr-0.6.7.dev130}/python/cpp/autodiff/bind_variable.cpp

@@ -3,8 +3,12 @@
 #include <nanobind/eigen/dense.h>
 #include <nanobind/nanobind.h>
 #include <nanobind/operators.h>
+#include <nanobind/stl/string.h>
 
+#include <algorithm>
+#include <array>
 #include <concepts>
+#include <string>
 #include <utility>
 
 #include <sleipnir/autodiff/variable.hpp>
@@ -13,6 +17,7 @@
 
 #include "docstrings.hpp"
 #include "for_each_type.hpp"
+#include "try_cast.hpp"
 
 namespace nb = nanobind;
 
@@ -93,7 +98,8 @@ void bind_variable(nb::module_& autodiff, nb::class_<Variable<double>>& cls) {
   using namespace nb::literals;
 
   cls.def(nb::init<>(), DOC(slp, Variable, Variable));
-  cls.def(nb::init<double>(), "value"_a, DOC(slp, Variable, Variable, 2));
+  cls.def(nb::init_implicit<double>(), "value"_a,
+          DOC(slp, Variable, Variable, 2));
   cls.def("set_value", nb::overload_cast<double>(&Variable<double>::set_value),
           "value"_a, DOC(slp, Variable, set_value));
   cls.def("value", &Variable<double>::value, DOC(slp, Variable, value));
@@ -116,6 +122,142 @@ void bind_variable(nb::module_& autodiff, nb::class_<Variable<double>>& cls) {
     }
   });
 
+  // Scalar-matrix multiplication and division with NumPy arrays. These are also
+  // handled by __array_ufunc__ below, but type checkers only see these
+  // overloads.
+  cls.def(
+      "__mul__",
+      [](const Variable<double>& lhs, nb::DRef<Eigen::MatrixXd> rhs) {
+        return lhs * rhs;
+      },
+      nb::is_operator(), "rhs"_a);
+  cls.def(
+      "__imul__",
+      [](const Variable<double>& lhs, nb::DRef<Eigen::MatrixXd> rhs) {
+        return lhs * rhs;
+      },
+      nb::is_operator(), "rhs"_a);
+  cls.def(
+      "__rmul__",
+      [](const Variable<double>& rhs, nb::DRef<Eigen::MatrixXd> lhs) {
+        return lhs * rhs;
+      },
+      nb::is_operator(), "lhs"_a);
+  cls.def(
+      "__rtruediv__",
+      [](const Variable<double>& rhs, nb::DRef<Eigen::MatrixXd> lhs) {
+        return lhs / rhs;
+      },
+      nb::is_operator(), "lhs"_a);
+
+  // Map NumPy ufunc reprs to the equivalent function names in Python's
+  // operator module
+  static constexpr std::array scalar_operators{
+      std::pair{"<ufunc 'add'>", "add"},
+      std::pair{"<ufunc 'subtract'>", "sub"},
+      std::pair{"<ufunc 'multiply'>", "mul"},
+      std::pair{"<ufunc 'divide'>", "truediv"},
+      std::pair{"<ufunc 'equal'>", "eq"},
+      std::pair{"<ufunc 'less'>", "lt"},
+      std::pair{"<ufunc 'less_equal'>", "le"},
+      std::pair{"<ufunc 'greater'>", "gt"},
+      std::pair{"<ufunc 'greater_equal'>", "ge"}};
+
+  // https://numpy.org/doc/stable/user/basics.dispatch.html
+  cls.def(
+      "__array_ufunc__",
+      [](Variable<double>& self, nb::object ufunc, nb::str method,
+         nb::args inputs, const nb::kwargs&) -> nb::object {
+        std::string method_name = nb::cast<std::string>(method);
+        std::string ufunc_name =
+            nb::cast<std::string>(ufunc.attr("__repr__")());
+
+        // VariableMatrix only represents 2D arrays
+        auto is_2d_array = [](const nb::handle& obj) {
+          return nb::isinstance<nb::ndarray<>>(obj) &&
+                 nb::cast<nb::ndarray<>>(obj).ndim() == 2;
+        };
+
+        if (method_name == "__call__" && inputs.size() == 2) {
+          if (auto op = std::ranges::find(
+                  scalar_operators, ufunc_name,
+                  [](const auto& pair) { return pair.first; });
+              op != scalar_operators.end()) {
+            auto np = nb::module_::import_("numpy");
+
+            auto is_0d = [&](const nb::handle& obj) {
+              // np.ndim() returns 0 for any non-array object
+              return !nb::isinstance<Variable<double>>(obj) &&
+                     nb::cast<int>(np.attr("ndim")(obj)) == 0;
+            };
+
+            // Convert NumPy scalars and 0D arrays to the equivalent Python
+            // scalar (e.g., np.float64 to float) and return other objects
+            // unchanged
+            auto numpy_to_python_scalar =
+                [&](const nb::handle& obj) -> nb::object {
+              nb::object numpy_types =
+                  nb::make_tuple(np.attr("generic"), np.attr("ndarray"));
+              if (PyObject_IsInstance(obj.ptr(), numpy_types.ptr()) == 1) {
+                return obj.attr("item")();
+              } else {
+                return nb::borrow(obj);
+              }
+            };
+
+            // If an operand is 0D, apply the Python operator to its Python
+            // scalar and the other operand, which dispatches to Variable's
+            // scalar overloads
+            if (is_0d(inputs[0])) {
+              return nb::module_::import_("operator")
+                  .attr(op->second)(numpy_to_python_scalar(inputs[0]),
+                                    inputs[1]);
+            } else if (is_0d(inputs[1])) {
+              return nb::module_::import_("operator")
+                  .attr(op->second)(inputs[0],
+                                    numpy_to_python_scalar(inputs[1]));
+            }
+          }
+
+          if (ufunc_name == "<ufunc 'multiply'>") {
+            if (is_2d_array(inputs[0])) {
+              if (auto result = apply_eigen_op(
+                      [&](auto& lhs) {
+                        return lhs.template cast<double>() * self;
+                      },
+                      inputs[0])) {
+                return result.value();
+              }
+            } else if (is_2d_array(inputs[1])) {
+              if (auto result = apply_eigen_op(
+                      [&](auto& rhs) {
+                        return self * rhs.template cast<double>();
+                      },
+                      inputs[1])) {
+                return result.value();
+              }
+            }
+          } else if (ufunc_name == "<ufunc 'divide'>") {
+            if (is_2d_array(inputs[0])) {
+              if (auto result = apply_eigen_op(
+                      [&](auto& lhs) {
+                        return lhs.template cast<double>() / self;
+                      },
+                      inputs[0])) {
+                return result.value();
+              }
+            }
+          }
+        }
+
+        return nb::not_implemented();
+      },
+      // nanobind can't annotate nb::args or nb::kwargs, so the stub needs an
+      // explicit signature
+      nb::sig("def __array_ufunc__(self, ufunc: object, method: str, "
+              "*inputs: object, **kwargs: object) -> object"),
+      "ufunc"_a, "method"_a, "inputs"_a, "kwargs"_a);
+
   cls.def(
       "__pow__",
       [](const Variable<double>& self, int power) { return pow(self, power); },

{sleipnirgroup_jormungandr-0.6.7.dev128 → sleipnirgroup_jormungandr-0.6.7.dev130}/python/cpp/autodiff/bind_variable_block.cpp

@@ -340,12 +340,11 @@ void bind_variable_block(
   cls.def(-nb::self);
 
   // Matrix-scalar/scalar-matrix operations
-  for_each_type<double, int, Variable<double>>([&]<typename T> {
-    cls.def(nb::self * T(), "rhs"_a);
-    cls.def(T() * nb::self, "lhs"_a);
-    cls.def(nb::self / T(), "rhs"_a);
-    cls.def(nb::self /= T(), "rhs"_a, DOC(slp, VariableBlock, operator, idiv));
-  });
+  cls.def(nb::self * Variable<double>(), "rhs"_a);
+  cls.def(Variable<double>() * nb::self, "lhs"_a);
+  cls.def(nb::self / Variable<double>(), "rhs"_a);
+  cls.def(nb::self /= Variable<double>(), "rhs"_a,
+          DOC(slp, VariableBlock, operator, idiv));
   for_each_type<double, int, const Variable<double>&,
                 nb::DRef<Eigen::MatrixXd>>([&]<typename T> {
     cls.def(

{sleipnirgroup_jormungandr-0.6.7.dev128 → sleipnirgroup_jormungandr-0.6.7.dev130}/python/cpp/autodiff/bind_variable_matrix.cpp

@@ -372,12 +372,11 @@ void bind_variable_matrix(nb::module_& autodiff,
   cls.def(-nb::self);
 
   // Matrix-scalar/scalar-matrix operations
-  for_each_type<double, int, Variable<double>>([&]<typename T> {
-    cls.def(nb::self * T(), "rhs"_a);
-    cls.def(T() * nb::self, "lhs"_a);
-    cls.def(nb::self / T(), "rhs"_a);
-    cls.def(nb::self /= T(), "rhs"_a, DOC(slp, VariableMatrix, operator, idiv));
-  });
+  cls.def(nb::self * Variable<double>(), "rhs"_a);
+  cls.def(Variable<double>() * nb::self, "lhs"_a);
+  cls.def(nb::self / Variable<double>(), "rhs"_a);
+  cls.def(nb::self /= Variable<double>(), "rhs"_a,
+          DOC(slp, VariableMatrix, operator, idiv));
   for_each_type<double, int, const Variable<double>&,
                 const VariableBlock<VariableMatrix<double>>,
                 nb::DRef<Eigen::MatrixXd>>([&]<typename T> {
@@ -529,6 +528,18 @@ void bind_variable_matrix(nb::module_& autodiff,
                nb::overload_cast<const VariableMatrix<double>&,
                                  const VariableMatrix<double>&>(&solve<double>),
                "A"_a, "B"_a, DOC(slp, solve));
+  autodiff.def(
+      "solve",
+      [](const VariableMatrix<double>& A, nb::DRef<Eigen::MatrixXd> B) {
+        return solve(A, VariableMatrix<double>(B));
+      },
+      "A"_a, "B"_a, DOC(slp, solve));
+  autodiff.def(
+      "solve",
+      [](nb::DRef<Eigen::MatrixXd> A, const VariableMatrix<double>& B) {
+        return solve(VariableMatrix<double>(A), B);
+      },
+      "A"_a, "B"_a, DOC(slp, solve));
 }
 
 }  // namespace slp

{sleipnirgroup_jormungandr-0.6.7.dev128 → sleipnirgroup_jormungandr-0.6.7.dev130}/python/cpp/docstrings.hpp

@@ -2959,6 +2959,21 @@ Template Args:
 
 )doc";
 
+constexpr char mkd_doc_slp_operator_div[] =
+R"doc(Binary division operator.
+
+Args:
+    lhs: Operator left-hand side.
+    rhs: Operator right-hand side.
+
+Template Args:
+    Scalar: Scalar type.
+
+Returns:
+    Result of division.
+
+)doc";
+
 constexpr char mkd_doc_slp_operator_eq[] =
 R"doc(Equality operator that returns an equality constraint for two
 Variables.
@@ -3049,6 +3064,30 @@ Args:
 
 )doc";
 
+constexpr char mkd_doc_slp_operator_mul[] =
+R"doc(Matrix-scalar multiplication operator.
+
+Args:
+    lhs: Operator left-hand side.
+    rhs: Operator right-hand side.
+
+Template Args:
+    Scalar: Scalar type.
+
+)doc";
+
+constexpr char mkd_doc_slp_operator_mul_2[] =
+R"doc(Scalar-matrix multiplication operator.
+
+Args:
+    lhs: Operator left-hand side.
+    rhs: Operator right-hand side.
+
+Template Args:
+    Scalar: Scalar type.
+
+)doc";
+
 constexpr char mkd_doc_slp_pow[] =
 R"doc(pow() for Variables.
 

{sleipnirgroup_jormungandr-0.6.7.dev128 → sleipnirgroup_jormungandr-0.6.7.dev130}/python/postprocess_stubs.py

@@ -5,9 +5,11 @@ Fixes up type stubs generated by nanobind's stubgen.
 
 stubgen can't express some things the bindings support at runtime:
 
-1. VariableBlock implicitly converts to VariableMatrix, so every VariableMatrix
-   parameter also accepts a VariableBlock. Return types and callback arguments
-   are left alone since those are always VariableMatrix.
+1. Implicit conversions widen what parameters accept: VariableBlock converts
+   to VariableMatrix, and float converts to Variable. Every VariableMatrix
+   parameter also accepts a VariableBlock, and every Variable parameter also
+   accepts a float. Return types and callback arguments are left alone since
+   those are always VariableMatrix or Variable.
 2. Sparse matrices are annotated as scipy.sparse.csc_matrix[float], but
    scipy-stubs expects a NumPy numeric scalar type argument.
 """
@@ -15,6 +17,7 @@ stubgen can't express some things the bindings support at runtime:
 import argparse
 import ast
 import re
+from dataclasses import dataclass
 from pathlib import Path
 
 PREFIX = r"(?:sleipnir\._sleipnir\.autodiff\.)?"
@@ -26,14 +29,53 @@ PREFIX = r"(?:sleipnir\._sleipnir\.autodiff\.)?"
 #   -> x: VariableMatrix | VariableBlock | None
 VARIABLE_MATRIX = re.compile(
     rf"""
-    (?<![\w.])                                  # Not part of another name
-    ({PREFIX})                                  # Optional module prefix
-    VariableMatrix\b                            # Type name
-    (?!\ \|\ {PREFIX}VariableBlock\b)           # Not already widened
+    (?<![\w.])                         # Not part of another name
+    ({PREFIX})                         # Optional module prefix
+    VariableMatrix\b                   # Type name
+    (?!\ \|\ {PREFIX}VariableBlock\b)  # Not already widened
     """,
     re.VERBOSE,
 )
 
+# Variable, optionally module-qualified. The word boundary excludes
+# VariableMatrix and VariableBlock, and the lookahead prevents already-widened
+# annotations from matching again.
+#
+#   x: Sequence[Variable]
+#   -> x: Sequence[Variable | float]
+VARIABLE = re.compile(
+    rf"""
+    (?<![\w.])         # Not part of another name
+    ({PREFIX})         # Optional module prefix
+    Variable\b         # Type name
+    (?!\ \|\ float\b)  # Not already widened
+    """,
+    re.VERBOSE,
+)
+
+
+@dataclass(frozen=True)
+class Widening:
+    """A type that also accepts another type via implicit conversion."""
+
+    # Matches the narrow type
+    pattern: re.Pattern[str]
+
+    # Replaces the narrow type with the type that converts to it
+    source: str
+
+    # Replaces the narrow type with the widened type
+    widened: str
+
+
+# \1 is "sleipnir._sleipnir.autodiff." prefix or empty
+WIDENINGS = [
+    Widening(
+        VARIABLE_MATRIX, r"\1VariableBlock", r"\1VariableMatrix | \1VariableBlock"
+    ),
+    Widening(VARIABLE, "float", r"\1Variable | float"),
+]
+
 
 def is_callable(annotation: ast.expr) -> bool:
     """Returns True if an annotation contains a Callable type."""
@@ -44,6 +86,17 @@ def is_callable(annotation: ast.expr) -> bool:
     )
 
 
+def is_nested(annotation: ast.expr, pattern: re.Pattern[str]) -> bool:
+    """
+    Returns True if a type matching the pattern appears inside a subscript of
+    an annotation (e.g., Variable in Sequence[Variable]).
+    """
+    return any(
+        isinstance(node, ast.Subscript) and pattern.search(ast.unparse(node.slice))
+        for node in ast.walk(annotation)
+    )
+
+
 def parameters(node: ast.FunctionDef) -> list[ast.arg | None]:
     """Returns a function's parameters in signature order."""
     args = node.args
@@ -68,11 +121,13 @@ def signature(node: ast.FunctionDef) -> list[str | None]:
 def widen_signatures(stub: str) -> str:
     """
     Widens every VariableMatrix in function parameter annotations to
-    VariableMatrix | VariableBlock.
+    VariableMatrix | VariableBlock, and every Variable to Variable | float.
 
     A parameter isn't widened if a sibling overload is identical except it
-    takes VariableBlock there; widening it would make that overload
-    unreachable.
+    takes the source type there (e.g., float instead of Variable); widening it
+    would make that overload unreachable. Types nested in a subscript (e.g.,
+    Sequence[Variable]) are always widened, since the sibling overloads don't
+    accept containers that mix both types.
     """
     # Group overloads by name within each module or class body
     overload_sets: dict[tuple[int, str], list[ast.FunctionDef]] = {}
@@ -81,25 +136,35 @@ def widen_signatures(stub: str) -> str:
             if isinstance(child, ast.FunctionDef):
                 overload_sets.setdefault((id(node), child.name), []).append(child)
 
-    annotations = []
+    # Annotations to widen, and the widenings that apply to each
+    annotations: list[tuple[ast.expr, list[Widening]]] = []
     for overloads in overload_sets.values():
         signatures = [signature(overload) for overload in overloads]
         for overload, sig in zip(overloads, signatures):
             for i, arg in enumerate(parameters(overload)):
-                # Callbacks receive VariableMatrix, so their annotations are
-                # left alone
+                # Callbacks receive VariableMatrix or Variable, so their
+                # annotations are left alone
                 if arg is None or arg.annotation is None or is_callable(arg.annotation):
                     continue
 
-                # Skip if a sibling overload already takes VariableBlock here
-                block_sig = sig.copy()
-                block_sig[i] = VARIABLE_MATRIX.sub(
-                    r"\1VariableBlock", ast.unparse(arg.annotation)
-                )
-                if block_sig != sig and block_sig in signatures:
-                    continue
-
-                annotations.append(arg.annotation)
+                widenings = []
+                for widening in WIDENINGS:
+                    # Skip if a sibling overload already takes the source type
+                    # here, unless the type is nested in a subscript
+                    source_sig = sig.copy()
+                    source_sig[i] = widening.pattern.sub(
+                        widening.source, ast.unparse(arg.annotation)
+                    )
+                    if (
+                        source_sig != sig
+                        and source_sig in signatures
+                        and not is_nested(arg.annotation, widening.pattern)
+                    ):
+                        continue
+
+                    widenings.append(widening)
+
+                annotations.append((arg.annotation, widenings))
 
     # Encode stub string into UTF-8 bytes
     source_bytes = stub.encode()
@@ -110,8 +175,10 @@ def widen_signatures(stub: str) -> str:
         line_byte_offsets.append(line_byte_offsets[-1] + len(line))
 
     # Edit back to front so earlier offsets stay valid
-    for annotation in sorted(
-        annotations, key=lambda node: (node.lineno, node.col_offset), reverse=True
+    for annotation, widenings in sorted(
+        annotations,
+        key=lambda item: (item[0].lineno, item[0].col_offset),
+        reverse=True,
     ):
         # Assertions to satisfy mypy. They can't fail in practice because
         # ast.parse() always sets those fields.
@@ -123,11 +190,9 @@ def widen_signatures(stub: str) -> str:
         end = line_byte_offsets[annotation.end_lineno - 1] + annotation.end_col_offset
 
         # Splice widened type annotations into source_bytes
-        #
-        # \1 is "sleipnir._sleipnir.autodiff." prefix or empty
-        widened = VARIABLE_MATRIX.sub(
-            r"\1VariableMatrix | \1VariableBlock", source_bytes[start:end].decode()
-        )
+        widened = source_bytes[start:end].decode()
+        for widening in widenings:
+            widened = widening.pattern.sub(widening.widened, widened)
         source_bytes = source_bytes[:start] + widened.encode() + source_bytes[end:]
     return source_bytes.decode()
 

{sleipnirgroup_jormungandr-0.6.7.dev128 → sleipnirgroup_jormungandr-0.6.7.dev130}/python/src/sleipnir/__init__.py

@@ -2,4 +2,4 @@
 
 from ._sleipnir import *
 
-__version__ = "0.6.7.dev128"
+__version__ = "0.6.7.dev130"