merged conflicts

Author: gisilvs

.github/workflows/continuous-integration.yml

@@ -68,3 +68,9 @@ jobs:
         run: |
           source ${TEST_VENV_PATH}/bin/activate
           ./testing/run_github_tests.sh
+      - name: Upload test logs
+        if: failure()
+        uses: actions/upload-artifact@v1
+        with:
+            name: testlogs-${{ matrix.shard }}
+            path: bazel-testlogs

discussion/fun_mcmc/BUILD

@@ -137,7 +137,6 @@ py_library(
     ],
 )
 
-# pytype
 py_library(
     name = "backend_jax",
     srcs = ["backend_jax.py"],
@@ -182,6 +181,7 @@ py_library(
 
 py_test(
     name = "fun_mcmc_test",
+    size = "large",
     srcs = ["fun_mcmc_test.py"],
     python_version = "PY3",
     shard_count = 8,

discussion/fun_mcmc/backend.py

@@ -40,6 +40,7 @@
     'get_backend',
     'JAX',
     'MANUAL_TRANSFORMS',
+    'prefer_static',
     'set_backend',
     'TENSORFLOW',
     'tf',
@@ -171,6 +172,7 @@ def __getattr__(self, attr):
     return ret
 
 
+prefer_static = _Dispatcher('prefer_static')
 tf = _Dispatcher('tf')
 tfp = _Dispatcher('tfp')
 util = _Dispatcher('util')

discussion/fun_mcmc/backend_jax.py

@@ -17,12 +17,14 @@
 from discussion.fun_mcmc import tf_on_jax
 from discussion.fun_mcmc import util_jax as util
 from tensorflow_probability.substrates import jax as tfp
+from tensorflow_probability.substrates.jax.internal import prefer_static
 
 tf = tf_on_jax.tf
 
 __all__ = [
     'BACKEND_NAME',
     'multi_backend_test',
+    'prefer_static',
     'tf',
     'tfp',
     'util',

discussion/fun_mcmc/backend_tf.py

@@ -17,10 +17,12 @@
 import tensorflow.compat.v2 as tf
 import tensorflow_probability as tfp
 from discussion.fun_mcmc import util_tf as util
+from tensorflow_probability.python.internal import prefer_static
 
 __all__ = [
     'BACKEND_NAME',
     'multi_backend_test',
+    'prefer_static',
     'tf',
     'tfp',
     'util',

discussion/fun_mcmc/fun_mcmc_lib.py

@@ -45,6 +45,7 @@
 
 from discussion.fun_mcmc import backend
 
+ps = backend.prefer_static
 tf = backend.tf
 tfp = backend.tfp
 util = backend.util
@@ -1349,7 +1350,7 @@ def hamiltonian_integrator(
   state_grads = int_state.state_grads
   state_extra = int_state.state_extra
 
-  num_steps = tf.convert_to_tensor(num_steps)
+  num_steps = ps.convert_to_shape_tensor(num_steps)
   is_ragged = len(num_steps.shape) > 0  # pylint: disable=g-explicit-length-test
 
   kinetic_energy, kinetic_energy_extra = call_potential_fn(
@@ -1358,7 +1359,7 @@ def hamiltonian_integrator(
 
   if is_ragged:
     step = 0
-    max_num_steps = tf.reduce_max(num_steps)
+    max_num_steps = ps.reduce_max(num_steps)
   else:
     step = []
     max_num_steps = num_steps

discussion/fun_mcmc/util_jax.py

@@ -71,7 +71,7 @@ def wrapper(_, *rest):
 
 def get_shallow_tree(is_leaf, tree):
   """Returns a shallow tree, expanding only when is_leaf(subtree) is False."""
-  return tree_util.tree_map(is_leaf, tree, is_leaf)
+  return tree_util.tree_map(is_leaf, tree, is_leaf=is_leaf)
 
 
 def assert_same_shallow_tree(shallow, tree):

spinoffs/fun_mc/fun_mc/BUILD

@@ -108,6 +108,7 @@ py_library(
 
 py_test(
     name = "fun_mc_test",
+    size = "large",
     srcs = ["fun_mc_test.py"],
     python_version = "PY3",
     shard_count = 8,

spinoffs/fun_mc/fun_mc/dynamic/backend_jax/backend.py

@@ -17,10 +17,12 @@
 from fun_mc.dynamic.backend_jax import tf_on_jax
 from fun_mc.dynamic.backend_jax import util
 from tensorflow_probability.substrates import jax as tfp
+from tensorflow_probability.substrates.jax.internal import prefer_static
 
 tf = tf_on_jax.tf
 
 __all__ = [
+    'prefer_static',
     'tf',
     'tfp',
     'util',

spinoffs/fun_mc/fun_mc/dynamic/backend_jax/util.py

@@ -71,7 +71,7 @@ def wrapper(_, *rest):
 
 def get_shallow_tree(is_leaf, tree):
   """Returns a shallow tree, expanding only when is_leaf(subtree) is False."""
-  return tree_util.tree_map(is_leaf, tree, is_leaf)
+  return tree_util.tree_map(is_leaf, tree, is_leaf=is_leaf)
 
 
 def assert_same_shallow_tree(shallow, tree):