code: def test_if_then_else():
      cases = [
          [(te.var("cond", dtype="bool"), "bool", "int32"), "int32"],
          [(True, "int32", "float32"), "float32"],
          [(False, "int32", "int64"), "int64"],
          [(te.var("cond", dtype="bool"), "uint32", "int32"), "uint32"],
          [
              (te.var("cond", dtype="int32"), "uint32", "int32"),
              "uint32",
          ],
      ]
      for (cond, lhs_dtype, rhs_dtype), out_dtype in cases:
          lhs = te.var("lhs", dtype=lhs_dtype)
          rhs = te.var("rhs", dtype=rhs_dtype)
          if cond is True or cond is False:
              out = tvm.tir.if_then_else(cond, lhs, rhs)
              out2 = tvm.tir.if_then_else(not cond, rhs, lhs)
              out3 = tvm.tir.if_then_else(not cond, lhs, rhs)
              assert tvm.ir.structural_equal(out, out2) == 1
              if cond:
                  assert (
                      tvm.ir.structural_equal(
                          out, lhs.astype(out_dtype)
                      )
                      == 1
                  )
                  assert (
                      tvm.ir.structural_equal(
                          out3, rhs.astype(out_dtype)
                      )
                      == 1
                  )
              else:
                  assert (
                      tvm.ir.structural_equal(
                          out, rhs.astype(out_dtype)
                      )
                      == 1
                  )
                  assert (
                      tvm.ir.structural_equal(
                          out3, lhs.astype(out_dtype)
                      )
                      == 1
                  )
          elif cond.dtype == "bool":
              out = tvm.tir.if_then_else(cond, lhs, rhs)
              assert out.dtype == out_dtype
              assert out.args[1].dtype == out_dtype
              assert out.args[2].dtype == out_dtype
          elif cond.dtype != "bool":
              check_throws(lambda: tvm.tir.if_then_else(cond, lhs, rhs))
          else:
              raise ValueError("Unknown combinations")
  
file_path: tvm/tests/python/tir-base/test_tir_ops.py
imports: import tvm
  import tvm.testing
  from tvm import te
  import pytest
label_desc: Use image node_id 7 for calling a global function with example usage: test_if_then_else() without return types