diff --git a/pystencils/cpu/vectorization.py b/pystencils/cpu/vectorization.py index 2f49d44a2b848d85bed7fe6ac896bc730351c8c2..4c632b1457c1e2778953a55d4378ee87f05ffb82 100644 --- a/pystencils/cpu/vectorization.py +++ b/pystencils/cpu/vectorization.py @@ -210,6 +210,8 @@ def insert_vector_casts(ast_node): # special treatment for the unary minus: make sure that the -1 has the same type as the argument dtype = int for arg in expr.args[1:]: + if type(arg) is sp.Pow: + arg = arg.args[0] if type(arg) is vector_memory_access and arg.dtype.base_type.is_float(): dtype = arg.dtype.base_type.numpy_dtype.type elif type(arg) is TypedSymbol and type(arg.dtype) is VectorType and arg.dtype.base_type.is_float(): diff --git a/pystencils/simp/assignment_collection.py b/pystencils/simp/assignment_collection.py index 95064408959830650df6402cf8f22983122638c7..33102dee5085549f7ce827efe99d15f67da2008d 100644 --- a/pystencils/simp/assignment_collection.py +++ b/pystencils/simp/assignment_collection.py @@ -437,9 +437,10 @@ class AssignmentCollection: class SymbolGen: """Default symbol generator producing number symbols ζ_0, ζ_1, ...""" - def __init__(self, symbol="xi"): + def __init__(self, symbol="xi", dtype=None): self._ctr = 0 self._symbol = symbol + self._dtype = dtype def __iter__(self): return self @@ -447,4 +448,6 @@ class SymbolGen: def __next__(self): name = f"{self._symbol}_{self._ctr}" self._ctr += 1 + if self._dtype is not None: + return pystencils.TypedSymbol(name, self._dtype) return sp.Symbol(name)