[Mlir-commits] [mlir] [mlir][python] fix up affine for (PR #74495)

[Oleksandr Alex Zinenko]

================
@@ -45,23 +46,50 @@ def __init__(
         - `lower_bound_operands` is the list of arguments to substitute the dimensions,
           then symbols in the `lower_bound` affine map, in an increasing order
         - `upper_bound_operands` is the list of arguments to substitute the dimensions,
-          then symbols in the `upper_bound` affine map, in an increasing order
+          then symbols in the `upper_bound` affine map, in an increasing order.
         """
 
+        if lower_bound_operands is None:
+            lower_bound_operands = []
+        if upper_bound_operands is None:
+            upper_bound_operands = []
+
+        if step is None:
+            step = 1
+        if upper_bound is None:
+            upper_bound, lower_bound = lower_bound, 0
+
+        if isinstance(lower_bound, int):
+            lower_bound = AffineMap.get_constant(lower_bound)
+        elif isinstance(lower_bound, (Operation, OpView, Value)):
+            if len(lower_bound_operands):
+                raise ValueError(
+                    f"Either a concrete lower bound or an AffineMap in combination "
+                    f"with lower bound operands, but not both, is supported."
+                )
+            lower_bound_operands.append(lower_bound)
+            lower_bound = AffineMap.get_identity(1)
+
+        if not isinstance(lower_bound, AffineMap):
+            raise ValueError(f"{lower_bound=} must be int | ResultValueT | AffineMap")
+
+        if isinstance(upper_bound, int):
+            upper_bound = AffineMap.get_constant(upper_bound)
+        elif isinstance(upper_bound, (Operation, OpView, Value)):
+            if len(upper_bound_operands):
+                raise ValueError(
+                    f"Either a concrete upper bound or an AffineMap in combination "
+                    f"with upper bound operands, but not both, is supported."
+                )
+            upper_bound_operands.append(upper_bound)
+            upper_bound = AffineMap.get_identity(1)
+
+        if not isinstance(upper_bound, AffineMap):
+            raise ValueError(f"{upper_bound=} must be int | ResultValueT | AffineMap")
+
         if iter_args is None:
             iter_args = []
         iter_args = _get_op_results_or_values(iter_args)
-        if len(lower_bound_operands) != lower_bound.n_inputs:
-            raise ValueError(
-                f"Wrong number of lower bound operands passed to AffineForOp. "
-                + "Expected {lower_bound.n_symbols}, got {len(lower_bound_operands)}."
-            )
-
-        if len(upper_bound_operands) != upper_bound.n_inputs:
-            raise ValueError(
-                f"Wrong number of upper bound operands passed to AffineForOp. "
-                + "Expected {upper_bound.n_symbols}, got {len(upper_bound_operands)}."
-            )
----------------
ftynse wrote:

These checks looked reasonable (except for using `n_symbols` instead `n_inputs` in the error message), why remove them?

https://github.com/llvm/llvm-project/pull/74495