Easyblock for JAX

Custom_EasyBlocks/jax.py

@@ -26,16 +26,21 @@ class EB_JAX(EasyBlock):
         if not self.cfg['prebuildopts']:
             self.cfg['prebuildopts'] = 'export BAZEL_LINKOPTS=-static-libstdc++:-static-libgcc; '
             self.cfg['prebuildopts'] = 'export BAZEL_LINKLIBS=-l%:libstdc++.a:-lm BAZEL_CXXOPTS=-std=gnu++0x ;'
-            self.cfg['prebuildopts'] = 'export TF_CUDA_PATHS="{}" '.format(cuda)
-            self.cfg['prebuildopts'] += 'GCC_HOST_COMPILER_PREFIX="{}/bin" '.format(binutils)
+            self.cfg['prebuildopts'] = 'export TF_CUDA_PATHS="{}" '.format(
+                cuda)
+            self.cfg['prebuildopts'] += 'GCC_HOST_COMPILER_PREFIX="{}/bin" '.format(
+                binutils)
             # To prevent bazel builds on different hosts/architectures conflicting with each other
             # we'll set HOME, inside which Bazel puts active files (in ~/.cache/bazel/...)
-            self.cfg['prebuildopts'] += 'HOME="{}/fake_home" && '.format(self.builddir)
+            self.cfg['prebuildopts'] += 'HOME="{}/fake_home" && '.format(
+                self.builddir)
 
         if not self.cfg['buildopts']:
-            self.cfg['buildopts'] = '--enable_cuda --cuda_path "{}" '.format(cuda)
+            self.cfg['buildopts'] = '--enable_cuda --cuda_path "{}" '.format(
+                cuda)
             self.cfg['buildopts'] += '--cudnn_path "{}" '.format(cudnn)
-            self.cfg['buildopts'] += '--bazel_path "{}/bin/bazel" '.format(bazel)
+            self.cfg['buildopts'] += '--bazel_path "{}/bin/bazel" '.format(
+                bazel)
             # Tell Bazel to pass PYTHONPATH through to what it's building, so it can find scipy etc.
             self.cfg['buildopts'] += '--bazel_options=--action_env=PYTHONPATH '
             self.cfg['buildopts'] += '--cuda_compute_capabilities "7.0,7.5,8.0" '
@@ -83,7 +88,9 @@ class EB_JAX(EasyBlock):
         txt = super(EB_JAX, self).make_module_extra()
         cuda = get_software_root('CUDA')
         if cuda:
-            txt += self.module_generator.set_environment('XLA_FLAGS', "--xla_gpu_cuda_data_dir={}/bin".format(cuda))
-        txt += self.module_generator.prepend_paths('PYTHONPATH', det_pylibdir())
+            txt += self.module_generator.set_environment(
+                'XLA_FLAGS', "--xla_gpu_cuda_data_dir={}/bin".format(cuda))
+        txt += self.module_generator.prepend_paths(
+            'PYTHONPATH', det_pylibdir())
 
         return txt