sleipnirgroup-jormungandr 0.6.7.dev55__tar.gz → 0.6.7.dev63__tar.gz

{sleipnirgroup_jormungandr-0.6.7.dev55 → sleipnirgroup_jormungandr-0.6.7.dev63}/CMakeLists.txt

@@ -301,7 +301,9 @@ if(SLEIPNIR_BUILD_PYTHON)
     if(DEFINED PY_BUILD_CMAKE_MODULE_NAME)
         set(PY_DEST ${PY_BUILD_CMAKE_MODULE_NAME})
     else()
-        set(PY_DEST lib/python${Python_VERSION_MAJOR}.${Python_VERSION_MINOR})
+        set(PY_DEST
+            lib/python${Python_VERSION_MAJOR}.${Python_VERSION_MINOR}/sleipnir
+        )
     endif()
 
     # nanobind dependency
@@ -343,47 +345,70 @@ if(SLEIPNIR_BUILD_PYTHON)
         LIBRARY DESTINATION ${PY_DEST}
     )
 
+    # py-build-cmake copies the pure Python sources into the wheel itself.
+    # Install them for plain CMake installs as well.
+    if(NOT DEFINED PY_BUILD_CMAKE_MODULE_NAME)
+        install(
+            DIRECTORY python/src/sleipnir/
+            DESTINATION ${PY_DEST}
+            COMPONENT python_modules
+            FILES_MATCHING
+            PATTERN "*.py"
+            PATTERN "__pycache__" EXCLUDE
+        )
+    endif()
+
+    # Generate stubs at install time because only the install tree has the
+    # sleipnir/ package layout needed to import the module by its fully
+    # qualified name `sleipnir._sleipnir`.
+    #
+    # Importing sleipnir._sleipnir requires the directory containing the
+    # sleipnir package (i.e., PY_DEST's parent) to be on the Python path. For
+    # py-build-cmake builds, PY_DEST has no parent. For plain CMake installs,
+    # PY_DEST's parent is lib/pythonX.Y.
+    cmake_path(GET PY_DEST PARENT_PATH PY_DEST_PARENT)
     nanobind_add_stub(
         _sleipnir_stub
         INSTALL_TIME
-        MARKER_FILE sleipnir/py.typed
-        MODULE _sleipnir
-        OUTPUT sleipnir/__init__.pyi
-        PYTHON_PATH $<TARGET_FILE_DIR:_sleipnir>
+        MARKER_FILE ${PY_DEST}/py.typed
+        MODULE sleipnir._sleipnir
+        OUTPUT ${PY_DEST}/_sleipnir/__init__.pyi
+        PYTHON_PATH ${PY_DEST_PARENT}
         DEPENDS _sleipnir
         COMPONENT python_modules
     )
     nanobind_add_stub(
         _sleipnir_autodiff_stub
         INSTALL_TIME
-        MODULE _sleipnir.autodiff
-        OUTPUT sleipnir/autodiff/__init__.pyi
-        PYTHON_PATH $<TARGET_FILE_DIR:_sleipnir>
+        MODULE sleipnir._sleipnir.autodiff
+        OUTPUT ${PY_DEST}/_sleipnir/autodiff.pyi
+        PYTHON_PATH ${PY_DEST_PARENT}
         DEPENDS _sleipnir
         COMPONENT python_modules
     )
     nanobind_add_stub(
         _sleipnir_optimization_stub
         INSTALL_TIME
-        MODULE _sleipnir.optimization
-        OUTPUT sleipnir/optimization/__init__.pyi
-        PYTHON_PATH $<TARGET_FILE_DIR:_sleipnir>
+        MODULE sleipnir._sleipnir.optimization
+        OUTPUT ${PY_DEST}/_sleipnir/optimization.pyi
+        PYTHON_PATH ${PY_DEST_PARENT}
         DEPENDS _sleipnir
         COMPONENT python_modules
     )
