diff --git a/backends/cadence/utils/facto_util.py b/backends/cadence/utils/facto_util.py index 4f0ab13e162..4dbc76e0e47 100644 --- a/backends/cadence/utils/facto_util.py +++ b/backends/cadence/utils/facto_util.py @@ -344,14 +344,25 @@ def random_size_constraint(deps: object, r: int, d: int) -> int: ] ) case "div.Tensor": - tensor_constraints.extend( - [ - cp.Value.Ne(lambda deps, dtype, struct: 0), - cp.Value.Le(lambda deps, dtype, struct: 2**3), - cp.Size.Le(lambda deps, r, d: 2**3), - cp.Rank.Le(lambda deps: 2**2), - ] - ) + if index == 1: # Only apply zero-prevention to divisor + tensor_constraints.extend( + [ + cp.Value.Ne( + lambda deps, dtype, struct: 0 + ), # Prevent division by zero + cp.Value.Le(lambda deps, dtype, struct: 2**3), + cp.Size.Le(lambda deps, r, d: 2**3), + cp.Rank.Le(lambda deps: 2**2), + ] + ) + else: + tensor_constraints.extend( + [ + cp.Value.Le(lambda deps, dtype, struct: 2**3), + cp.Size.Le(lambda deps, r, d: 2**3), + cp.Rank.Le(lambda deps: 2**2), + ] + ) case "pow.Tensor_Scalar": tensor_constraints.extend( [