@@ -237,14 +237,6 @@ def __new__(cls, *args, **kwargs):
237237 cond = d .relation (cond , GuardFactor (d ))
238238 conditionals [d ] = cond
239239
240- # Replace the ConditionalDimensions in `expr`
241- for d , cond in conditionals .items ():
242- # Replace dimension with index
243- index = d .index
244- index = index - relational_min (cond , d .parent )
245- shift = relational_shift (cond , d .parent )
246- expr = uxreplace (expr , {d : IntDiv (index , d .symbolic_factor ) + shift })
247-
248240 # Merge conditionals when possible. E.g if we have an implicit_dim
249241 # and there is a dimension with the same parent, we ca merged
250242 # its condition
@@ -254,14 +246,23 @@ def __new__(cls, *args, **kwargs):
254246 for cd in dict (conditionals ):
255247 if cd .parent == d .parent and cd != d :
256248 cond = conditionals .pop (d )
257- mode = cd .relation and d .relation
258- if issubclass (mode , sympy .Or ):
249+ if d .relation == 'strict' :
259250 conditionals [d ] = cond
260251 conditionals .pop (cd )
261252 else :
253+ mode = cd .relation and d .relation
262254 conditionals [cd ] = mode (cond , conditionals [cd ])
263255 break
264256
257+ # Replace the ConditionalDimensions in `expr`
258+ for d , cond in conditionals .items ():
259+ # Replace dimension with index
260+ index = d .index
261+ if d .condition is not None and d in expr .free_symbols :
262+ index = index - relational_min (cond , d .parent )
263+ shift = relational_shift (cond , d .parent )
264+ expr = uxreplace (expr , {d : IntDiv (index , d .symbolic_factor ) + shift })
265+
265266 # Lower all Differentiable operations into SymPy operations
266267 rhs = diff2sympy (expr .rhs )
267268
0 commit comments