def METHOD_NAME(self, expr):
    if expr.aggregate.name not in self.callee_arg_dict:
        return super().METHOD_NAME(expr)
    callee_arg_dim_tags = self.callee_arg_dict[expr.aggregate.name].dim_tags
    flattened_index = sum(dim_tag.stride*idx for dim_tag, idx in
            zip(callee_arg_dim_tags, expr.index_tuple))
    new_indices = []
    from operator import mul
    from functools import reduce
    stride = reduce(mul, self.desired_shape[expr.aggregate.name], 1)
    for length in self.desired_shape[expr.aggregate.name]:
        stride /= length
        ind = flattened_index // int(stride)
        flattened_index -= (int(stride) * ind)
        new_indices.append(simplify_via_aff(ind))
    return expr.aggregate.index(tuple(new_indices))