check backend compatability when astroNN is imported

src/astroNN/__init__.py

@@ -4,4 +4,20 @@
 
 from importlib.metadata import version
 
+import keras
+
 version = __version__ = version("astroNN")
+_KERAS_BACKEND = keras.backend.backend()
+
+# check if the backend is tensorflow or pytorch
+def check_backend():
+    """
+    Check if the backend is tensorflow or pytorch
+    """
+    if _KERAS_BACKEND != "tensorflow" and _KERAS_BACKEND != "torch":
+        raise ImportError(
+            f"astroNN only supports PyTorch or Tensorflow backend, but your current backend is {keras.backend.backend()}"
+        )
+
+
+check_backend()

src/astroNN/config.py

@@ -1,22 +1,12 @@
 import configparser
 import os
 import platform
-
-import keras
-import numpy as np
+from astroNN import _KERAS_BACKEND
 import importlib
 
+backend_framework = importlib.import_module(_KERAS_BACKEND)
 astroNN_CACHE_DIR = os.path.join(os.path.expanduser("~"), ".astroNN")
 _astroNN_MODEL_NAME = "model_weights.keras"  # default astroNN model filename
-_KERAS_BACKEND = keras.backend.backend()
-
-supported_backend = ["tensorflow", "torch", "jax"]
-if _KERAS_BACKEND not in supported_backend:
-    raise ImportError(
-        f"astroNN only support {supported_backend} backend, currently you have '{keras.backend.backend()}' as backend"
-    )
-else:
-    backend_framework = importlib.import_module(_KERAS_BACKEND)
 
 
 def config_path(flag=None):
@@ -249,3 +239,5 @@ def cpu_gpu_reader():
 MULTIPROCESS_FLAG = multiprocessing_flag_reader()
 ENVVAR_WARN_FLAG = envvar_warning_flag_reader()
 CUSTOM_MODEL_PATH = custom_model_path_reader()
+
+__all__ = [_KERAS_BACKEND, MAGIC_NUMBER, MULTIPROCESS_FLAG, ENVVAR_WARN_FLAG, CUSTOM_MODEL_PATH, backend_framework]

src/astroNN/models/base_master_nn.py

@@ -16,7 +16,10 @@
 from astroNN.config import _astroNN_MODEL_NAME, cpu_gpu_reader
 from astroNN.shared.nn_tools import cpu_fallback
 from astroNN.shared.nn_tools import folder_runnum
-
+from astroNN.config import (
+    _KERAS_BACKEND,
+    backend_framework,
+)
 epsilon, plot_model = keras.backend.epsilon, keras.utils.plot_model
 
 
@@ -508,24 +511,20 @@ def jacobian(self, x=None, mean_output=False, mc_num=1, denormalize=False):
 
         start_time = time.time()
 
-        if keras.backend.backend() == "tensorflow":
-            import tensorflow as tf
-
-            xtensor = tf.Variable(x_data)
+        if _KERAS_BACKEND == "tensorflow":
+            xtensor = backend_framework.Variable(x_data)
 
-            with tf.GradientTape(watch_accessed_variables=False) as tape:
+            with backend_framework.GradientTape(watch_accessed_variables=False) as tape:
                 tape.watch(xtensor)
                 temp = _model(xtensor)
                 if isinstance(temp, dict):
                     temp = temp["output"]
 
             jacobian = tape.batch_jacobian(temp, xtensor)
-        elif keras.backend.backend() == "torch":
-            import torch
-
+        elif _KERAS_BACKEND == "torch":
             # add new axis for vmap
-            xtensor = torch.tensor(x_data, requires_grad=True)[:, None, ...]
-            jacobian = torch.vmap(torch.func.jacrev(_model), randomness="different")(xtensor)
+            xtensor = backend_framework.tensor(x_data, requires_grad=True)[:, None, ...]
+            jacobian = backend_framework.vmap(backend_framework.func.jacrev(_model), randomness="different")(xtensor)
         else:
             raise ValueError("Only Tensorflow and PyTorch backend is supported")
 

