Add support for Eigen's Tensor module

cmake/nanobind-config.cmake

@@ -192,6 +192,7 @@ function (nanobind_build_library TARGET_NAME)
     ${NB_DIR}/include/nanobind/stl/vector.h
     ${NB_DIR}/include/nanobind/eigen/dense.h
     ${NB_DIR}/include/nanobind/eigen/sparse.h
+    ${NB_DIR}/include/nanobind/eigen/tensor.h
 
     ${NB_DIR}/src/buffer.h
     ${NB_DIR}/src/hash.h

docs/exchanging.rst

@@ -143,6 +143,8 @@ to external projects that provide further casters:
     - ``#include <nanobind/eigen/dense.h>``
   * - ``Eigen::SparseMatrix<..>``
     - ``#include <nanobind/eigen/sparse.h>``
+  * - ``Eigen::Tensor<..>``, ``Eigen::TensorMap<..>``, ``Eigen::TensorRef<..>``
+    - ``#include <nanobind/eigen/tensor.h>``
   * - Apache Arrow types
     -  `https://github.com/maximiliank/nanobind_pyarrow <https://github.com/maximiliank/nanobind_pyarrow>`__
   * - ...

include/nanobind/eigen/tensor.h

@@ -0,0 +1,324 @@
+/*
+    nanobind/eigen/tensor.h: type casters for Eigen tensors
+
+    Copyright (c) 2026 INRIA
+
+    Author(s): Wilson Jallet
+
+    All rights reserved. Use of this source code is governed by a
+    BSD-style license that can be found in the LICENSE file.
+*/
+
+#pragma once
+
+#include <nanobind/ndarray.h>
+
+#include <unsupported/Eigen/CXX11/Tensor>
+
+
+NAMESPACE_BEGIN(NB_NAMESPACE)
+
+NAMESPACE_BEGIN(detail)
+
+/// As of April 2026, Eigen::Tensor types support 16-byte alignment or no alignment.
+inline bool is_tensor_aligned(const void *data, std::size_t align = Eigen::Aligned) {
+    return (reinterpret_cast<std::size_t>(data) % align) == 0;
+}
+
+/// Type trait for inheriting from Eigen::TensorBase.
+/// All TensorBase specializations inherit from TensorBase<T, ReadOnlyAccessors>.
+template<typename T> constexpr bool is_eigen_tensor_v = 
+    std::is_base_of_v<Eigen::TensorBase<T, Eigen::ReadOnlyAccessors>, T>;
+
+template<typename T> constexpr bool eigen_tensor_is_row_major_v = T::Layout == Eigen::RowMajor;
+template<typename T> constexpr bool eigen_tensor_is_col_major_v = T::Layout == Eigen::ColMajor;
+
+template<typename T>
+constexpr bool is_eigen_tensor_map_v = false;
+
+// Covers const case
+template<typename T, int Options, template<class> class MakePointer>
+constexpr bool is_eigen_tensor_map_v<Eigen::TensorMap<T, Options, MakePointer>> = true;
+
+template<typename T>
+constexpr bool is_eigen_tensor_ref_v = false;
+
+// Covers const case
+template<typename T>
+constexpr bool is_eigen_tensor_ref_v<Eigen::TensorRef<T>> = true;
+
+template<typename T>
+constexpr bool is_eigen_tensor_plain_v = false;
+
+template<typename Scalar, int NumIndices, int Options, typename IndexType>
+constexpr bool is_eigen_tensor_plain_v<Eigen::Tensor<Scalar, NumIndices, Options, IndexType>>
+    = true;
+
+template<typename Scalar, std::ptrdiff_t... Indices, int Options, typename IndexType>
+constexpr bool is_eigen_tensor_plain_v<Eigen::TensorFixedSize<Scalar, Eigen::Sizes<Indices...>, Options, IndexType>> = true;
+
+template<typename T>
+constexpr bool is_eigen_tensor_xpr_v = 
+    is_eigen_tensor_v<T> &&
+    !is_eigen_tensor_plain_v<T> &&
+    !is_eigen_tensor_map_v<T> &&
+    !is_eigen_tensor_ref_v<T>;
+
+template<typename T, typename Scalar = typename T::Scalar>
+using ndarray_for_eigen_tensor_t = ndarray<
+    Scalar,
+    numpy,
+    ndim<T::NumDimensions>, 
+    std::conditional_t<
+        eigen_tensor_is_row_major_v<T>,
+        c_contig,
+        f_contig>>;
+
+/** \brief Type caster for ``Eigen::TensorMap<T>``
+ */
+template<typename T, int MapOptions, template<class> class MakePointer>
+struct type_caster<
+    Eigen::TensorMap<T, MapOptions, MakePointer>,
+    enable_if_t<is_ndarray_scalar_v<typename T::Scalar>>> {
+
+    using Scalar = typename T::Scalar;
+    using IndexType = typename T::Index;
+    static constexpr int NumIndices = T::NumIndices;
+    static constexpr int Options = T::Options;
+    using PlainTensor = Eigen::Tensor<Scalar, NumIndices, Options, IndexType>;
+    using Dimensions = typename T::Dimensions;
+    using MapType = Eigen::TensorMap<T, MapOptions, MakePointer>;
+    static constexpr bool IsAligned = MapType::IsAligned;
+
+    // Only partial specification. Dimensions not known at compile time...
+    using NDArray =
+        ndarray_for_eigen_tensor_t<T, std::conditional_t<std::is_const_v<T>,
+                                                         const Scalar,
+                                                         Scalar>>;
+    using NDArrayCaster = make_caster<NDArray>;
+    static constexpr auto Name = NDArrayCaster::Name;
+    template<typename T_> using Cast = MapType;
+    template<typename T_> static constexpr bool can_cast() { return true; };
+
+    NDArrayCaster caster;
+
+    bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
+        // Disable implicit conversions
+        flags &= ~(uint8_t)cast_flags::convert;
+        // Do not accept None
+        flags &= ~(uint8_t)cast_flags::accepts_none;
+
+        if (!caster.from_python(src, flags, cleanup))
+            return false;
+        if(IsAligned && !is_tensor_aligned(caster.value.data()))
+            return false;
+
+        return true;
+    }
+
+    static handle from_cpp(const MapType &v, rv_policy policy, cleanup_list *cleanup) noexcept {
+        size_t shape[NumIndices];
+        for (size_t i = 0 ; i < NumIndices; i++) {
+            shape[i] = static_cast<size_t>(v.dimension(i));
+        }
+
+        void* ptr = (void *)v.data();
+        if (policy == rv_policy::automatic || policy == rv_policy::automatic_reference)
+            policy = rv_policy::reference;
+        return NDArrayCaster::from_cpp(
+            NDArray {ptr, NumIndices, shape, handle()},
+            policy,
+            cleanup);
+    }
+
+    operator MapType() {
+        NDArray &t = caster.value;
+        std::array<long, NumIndices> shape;
+        for (size_t i = 0 ; i < NumIndices; i++) {
+            shape[i] = t.shape(i);
+        }
+        return MapType(t.data(), shape);
+    }
+};
+
+
+/** \brief Type caster for plain ``Eigen::Tensor<T>`` types.
+ */
+template<typename Scalar, int NumIndices, int Options, typename IndexType>
+struct type_caster<
+        Eigen::Tensor<Scalar, NumIndices, Options, IndexType>,
+        enable_if_t<is_ndarray_scalar_v<Scalar>>> {
+
+    using PlainTensor = Eigen::Tensor<Scalar, NumIndices, Options, IndexType>;
+    using Dimensions = typename PlainTensor::Dimensions;
+    using Coeffs = typename PlainTensor::CoeffReturnType;
+    static constexpr bool IsRowMajor = bool(Options & Eigen::RowMajorBit);
+    using NDArray = ndarray_for_eigen_tensor_t<PlainTensor>;
+    using NDArrayCaster = make_caster<NDArray>;
+
+    // PlainTensor value;
+    NB_TYPE_CASTER(PlainTensor, NDArrayCaster::Name);
+
+    bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
+        using NDArrayConst = ndarray_for_eigen_tensor_t<PlainTensor, const Scalar>;
+        make_caster<NDArrayConst> caster;
+        // Do not accept None
+        if (!caster.from_python(src, flags & ~(uint8_t)cast_flags::accepts_none, cleanup))
+            return false;
+
+        const NDArrayConst &array = caster.value;
+        // copy tensor dims
+        std::array<long, NumIndices> out_dims;
+        for(size_t i = 0; i < NumIndices; i++) {
+            out_dims[i] = array.shape(i);
+        }
+        value.resize(out_dims);
+
+        memcpy(value.data(), array.data(), array.size() * sizeof(Scalar));
+
+        return true;
+    }
+
+    static handle from_cpp(const PlainTensor &v, rv_policy policy, cleanup_list *cleanup) noexcept {
+        size_t shape[NumIndices];
+
+        for (size_t i = 0 ; i < NumIndices; i++) {
+            shape[i] = static_cast<size_t>(v.dimension(i));
+        }
+
+        void *ptr = (void *)v.data();
+
+        object owner;
+        if (policy == rv_policy::move) {
+            PlainTensor *tmp = new PlainTensor((PlainTensor&&)v);
+            owner = capsule(tmp, [](void* p) noexcept {
+                delete (PlainTensor*) p;
+            });
+            ptr = tmp->data();
+            policy = rv_policy::reference;
+        } else if (policy == rv_policy::reference_internal && cleanup->self()) {
+            owner = borrow(cleanup->self());
+            policy = rv_policy::reference;
+        }
+        return NDArrayCaster::from_cpp(
+            NDArray {ptr, NumIndices, shape, owner},
+            policy, cleanup);
+    }
+};
+
+/** \brief Type caster for Tensor expressions. From-cpp conversion just converts the expression to a plain Tensor object.
+ */
+template<typename T>
+struct type_caster<T, enable_if_t<is_eigen_tensor_xpr_v<T> && is_ndarray_scalar_v<typename T::Scalar>>> {
+    static constexpr int NumDimensions = T::NumDimensions;
+    static constexpr int Options = T::Options;
+    using IndexType = typename T::Index;
+    using XprTraits = typename Eigen::internal::traits<T>;
+    static constexpr int Layout = XprTraits::Layout;
+    using PlainTensor = Eigen::Tensor<typename T::Scalar, NumDimensions, Layout, IndexType>;
+    using Caster = make_caster<PlainTensor>;
+    static constexpr auto Name = Caster::Name;
+    template<typename T_> using Cast = T;
+    template<typename T_> static constexpr bool can_cast() { return true; }
+
+    /// Generating an expression template from a Python object is impossible.
+    bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept = delete;
+
+    template <typename T2>
+    static handle from_cpp(T2 &&v, rv_policy policy, cleanup_list *cleanup) noexcept {
+        return Caster::from_cpp(std::forward<T2>(v), policy, cleanup);
+    }
+};
+
+/** \brief Type caster for ``Eigen::TensorRef<T>``
+ */
+template<typename T>
+struct type_caster<
+    Eigen::TensorRef<T>,
+    enable_if_t<is_ndarray_scalar_v<typename T::Scalar>>> {
+
+    using Scalar = typename T::Scalar;
+    using IndexType = typename T::Index;
+    static constexpr int NumIndices = T::NumIndices;
+    static constexpr int Options = T::Options;
+    using PlainTensor = Eigen::Tensor<Scalar, NumIndices, Options, IndexType>;
+    using Dimensions = typename T::Dimensions;
+
+    // Only partial specification. Dimensions not known at compile time...
+    using NDArray =
+        ndarray_for_eigen_tensor_t<T, std::conditional_t<std::is_const_v<T>,
+                                                         const Scalar,
+                                                         Scalar>>;
+    using NDArrayCaster = make_caster<NDArray>;
+
+    using MapType = Eigen::TensorMap<T>;
+    using MapCaster = make_caster<MapType>;
+
+    using RefType = Eigen::TensorRef<T>;
+
+
+    static constexpr bool MaybeConvert = std::is_const_v<T>;
+    using PlainCaster = make_caster<PlainTensor>;
+
+    static constexpr auto Name = const_name<MaybeConvert>(PlainCaster::Name, MapCaster::Name);
+
+    template<typename T_> using Cast = RefType;
+    template<typename T_> static constexpr bool can_cast() { return true; };
+
+    MapCaster caster;
+    struct Empty {};
+    std::conditional_t<MaybeConvert, PlainCaster, Empty> plain_caster;
+
+    bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
+        // no conversion for mutable Ref
+        if constexpr (!std::is_const_v<T>)
+            flags &= ~(uint8_t) cast_flags::convert;
+
+        // Try direct cast
+        if (caster.from_python(src, flags, cleanup))
+            return true;
+
+        // if const T, attempt leveraging PlainTensor conversion
+        if constexpr (MaybeConvert) {
+            // we create a new temporary tensor object, and
+            // its lifetime is that of plain_caster.
+            // for manual conversion, disable conversion.
+            if ((flags & (uint8_t) cast_flags::manual))
+                flags &= ~(uint8_t) cast_flags::convert;
+            if (plain_caster.from_python(src, flags, cleanup))
+                return true;
+        }
+
+        return false;
+    }
+
+    static handle from_cpp(const RefType &v, rv_policy policy, cleanup_list *cleanup) noexcept {
+        size_t shape[NumIndices];
+
+        for (size_t i = 0; i < NumIndices; i++) {
+            shape[i] = static_cast<size_t>(v.dimension(i));
+        }
+
+        return NDArrayCaster::from_cpp(
+            NDArray((void *) v.data(), NumIndices, shape, handle()),
+            (policy == rv_policy::automatic ||
+             policy == rv_policy::automatic_reference)
+                ? rv_policy::reference
+                : policy,
+            cleanup);
+    }
+
+    operator RefType() {
+        if constexpr (MaybeConvert) {
+            // if there's a value, return it
+            if (plain_caster.caster.value.is_valid())
+                return RefType(plain_caster.operator PlainTensor&());
+        }
+        return RefType(caster.operator MapType());
+    }
+};
+
+
+NAMESPACE_END(detail)
+
+NAMESPACE_END(NB_NAMESPACE)

tests/CMakeLists.txt

@@ -134,6 +134,8 @@ find_package (Eigen3 3.3.1 NO_MODULE)
 if (TARGET Eigen3::Eigen)
   nanobind_add_module(test_eigen_ext test_eigen.cpp ${NB_EXTRA_ARGS})
   target_link_libraries(test_eigen_ext PRIVATE Eigen3::Eigen)
+  nanobind_add_module(test_eigen_tensor_ext test_eigen_tensor.cpp ${NB_EXTRA_ARGS})
+  target_link_libraries(test_eigen_tensor_ext PRIVATE Eigen3::Eigen)
 endif()
 
 add_library(
@@ -159,6 +161,7 @@ set(TEST_FILES
   test_callbacks.py
   test_classes.py
   test_eigen.py
+  test_eigen_tensor.py
   test_enum.py
   test_eval.py
   test_exception.py