Fix importing of _compile_env

src/pystencils_autodiff/tensorflow_jit.py

@@ -29,7 +29,7 @@ if get_compiler_config()['os'] != 'windows':
     _include_flags = ['-I' + sysconfig.get_paths()['include'], '-I' + get_pystencils_include_path()]
     _do_not_link_flag = "-c"
     _position_independent_flag = "-fPIC"
-    _compile_environment = os.environ.copy()
+    _compile_env = os.environ.copy()
 else:
     _do_not_link_flag = "/c"
     _output_flag = '/OUT:'
@@ -38,8 +38,8 @@ else:
     _position_independent_flag = "/DTHIS_FLAG_DOES_NOTHING"
     get_compiler_config()['command'] = 'cl.exe'
     config_env = get_compiler_config()['env'] if 'env' in get_compiler_config() else {}
-    _compile_environment = os.environ.copy()
-    _compile_environment.update(config_env)
+    _compile_env = os.environ.copy()
+    _compile_env.update(config_env)
 
 
 try:
@@ -79,7 +79,7 @@ def link(object_files, destination_file=None, overwrite_destination_file=True, a
                    _output_flag,
                    destination_file]  # /out: for msvc???
 
-        subprocess.check_call(command, env=_compile_environment)
+        subprocess.check_call(command, env=_compile_env)
 
     return destination_file
 
@@ -151,7 +151,7 @@ def compile_file(file, use_nvcc=False, nvcc='nvcc', overwrite_destination_file=T
                    destination_file]
 
     if not exists(destination_file) or overwrite_destination_file:
-        subprocess.check_call(command, env=_compile_environment)
+        subprocess.check_call(command, env=_compile_env)
 
     return destination_file
 

tests/test_native_tensorflow_compilation.py

@@ -16,17 +16,16 @@ from sysconfig import get_paths
 
 import pytest
 import sympy
-from pystencils_autodiff import create_backward_assignments
-from pystencils_autodiff._file_io import write_file
-from pystencils_autodiff.backends.astnodes import TensorflowModule
 
 import pystencils
 from pystencils.cpu.cpujit import get_compiler_config
 from pystencils.include import get_pystencils_include_path
+from pystencils_autodiff import create_backward_assignments
+from pystencils_autodiff._file_io import write_file
+from pystencils_autodiff.backends.astnodes import TensorflowModule
 from pystencils_autodiff.tensorflow_jit import _compile_env
 
 
-
 def test_detect_cpu_vs_cpu():
 
     for target in ('cpu', 'gpu'):