bug fixing

htjb · Jul 25, 2023 · f8735d3 · f8735d3
1 parent 2d4506d
commit f8735d3
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 10 deletions.
diff --git a/margarine/maf.py b/margarine/maf.py
@@ -420,6 +420,13 @@ def _train_step(self, x, w, loss_type, maf):
                     maf.trainable_variables))
             return loss
 
+    @tf.function(jit_compile=True)
+    def call_bij(self, bij, u, min, max):
+        x = _forward_transform(u)
+        x = bij(x)
+        x = _inverse_transform(x, min, max)
+        return x
+
     def __call__(self, u):
 
         r"""
@@ -434,13 +441,6 @@ def __call__(self, u):
 
         """
 
-        @tf.function(jit_compile=True)
-        def call_bij(bij, u, min=self.theta_min, max=self.theta_max):
-            x = _forward_transform(u)
-            x = bij(x)
-            x = _inverse_transform(x, min, max)
-            return x
-
 
         if self.cluster_number is not None:
             len_thetas = [len(self.theta[i]) for i in range(len(self.theta))]
@@ -456,12 +456,15 @@ def call_bij(bij, u, min=self.theta_min, max=self.theta_max):
 
             values = []
             for i in range(len(options)):
-                x = call_bij(self.bij[i], u, min=self.theta_min[i], max=self.theta_max[i]).numpy()
+                x = self.call_bij(self.bij[i], u, 
+                                  min=self.theta_min[i], 
+                                  max=self.theta_max[i]).numpy()
                 values.append(x)
 
             x = np.concatenate(values)
         else:
-            x = call_bij(self.bij, u).numpy()
+            x = self.call_bij(self.bij, u, self.theta_min, 
+                              self.theta_max).numpy()
 
         mask = np.isfinite(x).all(axis=-1)
         return x[mask, ...]

diff --git a/margarine/processing.py b/margarine/processing.py
@@ -26,6 +26,8 @@ def _forward_transform(x, min=0, max=1):
 
     """
     x = tf.cast(x, tf.float32)
+    min = tf.cast(min, tf.float32)
+    max = tf.cast(max, tf.float32)
     x = tfd.Uniform(min, max).cdf(x)
     x = tfd.Normal(0, 1).quantile(x)
     return x
@@ -54,7 +56,6 @@ def _inverse_transform(x, min, max):
 
     """
     x = tfd.Normal(0, 1).cdf(x)
-    print(min)
     min = tf.cast(min, tf.float32)
     max = tf.cast(max, tf.float32)
     x = tfd.Uniform(min, max).quantile(x)