diff --git a/python/lbmpy_walberla/walberla_lbm_generation.py b/python/lbmpy_walberla/walberla_lbm_generation.py
index 109e4ff21920ebc8376d61d9ccf7e774461ebebf..23f864051c828c4b72cd89823bbd3288a614115e 100644
--- a/python/lbmpy_walberla/walberla_lbm_generation.py
+++ b/python/lbmpy_walberla/walberla_lbm_generation.py
@@ -12,7 +12,7 @@ from lbmpy.updatekernels import create_lbm_kernel, create_stream_only_kernel
 from pystencils import AssignmentCollection, create_kernel, Target
 from pystencils.astnodes import SympyAssignment
 from pystencils.backends.cbackend import CBackend, CustomSympyPrinter, get_headers
-from pystencils.data_types import TypedSymbol, type_all_numbers, cast_func
+from pystencils.typing import TypedSymbol, cast_func
 from pystencils.field import Field
 from pystencils.stencil import offset_to_direction_string
 from pystencils.sympyextensions import get_symmetric_part
@@ -318,3 +318,8 @@ def equations_to_code(equations, variable_prefix="lm.", variables_without_prefix
                                                dtype=dtype))
         result.append(c_backend(assignment))
     return "\n".join(result)
+
+
+def type_all_numbers(expr, dtype):
+    substitutions = {a: cast_func(a, dtype) for a in expr.atoms(sp.Number)}
+    return expr.subs(substitutions)
\ No newline at end of file