Support Python 3.9

[matthewfeickert]

Description

TensorFlow finally is supporting Python 3.9 in v2.5.0 so it would be good to start supporting this release of TensorFlow and then to cut a release that supports Python 3.9.

The main roadblock to this is Issue #293, which has plagued pyhf for a long time. This Issue should get highest priority.

cc @lukasheinrich @kratsg

[matthewfeickert]

@lukasheinrich @kratsg If we jump to tensorflow v2.6.0 and tensorflow-probability v0.13.0 then we can get there now with the following

diff --git a/src/pyhf/tensor/tensorflow_backend.py b/src/pyhf/tensor/tensorflow_backend.py
index d5c56647..2547b3d2 100644
--- a/src/pyhf/tensor/tensorflow_backend.py
+++ b/src/pyhf/tensor/tensorflow_backend.py
@@ -2,6 +2,7 @@
 import logging
 import tensorflow as tf
 import tensorflow_probability as tfp
+from numpy import nan
 
 log = logging.getLogger(__name__)
 
@@ -119,7 +120,7 @@ def tile(self, tensor_in, repeats):
         """
         try:
             return tf.tile(tensor_in, repeats)
-        except tf.python.framework.errors_impl.InvalidArgumentError:
+        except tf.errors.InvalidArgumentError:
             shape = tf.shape(tensor_in).numpy().tolist()
             diff = len(repeats) - len(shape)
             if diff < 0:
@@ -426,8 +427,13 @@ def poisson_logpdf(self, n, lam):
             TensorFlow Tensor: Value of the continuous approximation to log(Poisson(n|lam))
         """
         lam = self.astensor(lam)
+        valid_obs_given_rate = tf.logical_or(
+            tf.math.not_equal(lam, n), tf.math.not_equal(n, 0)
+        )
 
-        return tfp.distributions.Poisson(lam).log_prob(n)
+        return tf.where(
+            valid_obs_given_rate, tfp.distributions.Poisson(lam).log_prob(n), nan
+        )
 
     def poisson(self, n, lam):
         r"""
@@ -457,8 +463,15 @@ def poisson(self, n, lam):
             TensorFlow Tensor: Value of the continuous approximation to Poisson(n|lam)
         """
         lam = self.astensor(lam)
+        valid_obs_given_rate = tf.logical_or(
+            tf.math.not_equal(lam, n), tf.math.not_equal(n, 0)
+        )
 
-        return tf.exp(tfp.distributions.Poisson(lam).log_prob(n))
+        return tf.where(
+            valid_obs_given_rate,
+            tf.exp(tfp.distributions.Poisson(lam).log_prob(n)),
+            nan,
+        )
 
     def normal_logpdf(self, x, mu, sigma):
         r"""

as tf.where and tf.logical_or are differentiable operations in TensorFlow. This would allow us to move forward with TensorFlow support for Python 3.9 (which will unblock things) still and then revisit Issue #293 later.

[kratsg]

Let's do that in v0.6.4+

[matthewfeickert]

Let's do that in v0.6.4+

Fair. I might still actually do everything tonight so that I get it all out of my head and into PRs, but we can get v0.6.3 out first.