fix jax backend tolist for tracers in logging

Author: pfackeldey

src/pyhf/tensor/jax_backend.py

@@ -2,6 +2,7 @@
 
 config.update('jax_enable_x64', True)
 
+import jax
 from jax import Array
 import jax.numpy as jnp
 from jax.scipy.special import gammaln, xlogy
@@ -14,6 +15,10 @@
 log = logging.getLogger(__name__)
 
 
+def currently_jitting():
+    return isinstance(jnp.array(1) + 1, jax.core.Tracer)
+
+
 class _BasicPoisson:
     def __init__(self, rate):
         self.rate = rate
@@ -184,6 +189,9 @@ def conditional(self, predicate, true_callable, false_callable):
         return true_callable() if predicate else false_callable()
 
     def tolist(self, tensor_in):
+        if currently_jitting():
+            # .aval is the abstract value and has a little nicer representation
+            return tensor_in.aval
         try:
             return jnp.asarray(tensor_in).tolist()
         except (TypeError, ValueError):