Draft: Tobi Dance

src/SI_Toolkit/Functions/TF/Compile.py

@@ -4,16 +4,21 @@
 import tensorflow as tf
 import torch
 
+from Control_Toolkit.others.globals_and_utils import get_logger
+log=get_logger(__name__)
+
 from SI_Toolkit.computation_library import ComputationLibrary
 
+
+
 try:
     from SI_Toolkit_ASF import GLOBALLY_DISABLE_COMPILATION, USE_JIT_COMPILATION
 except ImportError:
-    logging.warn("No compilation option set in SI_Toolkit_ASF. Setting GLOBALLY_DISABLE_COMPILATION to True.")
+    log.warn("No compilation option set in SI_Toolkit_ASF/__init.py__. Setting GLOBALLY_DISABLE_COMPILATION to True.")
     GLOBALLY_DISABLE_COMPILATION = True
 
 def tf_function_jit(func):
-    return tf.function(func=func, jit_compile=True)
+    return tf.function(func=func, jit_compile=True,)
 
 
 def tf_function_experimental(func):
@@ -25,6 +30,7 @@ def identity(func):
 
 
 if GLOBALLY_DISABLE_COMPILATION:
+    log.info('TensorFlow compilation is disabled by GLOBALLY_DISABLE_COMPILATION=True')
     CompileTF = identity
 else:
     if platform.machine() == 'arm64' and platform.system() == 'Darwin':  # For M1 Apple processor
@@ -33,9 +39,12 @@ def identity(func):
         CompileTF = tf.function
     else:
         CompileTF = tf_function_jit
+    log.info(f'using {CompileTF} compilation')
         # CompileTF = tf_function_experimental # Should be same as tf_function_jit, not appropriate for newer version of TF
 
 def CompileAdaptive(fun):
+    """ TODO add docstring to explain what it does and where it is used
+    """
     instance = fun.__self__
     assert hasattr(instance, "lib"), "Instance with this method has no computation library defined"
     computation_library: "type[ComputationLibrary]" = instance.lib
@@ -44,7 +53,8 @@ def CompileAdaptive(fun):
     if GLOBALLY_DISABLE_COMPILATION:
         return identity(fun)
     elif lib_name == 'TF':
+        log.info(f'compiling tensorflow {fun}')
         return CompileTF(fun)
     else:
-        print('Jit compilation for Pytorch not yet implemented.')
+        log.warning(f'JIT compilation for {lib_name} not yet implemented.')
         return identity(fun)

src/SI_Toolkit/Predictors/predictor_wrapper.py

@@ -29,6 +29,7 @@ def __init__(self):
 
     def configure(self, batch_size: int, horizon: int, dt: float, computation_library: "Optional[type[ComputationLibrary]]"=None, predictor_specification=None, compile_standalone=False):
 
+
         self.update_predictor_config_from_specification(predictor_specification)
 
         compile_standalone = {'disable_individual_compilation': not compile_standalone}
@@ -53,11 +54,16 @@ def configure(self, batch_size: int, horizon: int, dt: float, computation_librar
             self.predictor = predictor_ODE_tf(horizon=self.horizon, dt=dt, batch_size=self.batch_size, **self.predictor_config, **compile_standalone)
 
         else:
-            raise NotImplementedError('Type of the predictor not recognised.')
-        
+            raise NotImplementedError(f'Type of the predictor {self.predictor_type} is not recognised.')
+
         # computation_library defaults to None. In that case, do not check for conformity.
-        if computation_library is not None and computation_library not in self.predictor.supported_computation_libraries:
-            raise ValueError(f"Predictor {self.predictor.__class__.__name__} does not support {computation_library.__name__}")
+        # in other cases, check after we configure it to make sure it supports itself
+        if not computation_library is None and computation_library not in self.predictor.supported_computation_libraries:
+            raise ValueError(
+                f"Predictor {self.predictor.__class__.__name__} does not support {computation_library.__name__}")
+
+        self.predictor.lib=computation_library # set the library type on the predictor object so we can use it to assign attributes later
+
 
     def configure_with_compilation(self, batch_size, horizon, dt, predictor_specification=None):
         """

src/SI_Toolkit/computation_library.py

@@ -54,7 +54,7 @@ class ComputationLibrary:
     gather: Callable[[TensorType, TensorType, int], TensorType] = None
     gather_last: Callable[[TensorType, TensorType], TensorType] = None
     arange: Callable[[Optional[NumericType], NumericType, Optional[NumericType]], TensorType] = None
-    zeros: Callable[["tuple[int]"], TensorType] = None
+    zeros: Callable[["tuple[int,...]"], TensorType] = None
     zeros_like: Callable[[TensorType], TensorType] = None
     ones: Callable[["tuple[int]"], TensorType] = None
     sign: Callable[[TensorType], TensorType] = None
@@ -88,6 +88,9 @@ class ComputationLibrary:
     dot: Callable[[TensorType, TensorType], TensorType] = None
     stop_gradient: Callable[[TensorType], TensorType] = None
     assign: Callable[[Union[TensorType, tf.Variable], TensorType], Union[TensorType, tf.Variable]] = None
+    nan:TensorType=None
+    isnan:Callable[[TensorType],bool]=None
+    string = None
 
 
 class NumpyLibrary(ComputationLibrary):
@@ -158,14 +161,16 @@ class NumpyLibrary(ComputationLibrary):
     dot = np.dot
     stop_gradient = lambda x: x
     assign = LibraryHelperFunctions.set_to_value
-
+    nan = np.nan
+    isnan=np.isnan
+    string=str
 
 class TensorFlowLibrary(ComputationLibrary):
     lib = 'TF'
     reshape = tf.reshape
     permute = tf.transpose
     newaxis = tf.newaxis
-    shape = lambda x: x.get_shape()  # .as_list()
+    shape = tf.shape #  tobi does not understand reason for this previous definition: # lambda x: x.get_shape()  # .as_list()
     to_numpy = lambda x: x.numpy()
     to_variable = lambda x, dtype: tf.Variable(x, dtype=dtype)
     to_tensor = lambda x, dtype: tf.convert_to_tensor(x, dtype=dtype)
@@ -228,6 +233,9 @@ class TensorFlowLibrary(ComputationLibrary):
     dot = lambda a, b: tf.tensordot(a, b, 1)
     stop_gradient = tf.stop_gradient
     assign = LibraryHelperFunctions.set_to_variable
+    nan=tf.constant(np.nan)
+    isnan=tf.math.is_nan
+    string=tf.string
 
 
 class PyTorchLibrary(ComputationLibrary):
@@ -307,3 +315,6 @@ def gather_last_pytorch(a, index_vector):
     dot = torch.dot
     stop_gradient = tf.stop_gradient # FIXME: How to imlement this in torch?
     assign = LibraryHelperFunctions.set_to_value
+    nan=torch.nan
+    isnan=torch.isnan
+    string=lambda x: torch.ByteTensor(bytes(x,'utf8'))