23
23
24
24
25
25
def apply_tensor_contraints (op_name : str , index : int ) -> list [object ]:
26
+ # Constraint to limit tensor size product to < 4000
27
+ max_size_constraint = cp .Size .Le (lambda deps , r , d : max (1 , int ((3999 ) ** (1 / r ))))
28
+
26
29
tensor_constraints = (
27
30
[
28
31
cp .Dtype .In (
@@ -39,7 +42,7 @@ def apply_tensor_contraints(op_name: str, index: int) -> list[object]:
39
42
cp .Value .Le (lambda deps , dtype , struct : 2 ** 4 ),
40
43
cp .Rank .Ge (lambda deps : 1 ),
41
44
cp .Size .Ge (lambda deps , r , d : 1 ),
42
- cp . Size . Le ( lambda deps , r , d : 2 ** 9 ) ,
45
+ max_size_constraint ,
43
46
cp .Rank .Le (lambda deps : 2 ** 3 ),
44
47
]
45
48
if op_name
@@ -62,7 +65,7 @@ def apply_tensor_contraints(op_name: str, index: int) -> list[object]:
62
65
cp .Value .Le (lambda deps , dtype , struct : 2 ** 4 ),
63
66
cp .Rank .Ge (lambda deps : 1 ),
64
67
cp .Size .Ge (lambda deps , r , d : 1 ),
65
- cp . Size . Le ( lambda deps , r , d : 2 ** 9 ) ,
68
+ max_size_constraint ,
66
69
cp .Rank .Le (lambda deps : 2 ** 3 ),
67
70
]
68
71
)
@@ -76,7 +79,7 @@ def apply_tensor_contraints(op_name: str, index: int) -> list[object]:
76
79
cp .Value .Le (lambda deps , dtype , struct : 2 ** 4 ),
77
80
cp .Rank .Ge (lambda deps : 1 ),
78
81
cp .Size .Ge (lambda deps , r , d : 1 ),
79
- cp . Size . Le ( lambda deps , r , d : 2 ** 9 ) ,
82
+ max_size_constraint ,
80
83
]
81
84
else :
82
85
tensor_constraints = [
@@ -94,7 +97,7 @@ def apply_tensor_contraints(op_name: str, index: int) -> list[object]:
94
97
cp .Value .Le (lambda deps , dtype , struct : 2 ** 4 ),
95
98
cp .Rank .Ge (lambda deps : 1 ),
96
99
cp .Size .Ge (lambda deps , r , d : 1 ),
97
- cp . Size . Le ( lambda deps , r , d : 2 ** 9 ) ,
100
+ max_size_constraint ,
98
101
]
99
102
case "embedding.default" :
100
103
tensor_constraints = [
@@ -104,7 +107,7 @@ def apply_tensor_contraints(op_name: str, index: int) -> list[object]:
104
107
cp .Value .Le (lambda deps , dtype , struct : 2 ** 4 ),
105
108
cp .Rank .Ge (lambda deps : 1 ),
106
109
cp .Size .Ge (lambda deps , r , d : 1 ),
107
- cp . Size . Le ( lambda deps , r , d : 2 ** 9 ) ,
110
+ max_size_constraint ,
108
111
]
109
112
case "sigmoid.default" :
110
113
tensor_constraints .extend (
0 commit comments