| 1983 | // (`f_trivial`) storage buffer; returns `non_trivial` otherwise. |
| 1984 | template <size_t N> |
| 1985 | broadcast_trivial |
| 1986 | broadcast(const std::array<buffer_info, N> &buffers, ssize_t &ndim, std::vector<ssize_t> &shape) { |
| 1987 | ndim = std::accumulate( |
| 1988 | buffers.begin(), buffers.end(), ssize_t(0), [](ssize_t res, const buffer_info &buf) { |
| 1989 | return std::max(res, buf.ndim); |
| 1990 | }); |
| 1991 | |
| 1992 | shape.clear(); |
| 1993 | shape.resize((size_t) ndim, 1); |
| 1994 | |
| 1995 | // Figure out the output size, and make sure all input arrays conform (i.e. are either size 1 |
| 1996 | // or the full size). |
| 1997 | for (size_t i = 0; i < N; ++i) { |
| 1998 | auto res_iter = shape.rbegin(); |
| 1999 | auto end = buffers[i].shape.rend(); |
| 2000 | for (auto shape_iter = buffers[i].shape.rbegin(); shape_iter != end; |
| 2001 | ++shape_iter, ++res_iter) { |
| 2002 | const auto &dim_size_in = *shape_iter; |
| 2003 | auto &dim_size_out = *res_iter; |
| 2004 | |
| 2005 | // Each input dimension can either be 1 or `n`, but `n` values must match across |
| 2006 | // buffers |
| 2007 | if (dim_size_out == 1) { |
| 2008 | dim_size_out = dim_size_in; |
| 2009 | } else if (dim_size_in != 1 && dim_size_in != dim_size_out) { |
| 2010 | pybind11_fail("pybind11::vectorize: incompatible size/dimension of inputs!"); |
| 2011 | } |
| 2012 | } |
| 2013 | } |
| 2014 | |
| 2015 | bool trivial_broadcast_c = true; |
| 2016 | bool trivial_broadcast_f = true; |
| 2017 | for (size_t i = 0; i < N && (trivial_broadcast_c || trivial_broadcast_f); ++i) { |
| 2018 | if (buffers[i].size == 1) { |
| 2019 | continue; |
| 2020 | } |
| 2021 | |
| 2022 | // Require the same number of dimensions: |
| 2023 | if (buffers[i].ndim != ndim) { |
| 2024 | return broadcast_trivial::non_trivial; |
| 2025 | } |
| 2026 | |
| 2027 | // Require all dimensions be full-size: |
| 2028 | if (!std::equal(buffers[i].shape.cbegin(), buffers[i].shape.cend(), shape.cbegin())) { |
| 2029 | return broadcast_trivial::non_trivial; |
| 2030 | } |
| 2031 | |
| 2032 | // Check for C contiguity (but only if previous inputs were also C contiguous) |
| 2033 | if (trivial_broadcast_c) { |
| 2034 | ssize_t expect_stride = buffers[i].itemsize; |
| 2035 | auto end = buffers[i].shape.crend(); |
| 2036 | for (auto shape_iter = buffers[i].shape.crbegin(), |
| 2037 | stride_iter = buffers[i].strides.crbegin(); |
| 2038 | trivial_broadcast_c && shape_iter != end; |
| 2039 | ++shape_iter, ++stride_iter) { |
| 2040 | if (expect_stride == *stride_iter) { |
| 2041 | expect_stride *= *shape_iter; |
| 2042 | } else { |