diff --git a/pystencils_tests/test_address_of.py b/pystencils_tests/test_address_of.py index 6e23d5ffd6513cba16acb2d711a9de0cf16c222e..5b297fc63a8e088d858d2abd939b289961c2d7f6 100644 --- a/pystencils_tests/test_address_of.py +++ b/pystencils_tests/test_address_of.py @@ -3,20 +3,18 @@ Test of pystencils.data_types.address_of """ -import sympy - import pystencils -from pystencils.data_types import PointerType, address_of, cast_func +from pystencils.data_types import PointerType, address_of, cast_func, create_type from pystencils.simp.simplifications import sympy_cse def test_address_of(): x, y = pystencils.fields('x,y: int64[2d]') - s = pystencils.TypedSymbol('s', PointerType('int64')) + s = pystencils.TypedSymbol('s', PointerType(create_type('int64'))) assignments = pystencils.AssignmentCollection({ s: address_of(x[0, 0]), - y[0, 0]: cast_func(s, 'int64') + y[0, 0]: cast_func(s, create_type('int64')) }, {}) ast = pystencils.create_kernel(assignments) @@ -24,7 +22,7 @@ def test_address_of(): print(code) assignments = pystencils.AssignmentCollection({ - y[0, 0]: cast_func(address_of(x[0, 0]), 'int64') + y[0, 0]: cast_func(address_of(x[0, 0]), create_type('int64')) }, {}) ast = pystencils.create_kernel(assignments) @@ -34,11 +32,11 @@ def test_address_of(): def test_address_of_with_cse(): x, y = pystencils.fields('x,y: int64[2d]') - s = pystencils.TypedSymbol('s', PointerType('int64')) + s = pystencils.TypedSymbol('s', PointerType(create_type('int64'))) assignments = pystencils.AssignmentCollection({ - y[0, 0]: cast_func(address_of(x[0, 0]), 'int64'), - x[0, 0]: cast_func(address_of(x[0, 0]), 'int64') + 1 + y[0, 0]: cast_func(address_of(x[0, 0]), create_type('int64')) + s, + x[0, 0]: cast_func(address_of(x[0, 0]), create_type('int64')) + 1 }, {}) ast = pystencils.create_kernel(assignments)