+    install(
+        CODE
+            "execute_process(
+                COMMAND \"${Python_EXECUTABLE}\"
+                    \"${CMAKE_CURRENT_SOURCE_DIR}/python/postprocess_stubs.py\"
+                    ${PY_DEST}/_sleipnir/autodiff.pyi
+                    ${PY_DEST}/_sleipnir/optimization.pyi
+                WORKING_DIRECTORY \"\${CMAKE_INSTALL_PREFIX}\"
+                COMMAND_ERROR_IS_FATAL ANY
+            )"
+        COMPONENT python_modules
+    )
 
     if(NOT CMAKE_CROSSCOMPILING)
-        # pybind11_mkdoc dependency
-        FetchContent_Declare(
-            pybind11_mkdoc
-            GIT_REPOSITORY https://github.com/pybind/pybind11_mkdoc.git
-            GIT_TAG v3.0.0
-            GIT_SHALLOW ON
-            EXCLUDE_FROM_ALL
-            SYSTEM
-        )
-        FetchContent_MakeAvailable(pybind11_mkdoc)
-
         set(sleipnir_headers
             ${CMAKE_CURRENT_SOURCE_DIR}/include/sleipnir/autodiff/expression_type.hpp
             ${CMAKE_CURRENT_SOURCE_DIR}/include/sleipnir/autodiff/gradient.hpp

{sleipnirgroup_jormungandr-0.6.7.dev55 → sleipnirgroup_jormungandr-0.6.7.dev63}/PKG-INFO

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

{sleipnirgroup_jormungandr-0.6.7.dev55 → sleipnirgroup_jormungandr-0.6.7.dev63}/cmake/Pybind11Mkdoc.cmake

@@ -1,6 +1,20 @@
 function(pybind11_mkdoc target headers)
     find_package(Python3 REQUIRED COMPONENTS Interpreter)
 
+    # Check that pybind11_mkdoc is installed
+    execute_process(
+        COMMAND ${Python3_EXECUTABLE} -c "import pybind11_mkdoc"
+        RESULT_VARIABLE pybind11_mkdoc_result
+        OUTPUT_QUIET
+        ERROR_QUIET
+    )
+    if(NOT pybind11_mkdoc_result EQUAL 0)
+        message(
+            FATAL_ERROR
+            "pybind11_mkdoc not found for ${Python3_EXECUTABLE}. Install it with `pip install --group docstrings`."
+        )
+    endif()
+
     # Target compiler flags
     get_target_property(target_flags ${target} INCLUDE_DIRECTORIES)
     list(TRANSFORM target_flags PREPEND "-I")
@@ -60,7 +74,13 @@ function(pybind11_mkdoc target headers)
             ${CMAKE_CURRENT_SOURCE_DIR}/python/cpp/docstrings.hpp
             ${target_flags} ${clang_flags} ${eigen_flags} ${small_vector_flags}
             -std=c++23
-        DEPENDS ${headers}
+        COMMAND
+            ${CMAKE_COMMAND}
+            -DFILE=${CMAKE_CURRENT_SOURCE_DIR}/python/cpp/docstrings.hpp -P
+            ${CMAKE_CURRENT_SOURCE_DIR}/cmake/Pybind11MkdocPostprocess.cmake
+        DEPENDS
+            ${headers}
+            ${CMAKE_CURRENT_SOURCE_DIR}/cmake/Pybind11MkdocPostprocess.cmake
     )
 
     # Add docstrings target

sleipnirgroup_jormungandr-0.6.7.dev63/cmake/Pybind11MkdocPostprocess.cmake

@@ -0,0 +1,12 @@
+# Replaces pybind11_mkdoc's docstring pointers with constexpr arrays
+#
+# Usage: cmake -DFILE=<docstrings.hpp> -P Pybind11MkdocPostprocess.cmake
+
+file(READ ${FILE} contents)
+string(
+    REGEX REPLACE "\nstatic const char \\*([A-Za-z0-9_]+) ="
+    "\nconstexpr char \\1[] ="
+    contents
+    "${contents}"
+)
+file(WRITE ${FILE} "${contents}")

