Merge pull request #151 from rs-station/dw_cpu

Dw cpu resolves #150
rs-station · Jan 18, 2024 · 39941c7 · 39941c7
2 parents ecdf8c0 + 486136a
commit 39941c7
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/careless/models/priors/wilson.py b/careless/models/priors/wilson.py
@@ -136,7 +136,14 @@ def stddev(self):
         return self.wilson_prior.stddev()
 
     def log_prob(self, z):
-        z_parent = tf.gather(z, self.reflids, axis=-1)
+        mask = self.reflids >= 0
+        sanitized_reflids = tf.where(mask, self.reflids, 0)
+        z_parent = tf.where(
+            mask[None,:], 
+            tf.gather(z, sanitized_reflids, axis=-1),
+            0.,
+        )
+
         loc = tf.where(
             self.absent, 
             0.,

diff --git a/setup.py b/setup.py
@@ -41,7 +41,7 @@ def getVersionNumber():
     install_requires=[
         "reciprocalspaceship>=0.9.16",
         "tqdm",
-        "tensorflow>=2.8",
+        "tensorflow<2.16",
         "tensorflow-probability",
         "matplotlib",
         "seaborn",