Cache SizeInfo on ReduceNode

dwave/optimization/include/dwave-optimization/nodes/reduce.hpp

@@ -86,6 +86,9 @@ class ReduceNode : public ArrayOutputMixin<ArrayNode> {
     using ArrayOutputMixin::size;
     ssize_t size(const State& state) const override;
 
+    /// @copydoc Array::sizeinfo()
+    SizeInfo sizeinfo() const override { return sizeinfo_; };
+
     /// @copydoc Array::size_diff()
     ssize_t size_diff(const State& state) const override;
 
@@ -111,6 +114,8 @@ class ReduceNode : public ArrayOutputMixin<ArrayNode> {
     // as whether we're integral or not.
     const ValuesInfo values_info_;
 
+    const SizeInfo sizeinfo_;
+
     // During propagation, we need to take a predecessor's `update` and apply
     // it to the correct `index` in the ReduceNode buffer. This method
     // determines the correct buffer `index` given an `update.index`.

dwave/optimization/src/nodes/reduce.cpp

@@ -24,6 +24,8 @@
 #include <utility>
 
 #include "../functional_.hpp"
+#include "dwave-optimization/array.hpp"
+#include "dwave-optimization/state.hpp"
 
 namespace dwave::optimization {
 
@@ -281,7 +283,7 @@ class ReduceNodeData : public NodeStateData {
             if (flags_[index] == ReductionFlag::unchanged) {
                 reductions_diff_.emplace_back(index, reductions_[index]);
             }
-            flags_[index] = ReductionFlag::invalid;            
+            flags_[index] = ReductionFlag::invalid;
             return;
         }
         reduction = std::move(inverse.value());
@@ -581,6 +583,19 @@ ValuesInfo values_info(const Array* array_ptr, std::span<const ssize_t> axes,
     return bounds;
 }
 
+SizeInfo reducenode_calculate_sizeinfo(const Array* node_ptr, const Array* array_ptr,
+                                       std::span<const ssize_t> axes) {
+    // Node is statically sized. Note: If node_shape = {}, product(node_shape) = 1.
+    if (!node_ptr->dynamic()) return SizeInfo(product(node_ptr->shape()));
+    assert(node_ptr->shape().size() && node_ptr->shape().front() == -1);
+
+    // Node is dynamically sized but predecessor is always empty.
+    if (array_ptr->size() == 0) return SizeInfo(node_ptr);
+
+    // Node size is derived from its predecessor's size.
+    return array_ptr->sizeinfo() / product(keep_axes(array_ptr->shape(), axes));
+}
+
 template <class BinaryOp>
 ReduceNode<BinaryOp>::ReduceNode(ArrayNode* array_ptr) : ReduceNode(array_ptr, {}) {}
 
@@ -591,7 +606,8 @@ ReduceNode<BinaryOp>::ReduceNode(ArrayNode* array_ptr, std::span<const ssize_t>
           initial(initial),
           array_ptr_(array_ptr),
           axes_(normalize_axes(array_ptr, axes)),
-          values_info_(values_info<BinaryOp>(array_ptr_, axes_, initial)) {
+          values_info_(values_info<BinaryOp>(array_ptr_, axes_, initial)),
+          sizeinfo_(reducenode_calculate_sizeinfo(this, array_ptr_, axes_)) {
     add_predecessor(array_ptr);
 }
 

releasenotes/notes/reducenode_cache_sizeinfo-cba3557baf3f7901.yaml

@@ -0,0 +1,5 @@
+---
+features:
+  - |
+    The symbol: `All`, `Any`, `Max`, `Min`, `Prod`, and `Sum` can infer
+    their size from their predecessor's size.