Implements contract_iter4

- Matches other contract_iter functions

Author: ndgrigorian

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

@@ -872,6 +872,42 @@ int simplify_iteration_four_strides(const int nd,
     return nd_;
 }
 
+template <typename T, class Error, typename vecT = std::vector<T>>
+std::tuple<vecT, vecT, T, vecT, T, vecT, T, vecT, T>
+contract_iter4(vecT shape,
+               vecT strides1,
+               vecT strides2,
+               vecT strides3,
+               vecT strides4)
+{
+    const size_t dim = shape.size();
+    if (dim != strides1.size() || dim != strides2.size() ||
+        dim != strides3.size() || dim != strides4.size())
+    {
+        throw Error("Shape and strides must be of equal size.");
+    }
+    vecT out_shape = shape;
+    vecT out_strides1 = strides1;
+    vecT out_strides2 = strides2;
+    vecT out_strides3 = strides3;
+    vecT out_strides4 = strides4;
+    T disp1(0);
+    T disp2(0);
+    T disp3(0);
+    T disp4(0);
+
+    int nd = simplify_iteration_four_strides(
+        dim, out_shape.data(), out_strides1.data(), out_strides2.data(),
+        out_strides3.data(), out_strides4.data(), disp1, disp2, disp3, disp4);
+    out_shape.resize(nd);
+    out_strides1.resize(nd);
+    out_strides2.resize(nd);
+    out_strides3.resize(nd);
+    out_strides4.resize(nd);
+    return std::make_tuple(out_shape, out_strides1, disp1, out_strides2, disp2,
+                           out_strides3, disp3, out_strides4, disp4);
+}
+
 } // namespace strides
 } // namespace tensor
 } // namespace dpctl

dpctl/tensor/libtensor/source/simplify_iteration_space.cpp

@@ -674,7 +674,7 @@ void simplify_iteration_space_4(
         assert(simplified_shape.size() == static_cast<size_t>(nd));
         assert(simplified_src1_strides.size() == static_cast<size_t>(nd));
         assert(simplified_src2_strides.size() == static_cast<size_t>(nd));
-        assert(simplified_src3_strides.size() == static_cast < size_t(nd));
+        assert(simplified_src3_strides.size() == static_cast<size_t>(nd));
         assert(simplified_dst_strides.size() == static_cast<size_t>(nd));
     }
     shape = const_cast<const py::ssize_t *>(simplified_shape.data());

dpctl/tensor/libtensor/source/tensor_py.cpp

@@ -171,6 +171,17 @@ PYBIND11_MODULE(_tensor_impl, m)
         "as the original "
         "iterator, possibly in a different order.");
 
+    using dpctl::tensor::strides::contract_iter4;
+    m.def(
+        "_contract_iter4", &contract_iter4<py::ssize_t, py::value_error>,
+        "Simplifies iteration over elements of 4-tuple of arrays of given "
+        "shape "
+        "with strides stride1, stride2, stride3, and stride4. Returns "
+        "a 9-tuple: shape, stride and offset for the new iterator of possible "
+        "smaller dimension for each array, which traverses the same elements "
+        "as the original "
+        "iterator, possibly in a different order.");
+
     m.def("_copy_usm_ndarray_for_reshape", &copy_usm_ndarray_for_reshape,
           "Copies from usm_ndarray `src` into usm_ndarray `dst` with the same "
           "number of elements using underlying 'C'-contiguous order for flat "