Repository: onnx/onnx
Pull request: Clarify `cond` to `If` must contain a single element
### Description

Clarify `cond` to `If` must contain a single element. Previously it is unclear what the condition can be.

Fix reference runtime accordingly.

### Motivation and Context

Fixes https://github.com/onnx/onnx/issues/5520

Problem statement:
Non scalar as condition to If op
# Ask a Question

### Question
<!-- -->
I write a onnx If op case， the condition input is np.array([1, 2, 0, 3]).astype(bool), the output is then_branch, if input is np.array([0, 2, 0, 3]).astype(bool), the output is else_branch, so from result, it looks like the first element decides True or False, but in onnx source code op_if.py:
```
if len(cond.shape) > 0:
    try:
         evaluated_condition = all(cond)
```
all means check if all items in a list are True, so it doesn't matter the experiment result, because [1, 2, 0, 3] has 0.

### Further information
- Relevant Area: <!--e.g., model usage, backend, best practices, converters, shape_inference, version_converter, training, test, operators, IR, ONNX Hub, data preprocessing, CI pipelines. -->

- Is this issue related to a specific model?
**Model name**: <!-- *e.g. mnist* --> 
**Model opset**: <!-- *e.g. 17* -->
```
    # Define input and output tensors
    cond = onnx.helper.make_tensor_value_info('cond', onnx.TensorProto.BOOL, [])
    res = onnx.helper.make_tensor_value_info('res', onnx.TensorProto.FLOAT, [5])

    # Define then and else output tensors
    then_out = onnx.helper.make_tensor_value_info('then_out', onnx.TensorProto.FLOAT, [5])
    else_out = onnx.helper.make_tensor_value_info('else_out', onnx.TensorProto.FLOAT, [5])

    # Define constant nodes for then and else branches
    x = np.array([1, 2, 3, 4, 5]).astype(np.float32)
    y = np.array([5, 4, 3, 2, 1]).astype(np.float32)

    then_const_node = onnx.helper.make_node(
        'Constant',
        inputs=[],
        outputs=['then_out'],
        value=onnx.numpy_helper.from_array(x),
    )

    else_const_node = onnx.helper.make_node(
        'Constant',
        inputs=[],
        outputs=['else_out'],
        value=onnx.numpy_helper.from_array(y),
    )

    # Define then and else subgraphs
    then_body = onnx.helper.make_graph(
        [then_const_node],
        'then_body',
        [],
        [then_out]
    )

    else_body = onnx.helper.make_graph(
        [else_const_node],
        'else_body',
        [],
        [else_out]
    )

    # Define If node
    if_node = onnx.helper.make_node(
        'If',
        inputs=['cond'],
        outputs=['res'],
        then_branch=then_body,
        else_branch=else_body,
    )

    graph_def = helper.make_graph([if_node],
                                  'if_model',
                                  inputs=[cond],
                                  outputs=[res])

    # 将模型转换为ONNX格式
    model = helper.make_model(graph_def, producer_name='onnx-example', opset_imports=[helper.make_opsetid("", 11)])
    onnx.save(model, 'if_model.onnx')

    # Print the model
    # print(model)

    # Run the model using onnxruntime
    sess = onnxruntime.InferenceSession(model.SerializeToString())
    cond_tensor = np.array([1, 2, 0, 3]).astype(bool)
    print("cond_tensor: ", cond_tensor)
    # print("all(cond_tensor): ", all(cond_tensor))
    input_data = {'cond': cond_tensor}
    output = sess.run([], input_data)
    print("if_output: ", output)
```

### Notes
<!-- Any additional information, code snippets. -->

Hints from issue discussion:
When do we expect the cond in if is not a scalar Boolean?
> When do we expect the cond in if is not a scalar Boolean?

in real case, the front operator may give a non scalar boolean tensor to If operator, so I want to know onnx if operator how to handle nonscalar boolean tensor
Possible to just take the boolean out? Sounds like we should only allow a single boolean as the condition
> 

However, onnx doesn't report error when I set nonscalar boolean, it works, so I think the onnx support non scalar boolean as cond. Could you tell me when the cond is nonscalar boolean, onnx how to decide the branch?
cc @gramalingam for spec clarification 
That sounds like a bug (in both cases) to me. The condition should have only one boolean value. Ideally, it should be a tensor of rank zero (but I can understand some implementations supporting other ranks, like rank 1, due to historical ambiguity, but allowing more than 1 boolean value seems inappropriate, with potential for confusion just like here).
> That sounds like a bug (in both cases) to me. The condition should have only one boolean value. Ideally, it should be a tensor of rank zero (but I can understand some implementations supporting other ranks, like rank 1, due to historical ambiguity, but allowing more than 1 boolean value seems inappropriate, with potential for confusion just like here).

so in the onnx current version, if the cond is a NonScalar Boolean, does onnx use the first element of cond to decide the branch?
Seems like that’s what ONNX Runtime does. It’s currently unspecified by ONNX

Commit message: Update onnx/reference/ops/op_if.py

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