{sleipnirgroup_jormungandr-0.6.7.dev55 → sleipnirgroup_jormungandr-0.6.7.dev63}/pyproject.toml

@@ -21,8 +21,11 @@ build-backend = "py_build_cmake.build"
 
 [dependency-groups]
 dev = [
+    "mypy>=2.4.0",
     "pytest>=9.1.1",
+    "scipy-stubs",
 ]
+docstrings = ["pybind11-mkdoc~=3.0.0"]
 
 [tool.py-build-cmake.cmake]
 build_type = "Release"
@@ -54,8 +57,29 @@ exclude = []
 python_abi = "abi3"
 abi3_minimum_cpython_version = 312
 
+[tool.mypy]
+files = [
+  "benchmarks",
+  "docs/generate_website.py",
+  "examples",
+  "python/postprocess_stubs.py",
+  "python/test",
+  "tools",
+]
+allow_redefinition_new = true
+disallow_incomplete_defs = true
+disallow_untyped_defs = true
+explicit_package_bases = true
+local_partial_types = true
+mypy_path = ["examples/current_manager", "python/test/optimization"]
+
+  [[tool.mypy.overrides]]
+  module = "mpl_toolkits.*"
+  ignore_missing_imports = true
+
 [tool.pytest.ini_options]
 minversion = "6.0"
+pythonpath = ["examples/current_manager"]
 testpaths = ["python/test", "examples/current_manager/test"]
 
 [tool.uv]
