diff --git a/sumpy/kernel.py b/sumpy/kernel.py index 81556bf92..d6f7ab6d7 100644 --- a/sumpy/kernel.py +++ b/sumpy/kernel.py @@ -1092,10 +1092,10 @@ def __init__(self, vec_name, additional_indices): def map_subscript(self, expr): from pymbolic.primitives import CommonSubexpression, cse_scope - if (expr.aggregate.name == self.vec_name - and isinstance(expr.index, int)): + if (expr.aggregate.name == self.vec_name and isinstance(expr.index, int)): + index = (expr.index, *self.additional_indices) return CommonSubexpression( - expr.aggregate[(expr.index, *self.additional_indices)], + expr.aggregate[index], prefix=None, scope=cse_scope.EVALUATION) else: return IdentityMapper.map_subscript(self, expr)