Merge pull request #433 from arcondello/fix/BinaryOp-broadcasting

Fix shape when broadcasting with BinaryOpNode

Author: arcondello

dwave/optimization/src/array.cpp

@@ -264,6 +264,12 @@ std::vector<ssize_t> broadcast_shape(const std::span<const ssize_t> lhs,
     }
     assert(sit == shape.rend());
 
+    // Check that we haven't put a dynamic axis anywhere except axis 0
+    if (std::ranges::any_of(shape | std::views::drop(1), [](const auto& val) { return val < 0; })) {
+        throw std::invalid_argument("operands could not be broadcast together with shapes " +
+                                    shape_to_string(lhs) + " " + shape_to_string(rhs));
+    }
+
     return shape;
 }
 std::vector<ssize_t> broadcast_shape(std::initializer_list<ssize_t> lhs,

dwave/optimization/src/nodes/binaryop.cpp

@@ -367,25 +367,40 @@ template <class BinaryOp>
 std::span<const ssize_t> BinaryOpNode<BinaryOp>::shape(const State& state) const {
     if (!this->dynamic()) return this->shape();
 
-    const ssize_t lhs_size = operands_[0]->size(state);
+    const auto [lhs, rhs] = operands_;
 
-    if (lhs_size == operands_[1]->size(state)) return operands_[0]->shape(state);
+    // If we're broadcasting we know which size we're looking at
+    if (lhs->size() == 1) return rhs->shape(state);
+    if (rhs->size() == 1) return lhs->shape(state);
 
-    return (lhs_size == 1) ? operands_[1]->shape(state) : operands_[0]->shape(state);
+    // Dev note: it's very tempting to put an assert here that checks
+    // that lhs.shape == rhs.shape, but that can make calling shape()
+    // very expensive in some cases.
+    // Having this commented out makes the second if-branch above redundant, but
+    // I am leaving this in for clarity and future debugging.
+    // assert(std::ranges::equal(lhs->shape(state), rhs->shape(state)));
+
+    return lhs->shape(state);
 }
 
 template <class BinaryOp>
 ssize_t BinaryOpNode<BinaryOp>::size(const State& state) const {
-    if (ssize_t size = this->size(); size >= 0) {
-        return size;
-    }
+    if (const ssize_t size = this->size(); size >= 0) return size;
+
+    const auto [lhs, rhs] = operands_;
 
-    const ssize_t lhs_size = operands_[0]->size(state);
-    const ssize_t rhs_size = operands_[1]->size(state);
+    // If we're broadcasting we know which size we're looking at
+    if (lhs->size() == 1) return rhs->size(state);
+    if (rhs->size() == 1) return lhs->size(state);
 
-    if (lhs_size == rhs_size) return lhs_size;
+    // Dev note: it's very tempting to put an assert here that checks
+    // that lhs.size == rhs.size, but that can make calling size()
+    // very expensive in some cases.
+    // Having this commented out makes the second if-branch above redundant, but
+    // I am leaving this in for clarity and future debugging.
+    // assert(lhs->size(state) == rhs->size(state)));
 
-    return (lhs_size == 1) ? rhs_size : lhs_size;
+    return lhs->size(state);
 }
 
 template <class BinaryOp>

releasenotes/notes/fix-BinaryOp-shape-2c3447f7a1f41a0f.yaml

@@ -0,0 +1,11 @@
+---
+fixes:
+  - |
+    Change C++ function ``broadcast_shape()`` correctly throw an error when
+    an attempt to broadcast a dynamically sized shape would result in an invalid
+    result.
+    See `#429 <https://github.com/dwavesystems/dwave-optimization/issues/429>`_.
+  - |
+    Fix ``BinaryNode::shape(const State&)`` and ``BinaryNode::size(const State&)``
+    so that they do not propagate the wrong shape/size when both predecessors
+    have the same shape.

tests/cpp/nodes/test_binaryop.cpp

@@ -466,6 +466,22 @@ TEST_CASE("BinaryOpNode - MultiplyNode") {
             CHECK(y_ptr->integral());
         }
     }
+
+    GIVEN("x = SetNode(5), y = Constant(0), z = x * y") {
+        // This test is for a specific bug we had where it would accidentlly
+        // pick up the shape from the lhs if the size of the lhs and rhs were
+        // both 1.
+
+        auto x_ptr = graph.emplace_node<SetNode>(5);
+        auto y_ptr = graph.emplace_node<ConstantNode>(0);
+        auto z_ptr = graph.emplace_node<MultiplyNode>(y_ptr, x_ptr);
+
+        auto state = graph.empty_state();
+        x_ptr->initialize_state(state, {0});
+        graph.initialize_state(state);
+
+        CHECK_THAT(z_ptr->shape(state), RangeEquals({1}));
+    }
 }
 
 TEST_CASE("BinaryOpNode - DivideNode") {

tests/cpp/test_array.cpp

@@ -481,32 +481,39 @@ TEST_CASE("Update") {
     }
 }
 
