Add TODO

src/pystencils_autodiff/framework_integration/astnodes.py

@@ -208,6 +208,8 @@ class JinjaCppFile(Node):
                               else a)
                              for a in v]
                        for k, v in self.ast_dict.items()}
+
+        # TODO: possibly costly tree traversal
         render_dict.update({"headers": pystencils.backends.cbackend.get_headers(self)})
         render_dict.update({"globals":
                             {self.printer(g) for g in pystencils.backends.cbackend.get_global_declarations(self)}})

tests/test_module_printing.py

@@ -122,3 +122,4 @@ def test_custom_printer():
     backward_ast.function_name = 'backward'
     module = TorchModule("hallo", [forward_ast, backward_ast])
     print(DoesNotLikeTorchPrinter()(module))
+