@@ -66,8 +90,11 @@ cache-keys = [
     { file = "include/**" },
     { file = "pyproject.toml" },
     { file = "python/cpp/**/*.cpp" },
+    { file = "python/cpp/arity.hpp" },
     { file = "python/cpp/binders.hpp" },
     { file = "python/cpp/for_each_type.hpp" },
+    { file = "python/cpp/slice_tuple.hpp" },
+    { file = "python/cpp/text_wrap.hpp" },
     { file = "python/cpp/try_cast.hpp" },
     { file = "python/src/**" },
     { file = "src/**" },

sleipnirgroup_jormungandr-0.6.7.dev63/python/cpp/arity.hpp

@@ -0,0 +1,42 @@
+// Copyright (c) Sleipnir contributors
+
+#pragma once
+
+#include <cstddef>
+#include <type_traits>
+
+namespace slp {
+
+/// Type implicitly convertible to any type. Used for counting aggregate
+/// members.
+struct AnyType {
+  template <typename T>
+  // NOLINTNEXTLINE (google-explicit-constructor)
+  operator T() const;
+};
+
+/// Returns the number of members in aggregate type T.
+///
+/// @tparam T Aggregate type.
+/// @tparam Args Initializers for the members counted so far.
+template <typename T, typename... Args>
+  requires std::is_aggregate_v<T>
+consteval size_t member_count() {
+  if constexpr (requires { T{Args{}..., AnyType{}}; }) {
+    return member_count<T, Args..., AnyType>();
+  } else {
+    return sizeof...(Args);
+  }
+}
+
+/// Returns the number of parameters of a lambda.
+///
+/// @tparam F Lambda type.
+template <typename F>
+consteval size_t lambda_param_count() {
+  return []<typename R, typename C, typename... Args>(R (C::*)(Args...) const) {
+    return sizeof...(Args);
+  }(&F::operator());
+}
+
+}  // namespace slp

{sleipnirgroup_jormungandr-0.6.7.dev55 → sleipnirgroup_jormungandr-0.6.7.dev63}/python/cpp/autodiff/bind_variable_block.cpp

@@ -6,6 +6,7 @@
 #include <nanobind/operators.h>
 #include <nanobind/stl/function.h>
 #include <nanobind/stl/string.h>
+#include <nanobind/stl/tuple.h>
 #include <nanobind/stl/vector.h>
 
 #include <stdint.h>
@@ -15,6 +16,7 @@
 #include <functional>
 #include <stdexcept>
 #include <string>
+#include <tuple>
 
 #include <Eigen/Core>
 #include <sleipnir/autodiff/slice.hpp>
@@ -25,6 +27,7 @@
 
 #include "docstrings.hpp"
 #include "for_each_type.hpp"
+#include "slice_tuple.hpp"
 #include "try_cast.hpp"
 
 namespace nb = nanobind;
@@ -68,7 +71,7 @@ void bind_variable_block(
       "row"_a, "value"_a);
   cls.def(
       "__setitem__",
-      [](VariableBlock<VariableMatrix<double>>& self, nb::tuple slices,
+      [](VariableBlock<VariableMatrix<double>>& self, SliceTuple slices,
          nb::object value) {
         if (slices.size() != 2) {
           throw nb::index_error(
@@ -151,31 +154,32 @@ void bind_variable_block(
   cls.def(
       "__getitem__",
       [](VariableBlock<VariableMatrix<double>>& self,
-         nb::tuple slices) -> nb::object {
+         std::tuple<int, int> slices) -> Variable<double>& {
+        auto [row, col] = slices;
+
+        if (row >= self.rows() || col >= self.cols()) {
+          throw std::out_of_range("Index out of bounds");
+        }
+
+        if (row < 0) {
+          row += self.rows();
+        }
+        if (col < 0) {
+          col += self.cols();
+        }
+        return self[row, col];
+      },
+      nb::rv_policy::reference_internal, "slices"_a,
+      DOC(slp, VariableBlock, operator, array));
+  cls.def(
+      "__getitem__",
+      [](VariableBlock<VariableMatrix<double>>& self,
+         SliceTuple slices) -> VariableBlock<VariableMatrix<double>> {
         if (slices.size() != 2) {
           throw nb::index_error(
               std::format("Expected 2 slices, got {}.", slices.size()).c_str());
         }
 
-        // If both indices are integers instead of slices, return Variable
-        // instead of VariableBlock
-        if (nb::isinstance<int>(slices[0]) && nb::isinstance<int>(slices[1])) {
-          int row = nb::cast<int>(slices[0]);
-          int col = nb::cast<int>(slices[1]);
-
-          if (row >= self.rows() || col >= self.cols()) {
-            throw std::out_of_range("Index out of bounds");
-          }
-
-          if (row < 0) {
-            row += self.rows();
-          }
-          if (col < 0) {
-            col += self.cols();
-          }
-          return nb::cast(self[row, col]);
-        }
-
         Slice row_slice;
         int row_slice_length;
         Slice col_slice;
@@ -211,8 +215,7 @@ void bind_variable_block(
           col_slice_length = 1;
         }
 
-        return nb::cast(
-            self[row_slice, row_slice_length, col_slice, col_slice_length]);
+        return self[row_slice, row_slice_length, col_slice, col_slice_length];
       },
       nb::keep_alive<0, 1>(), DOC(slp, VariableBlock, operator, array),
       "slices"_a);
@@ -223,6 +226,31 @@ void bind_variable_block(
           nb::overload_cast<int>(&VariableBlock<VariableMatrix<double>>::col),
           "col"_a, DOC(slp, VariableBlock, col));
 
+  // Matrix-matrix multiplication. Products with a VariableMatrix on the left
+  // are bound on VariableMatrix.
+  for (const char* name : {"__mul__", "__matmul__"}) {
+    cls.def(
+        name,
+        [](const VariableBlock<VariableMatrix<double>>& lhs,
+           const VariableBlock<VariableMatrix<double>>& rhs) {
+          return lhs * rhs;
+        },
+        nb::is_operator(), "rhs"_a);
+  }
+
+  // Matrix multiplication with NumPy arrays. These are also handled by
+  // __array_ufunc__ below, but type checkers only see these overloads.
+  cls.def(
+      "__matmul__",
+      [](const VariableBlock<VariableMatrix<double>>& lhs,
+         nb::DRef<Eigen::MatrixXd> rhs) { return lhs * rhs; },
+      nb::is_operator(), "rhs"_a);
+  cls.def(
+      "__rmatmul__",
+      [](const VariableBlock<VariableMatrix<double>>& rhs,
+         nb::DRef<Eigen::MatrixXd> lhs) { return lhs * rhs; },
+      nb::is_operator(), "lhs"_a);
+
   // https://numpy.org/doc/stable/user/basics.dispatch.html
   cls.def(
       "__array_ufunc__",
@@ -302,6 +330,10 @@ void bind_variable_block(
 
         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(nb::self + nb::self, "rhs"_a);

{sleipnirgroup_jormungandr-0.6.7.dev55 → sleipnirgroup_jormungandr-0.6.7.dev63}/python/cpp/autodiff/bind_variable_matrix.cpp

@@ -6,6 +6,7 @@
 #include <nanobind/operators.h>
 #include <nanobind/stl/function.h>
 #include <nanobind/stl/string.h>
+#include <nanobind/stl/tuple.h>
 #include <nanobind/stl/vector.h>
 
 #include <stdint.h>
@@ -15,6 +16,7 @@
 #include <functional>
 #include <stdexcept>
 #include <string>
+#include <tuple>
 #include <vector>
 
 #include <Eigen/Core>
@@ -26,6 +28,7 @@
 
 #include "docstrings.hpp"
 #include "for_each_type.hpp"
+#include "slice_tuple.hpp"
 #include "try_cast.hpp"
 
 namespace nb = nanobind;
@@ -90,7 +93,7 @@ void bind_variable_matrix(nb::module_& autodiff,
       "row"_a, "value"_a);
   cls.def(
       "__setitem__",
-      [](VariableMatrix<double>& self, nb::tuple slices, nb::object value) {
+      [](VariableMatrix<double>& self, SliceTuple slices, nb::object value) {
         if (slices.size() != 2) {
           throw nb::index_error(
               std::format("Expected 2 slices, got {}.", slices.size()).c_str());
@@ -170,31 +173,33 @@ void bind_variable_matrix(nb::module_& autodiff,
       DOC(slp, VariableMatrix, operator, array, 3));
   cls.def(
       "__getitem__",
-      [](VariableMatrix<double>& self, nb::tuple slices) -> nb::object {
+      [](VariableMatrix<double>& self,
+         std::tuple<int, int> slices) -> Variable<double>& {
+        auto [row, col] = slices;
+
+        if (row >= self.rows() || col >= self.cols()) {
+          throw std::out_of_range("Index out of bounds");
+        }
+
+        if (row < 0) {
+          row += self.rows();
+        }
+        if (col < 0) {
+          col += self.cols();
+        }
+        return self[row, col];
+      },
+      nb::rv_policy::reference_internal, "slices"_a,
+      DOC(slp, VariableMatrix, operator, array));
+  cls.def(
+      "__getitem__",
+      [](VariableMatrix<double>& self,
+         SliceTuple slices) -> VariableBlock<VariableMatrix<double>> {
         if (slices.size() != 2) {
           throw nb::index_error(
               std::format("Expected 2 slices, got {}.", slices.size()).c_str());
         }
 
-        // If both indices are integers instead of slices, return Variable
-        // instead of VariableBlock
-        if (nb::isinstance<int>(slices[0]) && nb::isinstance<int>(slices[1])) {
-          int row = nb::cast<int>(slices[0]);
-          int col = nb::cast<int>(slices[1]);
-
-          if (row >= self.rows() || col >= self.cols()) {
-            throw std::out_of_range("Index out of bounds");
-          }
-
-          if (row < 0) {
-            row += self.rows();
-          }
-          if (col < 0) {
-            col += self.cols();
-          }
-          return nb::cast(self[row, col]);
-        }
-
         Slice row_slice;
         int row_slice_length;
         Slice col_slice;
@@ -227,8 +232,7 @@ void bind_variable_matrix(nb::module_& autodiff,
           col_slice_length = 1;
         }
 
-        return nb::cast(
-            self[row_slice, row_slice_length, col_slice, col_slice_length]);
+        return self[row_slice, row_slice_length, col_slice, col_slice_length];
       },
       nb::keep_alive<0, 1>(), "slices"_a,
       DOC(slp, VariableMatrix, operator, array));
@@ -237,43 +241,46 @@ void bind_variable_matrix(nb::module_& autodiff,
   cls.def("col", nb::overload_cast<int>(&VariableMatrix<double>::col), "col"_a,
           DOC(slp, VariableMatrix, col));
 
-  // Matrix-matrix multiplication
-  cls.def(
-      "__mul__",
-      [](const VariableMatrix<double>& lhs, const VariableMatrix<double>& rhs) {
-        return lhs * rhs;
-      },
-      nb::is_operator(), "rhs"_a);
+  // Matrix-matrix multiplication. Products with a VariableBlock on the left are
+  // bound on VariableBlock.
+  for (const char* name : {"__mul__", "__matmul__"}) {
+    cls.def(
+        name,
+        [](const VariableMatrix<double>& lhs,
+           const VariableMatrix<double>& rhs) { return lhs * rhs; },
+        nb::is_operator(), "rhs"_a);
+    cls.def(
+        name,
+        [](const VariableMatrix<double>& lhs,
+           const VariableBlock<VariableMatrix<double>>& rhs) {
+          return lhs * rhs;
+        },
+        nb::is_operator(), "rhs"_a);
+  }
+  for (const char* name : {"__rmul__", "__rmatmul__"}) {
+    cls.def(
+        name,
+        [](const VariableMatrix<double>& rhs,
+           const VariableBlock<VariableMatrix<double>>& lhs) {
+          return lhs * rhs;
+        },
+        nb::is_operator(), "lhs"_a);
+  }
+
+  // Matrix multiplication with NumPy arrays. These are also handled by
+  // __array_ufunc__ below, but type checkers only see these overloads.
   cls.def(
-      "__mul__",
-      [](const VariableMatrix<double>& lhs,
-         const VariableBlock<VariableMatrix<double>>& rhs) {
+      "__matmul__",
+      [](const VariableMatrix<double>& lhs, nb::DRef<Eigen::MatrixXd> rhs) {
         return lhs * rhs;
       },
       nb::is_operator(), "rhs"_a);
   cls.def(
-      "__rmul__",
-      [](const VariableMatrix<double>& rhs,
-         const VariableBlock<VariableMatrix<double>>& lhs) {
+      "__rmatmul__",
+      [](const VariableMatrix<double>& rhs, nb::DRef<Eigen::MatrixXd> lhs) {
         return lhs * rhs;
       },
       nb::is_operator(), "lhs"_a);
-  cls.def(
-      "__mul__",
-      [](const VariableBlock<VariableMatrix<double>>& lhs,
-         const VariableBlock<VariableMatrix<double>>& rhs) {
-        return lhs * rhs;
-      },
-      nb::is_operator(), "rhs"_a);
-  for_each_type<VariableMatrix<double>, VariableBlock<VariableMatrix<double>>>(
-      [&]<typename L> {
-        for_each_type<VariableMatrix<double>,
-                      VariableBlock<VariableMatrix<double>>>([&]<typename R> {
-          cls.def(
-              "__matmul__", [](L&& lhs, R&& rhs) { return lhs * rhs; },
-              nb::is_operator(), "rhs"_a);
-        });
-      });
 
   // https://numpy.org/doc/stable/user/basics.dispatch.html
   cls.def(
@@ -354,6 +361,10 @@ void bind_variable_matrix(nb::module_& autodiff,
 
         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(nb::self + nb::self, "rhs"_a);