-TEST_CASE("Test resulting_shape()") {
+TEST_CASE("Test broadcast_shape()") {
     SECTION("(256,256,3) x (3,) -> (256,256,3)") {
-        CHECK(std::ranges::equal(broadcast_shape({256, 256, 3}, {3}), std::vector{256, 256, 3}));
+        CHECK_THAT(broadcast_shape({256, 256, 3}, {3}), RangeEquals({256, 256, 3}));
     }
     SECTION("(8,1,6,1) x (7,1,5) -> (8,7,6,5)") {
-        CHECK(std::ranges::equal(broadcast_shape({8, 1, 6, 1}, {7, 1, 5}),
-                                 std::vector{8, 7, 6, 5}));
+        CHECK_THAT(broadcast_shape({8, 1, 6, 1}, {7, 1, 5}), RangeEquals({8, 7, 6, 5}));
     }
     SECTION("(7,1,5) x (8,1,6,1) -> (8,7,6,5)") {
-        CHECK(std::ranges::equal(broadcast_shape({7, 1, 5}, {8, 1, 6, 1}),
-                                 std::vector{8, 7, 6, 5}));
+        CHECK_THAT(broadcast_shape({7, 1, 5}, {8, 1, 6, 1}), RangeEquals({8, 7, 6, 5}));
     }
     SECTION("(5,4) x (1,) -> (5, 4)") {
-        CHECK(std::ranges::equal(broadcast_shape({5, 4}, {1}), std::vector{5, 4}));
+        CHECK_THAT(broadcast_shape({5, 4}, {1}), RangeEquals({5, 4}));
     }
     SECTION("(5,4) x (4,) -> (5, 4)") {
-        CHECK(std::ranges::equal(broadcast_shape({5, 4}, {4}), std::vector{5, 4}));
+        CHECK_THAT(broadcast_shape({5, 4}, {4}), RangeEquals({5, 4}));
     }
     SECTION("(15,3,5) x (15,1,5) -> (15,3,5)") {
-        CHECK(std::ranges::equal(broadcast_shape({15, 3, 5}, {15, 1, 5}), std::vector{15, 3, 5}));
+        CHECK_THAT(broadcast_shape({15, 3, 5}, {15, 1, 5}), RangeEquals({15, 3, 5}));
     }
     SECTION("(15,3,5) x (3,5) -> (15,3,5)") {
-        CHECK(std::ranges::equal(broadcast_shape({15, 3, 5}, {3, 5}), std::vector{15, 3, 5}));
+        CHECK_THAT(broadcast_shape({15, 3, 5}, {3, 5}), RangeEquals({15, 3, 5}));
     }
     SECTION("(15,3,5) x (3,1) -> (15,3,5)") {
-        CHECK(std::ranges::equal(broadcast_shape({15, 3, 5}, {3, 1}), std::vector{15, 3, 5}));
+        CHECK_THAT(broadcast_shape({15, 3, 5}, {3, 1}), RangeEquals({15, 3, 5}));
+    }
+    SECTION("(-1,) x (1,) -> (-1,)") { CHECK_THAT(broadcast_shape({-1}, {1}), RangeEquals({-1})); }
+    SECTION("(1,) x (-1,) -> (-1,)") { CHECK_THAT(broadcast_shape({1}, {-1}), RangeEquals({-1})); }
+    SECTION("(1,2) x (-1,1) -> (-1,2)") {
+        CHECK_THAT(broadcast_shape({1, 2}, {-1, 1}), RangeEquals({-1, 2}));
+    }
+    SECTION("(1,1) x (-1,) -> invalid") {
+        CHECK_THROWS_WITH(broadcast_shape({1, 1}, {-1}),
+                          "operands could not be broadcast together with shapes (1, 1) (-1,)");
     }
     SECTION("(3,) x (4,)") {
         CHECK_THROWS_WITH(broadcast_shape({3}, {4}),

tests/test_symbols.py

@@ -164,6 +164,11 @@ def test_broadcasting(self):
         with self.assertRaises(ValueError):
             a + b
 
+        c = model.constant([[[1]]])
+        d = model.set(5)
+        with self.assertRaises(ValueError):
+            c + d
+
     def test_scalar_addition(self):
         model = Model()
         a = model.constant(5)