diff --git a/include/xtensor/core/xfunction.hpp b/include/xtensor/core/xfunction.hpp index 99144b377..38982e479 100644 --- a/include/xtensor/core/xfunction.hpp +++ b/include/xtensor/core/xfunction.hpp @@ -710,7 +710,9 @@ namespace xt template inline bool xfunction::broadcast_shape(S& shape, bool reuse_cache) const { - if (m_cache.is_initialized && reuse_cache) + // The cached shape can only be copied into a shape of the same rank; any other + // rank goes through the full broadcast, which also reports incompatible ranks. + if (m_cache.is_initialized && reuse_cache && shape.size() == m_cache.shape.size()) { std::copy(m_cache.shape.cbegin(), m_cache.shape.cend(), shape.begin()); return m_cache.is_trivial; diff --git a/test/test_xfunction.cpp b/test/test_xfunction.cpp index ddf08758c..a9f2dccef 100644 --- a/test/test_xfunction.cpp +++ b/test/test_xfunction.cpp @@ -110,6 +110,31 @@ namespace xt XT_EXPECT_ANY_THROW(xt::xarray res = arr1 * arr2); } + TEST(xfunction, broadcast_shape_cached_other_rank) + { + xt::xarray a = {{1.0, 2.0, 3.0}, {4.0, 5.0, 6.0}}; + auto f = a + 1.0; + // Fills the shape cache of f + EXPECT_EQ(f.shape().size(), 2u); + + SUBCASE("higher rank") + { + using shape_type = xt::dynamic_shape; + shape_type cached = uninitialized_shape(3); + shape_type computed = uninitialized_shape(3); + bool cached_trivial = f.broadcast_shape(cached, true); + bool computed_trivial = f.broadcast_shape(computed, false); + EXPECT_EQ(cached, computed); + EXPECT_EQ(cached_trivial, computed_trivial); + } + + SUBCASE("lower rank") + { + xt::xtensor res; + XT_EXPECT_ANY_THROW(res = f); + } + } + TEST(xfunction, shape) { xfunction_features f;