Force-set SetConstStride* when needed
@@ -1048,6 +1048,25 @@ class Solution:
assert(isPackedIndex(state, problemType["Index0"], 0x1))
assert(isPackedIndex(state, problemType["Index1"], 0x2))
+ if state["PackBatchDims"]==1:
+ for bi in problemType["IndicesBatch"]:
+ found = False
+ for setc in problemType["SetConstStrideB"]:
+ if setc[0]==bi and setc[1]==0:
+ found = True
+ if not found:
+ print ("Warning: batch index [%s,0] should be in SetConstStrideB"%bi)
+ problemType["SetConstStrideB"].append([bi,0])
+ if state["PackBatchDims"]==2:
+ for bi in problemType["IndicesBatch"]:
+ found = False
+ for setc in problemType["SetConstStrideA"]:
+ if setc[0]==bi and setc[1]==0:
+ found = True
+ if not found:
+ print ("Warning: batch index [%s,0] should be in SetConstStrideA"%bi)
+ problemType["SetConstStrideA"].append([bi,0])
+
for idx in problemType["IndexAssignmentsA"]:
if isPackedIndex(state, idx, 0x1):
assert (idx < problemType["NumIndicesC"])
