1 files changed, 6 insertions, 1 deletions
diff --git a/nn/runtime/ModelBuilder.h b/nn/runtime/ModelBuilder.h
index c474ea7e8..94baab70b 100644
--- a/nn/runtime/ModelBuilder.h
+++ b/nn/runtime/ModelBuilder.h
@@ -133,7 +133,8 @@ class ModelBuilder {
    private:
     // TODO(b/132322449): move partitionTheWork, findBestDeviceForEachOperation,
     // getPerformance, supportedByControlFlowInterpreter,
-    // partitionTheWorkInternal, sortIntoRunOrder to CompilationBuilder?
+    // isControlFlowOperationWithOperandOfUnknownSize, partitionTheWorkInternal,
+    // sortIntoRunOrder to CompilationBuilder?
 
     // Populates bestDeviceForOperation
     //
@@ -152,6 +153,10 @@ class ModelBuilder {
                          uint32_t operationIndex) const;
     bool supportedByControlFlowInterpreter(uint32_t operationIndex) const;
 
+    // Returns true if the operation is IF or WHILE and has an inner or outer
+    // input or output of unknown size.
+    bool isControlFlowOperationWithOperandOfUnknownSize(uint32_t operationIndex) const;
+
     int partitionTheWorkInternal(uint32_t sourceModelIndex,
                                  const std::vector<std::shared_ptr<Device>>& devices,
                                  uint32_t preference, uint32_t priority,