Merge pull request #933 from IntelPython/cleanup-tensor-step4

[cleanup/tensor, part 4] Made contract_iter and contract_iter2 functions templated to avoid needing to include pybind11 headers

Author: oleksandr-pavlyk

dpctl/tensor/libtensor/include/utils/strided_iters.hpp

@@ -499,16 +499,16 @@ int simplify_iteration_two_strides(const int nd,
     return nd_;
 }
 
-using vecT = std::vector<py::ssize_t>;
-std::tuple<vecT, vecT, py::size_t> contract_iter(vecT shape, vecT strides)
+template <typename T, class Error, typename vecT = std::vector<T>>
+std::tuple<vecT, vecT, T> contract_iter(vecT shape, vecT strides)
 {
     const size_t dim = shape.size();
     if (dim != strides.size()) {
-        throw py::value_error("Shape and strides must be of equal size.");
+        throw Error("Shape and strides must be of equal size.");
     }
     vecT out_shape = shape;
     vecT out_strides = strides;
-    py::ssize_t disp(0);
+    T disp(0);
 
     int nd = simplify_iteration_stride(dim, out_shape.data(),
                                        out_strides.data(), disp);
@@ -517,18 +517,19 @@ std::tuple<vecT, vecT, py::size_t> contract_iter(vecT shape, vecT strides)
     return std::make_tuple(out_shape, out_strides, disp);
 }
 
-std::tuple<vecT, vecT, py::size_t, vecT, py::ssize_t>
+template <typename T, class Error, typename vecT = std::vector<T>>
+std::tuple<vecT, vecT, T, vecT, T>
 contract_iter2(vecT shape, vecT strides1, vecT strides2)
 {
     const size_t dim = shape.size();
     if (dim != strides1.size() || dim != strides2.size()) {
-        throw py::value_error("Shape and strides must be of equal size.");
+        throw Error("Shape and strides must be of equal size.");
     }
     vecT out_shape = shape;
     vecT out_strides1 = strides1;
     vecT out_strides2 = strides2;
-    py::ssize_t disp1(0);
-    py::ssize_t disp2(0);
+    T disp1(0);
+    T disp2(0);
 
     int nd = simplify_iteration_two_strides(dim, out_shape.data(),
                                             out_strides1.data(),

dpctl/tensor/libtensor/source/tensor_py.cpp

@@ -1449,7 +1449,7 @@ PYBIND11_MODULE(_tensor_impl, m)
     import_dpctl();
 
     m.def(
-        "_contract_iter", &contract_iter,
+        "_contract_iter", &contract_iter<py::ssize_t, py::value_error>,
         "Simplifies iteration of array of given shape & stride. Returns "
         "a triple: shape, stride and offset for the new iterator of possible "
         "smaller dimension, which traverses the same elements as the original "
@@ -1464,7 +1464,7 @@ PYBIND11_MODULE(_tensor_impl, m)
           py::arg("depends") = py::list());
 
     m.def(
-        "_contract_iter2", &contract_iter2,
+        "_contract_iter2", &contract_iter2<py::ssize_t, py::value_error>,
         "Simplifies iteration over elements of pair of arrays of given shape "
         "with strides stride1 and stride2. Returns "
         "a 5-tuple: shape, stride and offset for the new iterator of possible "