Patch to review:
diff --git a/docs/Changelog.md b/docs/Changelog.md
index ffb2175d7e7..93f36c5f512 100644
--- a/docs/Changelog.md
+++ b/docs/Changelog.md
@@ -1579,7 +1579,7 @@ This version of the operator has been available since version 1 of the default O
 
 <dl>
 <dt><tt>cond</tt> : B</dt>
-<dd>Condition for the if</dd>
+<dd>Condition for the if. The tensor must contain a single element.</dd>
 </dl>
 
 #### Outputs (1 - &#8734;)
@@ -11525,7 +11525,7 @@ This version of the operator has been available since version 11 of the default
 
 <dl>
 <dt><tt>cond</tt> : B</dt>
-<dd>Condition for the if</dd>
+<dd>Condition for the if. The tensor must contain a single element.</dd>
 </dl>
 
 #### Outputs (1 - &#8734;)
@@ -16162,7 +16162,7 @@ This version of the operator has been available since version 13 of the default
 
 <dl>
 <dt><tt>cond</tt> : B</dt>
-<dd>Condition for the if</dd>
+<dd>Condition for the if. The tensor must contain a single element.</dd>
 </dl>
 
 #### Outputs (1 - &#8734;)
@@ -20057,7 +20057,7 @@ This version of the operator has been available since version 16 of the default
 
 <dl>
 <dt><tt>cond</tt> : B</dt>
-<dd>Condition for the if</dd>
+<dd>Condition for the if. The tensor must contain a single element.</dd>
 </dl>
 
 #### Outputs (1 - &#8734;)
@@ -23069,7 +23069,7 @@ This version of the operator has been available since version 19 of the default
 
 <dl>
 <dt><tt>cond</tt> : B</dt>
-<dd>Condition for the if</dd>
+<dd>Condition for the if. The tensor must contain a single element.</dd>
 </dl>
 
 #### Outputs (1 - &#8734;)
diff --git a/docs/Operators.md b/docs/Operators.md
index 3173b32c4c4..1f925418624 100644
--- a/docs/Operators.md
+++ b/docs/Operators.md
@@ -11905,7 +11905,7 @@ Other versions of this operator: <a href="Changelog.md#If-1">1</a>, <a href="Cha
 
 <dl>
 <dt><tt>cond</tt> : B</dt>
-<dd>Condition for the if</dd>
+<dd>Condition for the if. The tensor must contain a single element.</dd>
 </dl>
 
 #### Outputs (1 - &#8734;)