src/astroNN/models/base_vae.py

@@ -289,7 +289,7 @@ def custom_train_step(self, data):
         y = keras.ops.cast(y["output"], backend_framework.float32)
 
         # Run forward pass.
-        if keras.backend.backend() == "tensorflow":
+        if _KERAS_BACKEND == "tensorflow":
             with backend_framework.GradientTape() as tape:
                 encoder_output = self.keras_encoder(x, training=True)
                 if isinstance(encoder_output, dict):
@@ -313,7 +313,7 @@ def custom_train_step(self, data):
             self.keras_model.optimizer.apply_gradients(
                 zip(gradients, self.keras_model.trainable_weights)
             )
-        elif keras.backend.backend() == "torch":
+        elif _KERAS_BACKEND == "torch":
             self.keras_model.zero_grad()
             encoder_output = self.keras_encoder(x, training=True)
             if isinstance(encoder_output, dict):
@@ -936,28 +936,24 @@ def jacobian_latent(self, x=None, mean_output=False, mc_num=1, denormalize=False
 
         start_time = time.time()
 
-        if keras.backend.backend() == "tensorflow":
-            import tensorflow as tf
+        if _KERAS_BACKEND == "tensorflow":
+            xtensor = backend_framework.Variable(x_data)
 
-            xtensor = tf.Variable(x_data)
-
-            with tf.GradientTape(watch_accessed_variables=False) as tape:
+            with backend_framework.GradientTape(watch_accessed_variables=False) as tape:
                 tape.watch(xtensor)
                 temp = _model(xtensor)
 
-            jacobian = tf.squeeze(tape.batch_jacobian(temp, xtensor))
-        elif keras.backend.backend() == "torch":
-            import torch
-
-            xtensor = torch.tensor(x_data, requires_grad=True)
-            jacobian = torch.autograd.functional.jacobian(_model, xtensor)
+            jacobian = backend_framework.squeeze(tape.batch_jacobian(temp, xtensor))
+        elif _KERAS_BACKEND == "torch":
+            xtensor = backend_framework.tensor(x_data, requires_grad=True)
+            jacobian = backend_framework.autograd.functional.jacobian(_model, xtensor)
         else:
             raise ValueError("Only Tensorflow and PyTorch backend is supported")
 
-        jacobian = tf.squeeze(tape.batch_jacobian(temp, xtensor))
+        jacobian = backend_framework.squeeze(tape.batch_jacobian(temp, xtensor))
 
         if mean_output is True:
-            jacobian_master = tf.reduce_mean(jacobian, axis=0).numpy()
+            jacobian_master = backend_framework.reduce_mean(jacobian, axis=0).numpy()
         else:
             jacobian_master = jacobian.numpy()
 

src/astroNN/nn/layers.py

@@ -1,7 +1,7 @@
 import math
 
 import keras
-from astroNN.config import backend_framework
+from astroNN.config import _KERAS_BACKEND, backend_framework
 
 class KLDivergenceLayer(keras.layers.Layer):
     """
@@ -337,12 +337,12 @@ def loop_fn(i):
             return self.layer(inputs)
 
         # vectorizing operation depends on backend
-        if keras.backend.backend() == "torch":
+        if _KERAS_BACKEND == "torch":
             outputs = backend_framework.vmap(
                 loop_fn, randomness="different", in_dims=0
             )(self.arange_n)
         # TODO: tensorflow vectorized_map traced operation so there is no randomness which affects e.g., dropout
-        # elif keras.backend.backend() == "tensorflow":
+        # elif _KERAS_BACKEND == "tensorflow":
         #     outputs = backend_framework.vectorized_map(loop_fn, self.arange_n)
         else:  # fallback to simple for loop
             outputs = [self.layer(inputs) for _ in range(self.n)]