diff --git a/onnx/defs/controlflow/defs.cc b/onnx/defs/controlflow/defs.cc
index 8628f35bc77..f6bd6d64c7a 100644
--- a/onnx/defs/controlflow/defs.cc
+++ b/onnx/defs/controlflow/defs.cc
@@ -26,7 +26,7 @@ ONNX_OPERATOR_SET_SCHEMA(
     19,
     OpSchema()
         .SetDoc("If conditional")
-        .Input(0, "cond", "Condition for the if", "B")
+        .Input(0, "cond", "Condition for the if. The tensor must contain a single element.", "B")
         .Output(
             0,
             "outputs",
diff --git a/onnx/defs/controlflow/old.cc b/onnx/defs/controlflow/old.cc
index ff9f5bcee6d..2ade6bb8cb6 100644
--- a/onnx/defs/controlflow/old.cc
+++ b/onnx/defs/controlflow/old.cc
@@ -25,7 +25,7 @@ ONNX_OPERATOR_SET_SCHEMA(
     16,
     OpSchema()
         .SetDoc("If conditional")
-        .Input(0, "cond", "Condition for the if", "B")
+        .Input(0, "cond", "Condition for the if. The tensor must contain a single element.", "B")
         .Output(
             0,
             "outputs",
@@ -1747,7 +1747,7 @@ ONNX_OPERATOR_SET_SCHEMA(
     1,
     OpSchema()
         .SetDoc("If conditional")
-        .Input(0, "cond", "Condition for the if", "B")
+        .Input(0, "cond", "Condition for the if. The tensor must contain a single element.", "B")
         .Output(
             0,
             "outputs",
@@ -1841,7 +1841,7 @@ ONNX_OPERATOR_SET_SCHEMA(
     11,
     OpSchema()
         .SetDoc("If conditional")
-        .Input(0, "cond", "Condition for the if", "B")
+        .Input(0, "cond", "Condition for the if. The tensor must contain a single element.", "B")
         .Output(
             0,
             "outputs",
@@ -1933,7 +1933,7 @@ ONNX_OPERATOR_SET_SCHEMA(
     13,
     OpSchema()
         .SetDoc("If conditional")
-        .Input(0, "cond", "Condition for the if", "B")
+        .Input(0, "cond", "Condition for the if. The tensor must contain a single element.", "B")
         .Output(
             0,
             "outputs",
diff --git a/onnx/reference/ops/op_if.py b/onnx/reference/ops/op_if.py
index 528ea44f6e4..17668aeab20 100644
--- a/onnx/reference/ops/op_if.py
+++ b/onnx/reference/ops/op_if.py
@@ -2,6 +2,9 @@
 
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
+import numpy as np
 
 from onnx.reference.op_run import OpRun
 
@@ -23,28 +26,20 @@ def need_context(self) -> bool:
         """
         return True
 
-    def _run(self, cond, context=None, else_branch=None, then_branch=None, attributes=None):  # type: ignore
-        if len(cond.shape) > 0:
-            try:
-                evaluated_condition = all(cond)
-            except ValueError as e:
-                raise ValueError(
-                    f"Unable to evaluate the condition with {type(cond)}, "
-                    f"shape={cond.shape}, dtype={cond.dtype}."
-                ) from e
-            if evaluated_condition:
-                self._log("  -- then> {%r}", context)
-                outputs = self._run_then_branch(context, attributes=attributes)  # type: ignore
-                self._log("  -- then<")
-                final = tuple(outputs)
-                branch = "then"
-            else:
-                self._log("  -- else> {%r}", context)
-                outputs = self._run_else_branch(context, attributes=attributes)  # type: ignore
-                self._log("  -- else<")
-                final = tuple(outputs)
-                branch = "else"
-        elif cond:
+    def _run(
+        self,
+        cond: np.ndarray | np.bool_,
+        context=None,
+        else_branch=None,
+        then_branch=None,
+        attributes=None,
+    ):
+        if cond.size != 1:
+            raise RuntimeError(
+                f"Operator If ({self.onnx_node.name!r}) expects a single element condition, but size of 'cond' is {len(cond)}."
+            )
+        cond_ = cond.item(0)
+        if cond_:
             self._log("  -- then> {%r}", context)
             outputs = self._run_then_branch(context, attributes=attributes)  # type: ignore
             self._log("  -- then<")
@@ -58,7 +53,7 @@ def _run(self, cond, context=None, else_branch=None, then_branch=None, attribute
             branch = "else"
 
         if not final:
-            raise RuntimeError(  # pragma: no cover
+            raise RuntimeError(
                 f"Operator If ({self.onnx_node.name!r}) does not have any output."
             )
         for i, f in enumerate(final):
@@ -66,7 +61,7 @@ def _run(self, cond, context=None, else_branch=None, then_branch=None, attribute
                 br = self.then_branch if branch == "then" else self.else_branch  # type: ignore
                 names = br.output_names
                 inits = [i.name for i in br.obj.graph.initializer]
-                raise RuntimeError(  # pragma: no cover
+                raise RuntimeError(
                     f"Output {i!r} (branch={branch!r}, name={names[i]!r}) is None, "
                     f"available inputs={sorted(context)}, initializers={inits}."
                 )
diff --git a/onnx/test/reference_evaluator_test.py b/onnx/test/reference_evaluator_test.py
index 3b9f88e018f..06545db804b 100644
--- a/onnx/test/reference_evaluator_test.py
+++ b/onnx/test/reference_evaluator_test.py
@@ -785,11 +785,11 @@ def test_if(self):
         sess = ReferenceEvaluator(model_def)
         self.assertEqual(str(sess), "ReferenceEvaluator(X) -> Z")
 
-        x = np.array([1, 2], dtype=np.float32)
+        x = np.array([1], dtype=np.float32)
         got = sess.run(None, {"X": x})[0]
         assert_allclose(np.array([1], dtype=np.float32), got)
 
-        x = np.array([-1, -2], dtype=np.float32)
+        x = np.array([-1], dtype=np.float32)
         got = sess.run(None, {"X": x})[0]
         assert_allclose(np.array([0], dtype=np.float32), got)
 


Code review comments:
diff_hunk: @@ -23,28 +26,20 @@ def need_context(self) -> bool:
         """
         return True
 
-    def _run(self, cond, context=None, else_branch=None, then_branch=None, attributes=None):  # type: ignore
-        if len(cond.shape) > 0:
-            try:
-                evaluated_condition = all(cond)
-            except ValueError as e:
-                raise ValueError(
-                    f"Unable to evaluate the condition with {type(cond)}, "
-                    f"shape={cond.shape}, dtype={cond.dtype}."
-                ) from e
-            if evaluated_condition:
-                self._log("  -- then> {%r}", context)
-                outputs = self._run_then_branch(context, attributes=attributes)  # type: ignore
-                self._log("  -- then<")
-                final = tuple(outputs)
-                branch = "then"
-            else:
-                self._log("  -- else> {%r}", context)
-                outputs = self._run_else_branch(context, attributes=attributes)  # type: ignore
-                self._log("  -- else<")
-                final = tuple(outputs)
-                branch = "else"
-        elif cond:
+    def _run(
+        self,
+        cond: np.ndarray | np.bool_,
+        context=None,
+        else_branch=None,
+        then_branch=None,
+        attributes=None,
+    ):
+        if cond.size != 1:
+            raise RuntimeError(
path: onnx/reference/ops/op_if.py
text: @author:
```suggestion
            raise ValueError(
```