Merge branch 'main' into frighterafix#1505

Author: ColCarroll

discussion/examples/windowed_sampling.ipynb

@@ -1837,9 +1837,9 @@
             "WARNING:tensorflow:Note that RandomStandardNormal inside pfor op may not give same output as inside a sequential loop.\n",
             "Fast window 75\n",
             "Slow window 25\n",
-            "WARNING:tensorflow:5 out of the last 5 calls to \u003cfunction slow_window at 0x7f9456031ea0\u003e triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.\n",
+            "WARNING:tensorflow:5 out of the last 5 calls to \u003cfunction slow_window at 0x7f9456031ea0\u003e triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.\n",
             "Slow window 50\n",
-            "WARNING:tensorflow:6 out of the last 6 calls to \u003cfunction slow_window at 0x7f9456031ea0\u003e triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.\n",
+            "WARNING:tensorflow:6 out of the last 6 calls to \u003cfunction slow_window at 0x7f9456031ea0\u003e triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.\n",
             "Slow window 100\n",
             "Slow window 200\n",
             "Fast window 75\n",

discussion/meads/README.md

@@ -46,7 +46,7 @@ def _sample_posterior(target_log_prob_unconstrained,
                       parallel_iterations=10,
                       jit_compile=True,
                       use_input_signature=False,
-                      experimental_relax_shapes=False):
+                      reduce_retracing=False):
   """MCMC sampling with HMC/NUTS using an expanding epoch tuning scheme."""
 
   seed_stream = tfp.util.SeedStream(seed, 'window_tune_nuts_sampling')
@@ -117,7 +117,7 @@ def _sample_posterior(target_log_prob_unconstrained,
       input_signature=input_signature,
       autograph=False,
       jit_compile=jit_compile,
-      experimental_relax_shapes=experimental_relax_shapes)
+      reduce_retracing=reduce_retracing)
   def fast_adaptation_interval(num_steps, previous_state):
     """Step size only adaptation interval.
 
@@ -179,7 +179,7 @@ def body_fn_window2(
       input_signature=input_signature,
       autograph=False,
       jit_compile=jit_compile,
-      experimental_relax_shapes=experimental_relax_shapes)
+      reduce_retracing=reduce_retracing)
   def slow_adaptation_interval(num_steps, previous_n, previous_state,
                                previous_mean, previous_cov):
     """Interval that tunes the mass matrix and step size simultaneously.
@@ -328,7 +328,7 @@ def window_tune_nuts_sampling(target_log_prob,
                               parallel_iterations=10,
                               jit_compile=True,
                               use_input_signature=True,
-                              experimental_relax_shapes=False):
+                              reduce_retracing=False):
   """Sample from a density with NUTS and an expanding window tuning scheme.
 
   This function implements a turnkey MCMC sampling routine using NUTS and an
@@ -347,7 +347,7 @@ def window_tune_nuts_sampling(target_log_prob,
   of the tuning epoch (window 1, 2, and 3 in Stan [1]) run with two @tf.function
   compiled functions. The user can control the compilation options using the
   kwargs `jit_compile`, `use_input_signature`, and
-  `experimental_relax_shapes`.  Setting all to True would compile to XLA and
+  `reduce_retracing`.  Setting all to True would compile to XLA and
   potentially avoid the small overhead of function recompilation (note that it
   is not yet the case in XLA right now). It is not yet clear whether doing it
   this way is better than just wrapping the full inference routine in
@@ -403,7 +403,7 @@ def window_tune_nuts_sampling(target_log_prob,
       function is always compiled by XLA.
     use_input_signature: If True, generate an input_signature kwarg to pass to
       tf.function decorator.
-    experimental_relax_shapes: kwarg pass to tf.function decorator. When True,
+    reduce_retracing: kwarg pass to tf.function decorator. When True,
       tf.function may generate fewer, graphs that are less specialized on input
       shapes.
 
@@ -564,6 +564,6 @@ def target_log_prob_unconstrained_concated(x):
       parallel_iterations=parallel_iterations,
       jit_compile=jit_compile,
       use_input_signature=use_input_signature,
-      experimental_relax_shapes=experimental_relax_shapes)
+      reduce_retracing=reduce_retracing)
   return forward_transform(
       split_and_reshape(nuts_samples)), diagnostic, conditioning_bijector

discussion/meads/meads.ipynb

@@ -47,7 +47,7 @@
 
 def map_tree(fn, tree, *args):
   """Maps `fn` over the leaves of a nested structure."""
-  return tree_util.tree_multimap(fn, tree, *args)
+  return tree_util.tree_map(fn, tree, *args)
 
 
 def flatten_tree(tree):
@@ -66,7 +66,7 @@ def map_tree_up_to(shallow, fn, tree, *rest):
   def wrapper(_, *rest):
     return fn(*rest)
 
-  return tree_util.tree_multimap(wrapper, shallow, tree, *rest)
+  return tree_util.tree_map(wrapper, shallow, tree, *rest)
 
 
 def get_shallow_tree(is_leaf, tree):
@@ -76,7 +76,7 @@ def get_shallow_tree(is_leaf, tree):
 
 def assert_same_shallow_tree(shallow, tree):
   """Asserts that `tree` has the same shallow structure as `shallow`."""
-  # Do a dummy multimap for the side-effect of verifying that the structures are
+  # Do a dummy map for the side-effect of verifying that the structures are
   # the same. This doesn't catch all the errors we actually care about, sadly.
   map_tree_up_to(shallow, lambda *args: (), tree)
 

discussion/turnkey_inference_candidate/window_tune_nuts_sampling.py

@@ -1,4 +1,3 @@
-# Lint as: python3
 # Copyright 2020 The TensorFlow Probability Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

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

@@ -1,4 +1,3 @@
-# Lint as: python3
 # Copyright 2020 The TensorFlow Probability Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

spinoffs/inference_gym/inference_gym/backends/jax_integration_test.py

@@ -1,4 +1,3 @@
-# Lint as: python3
 # Copyright 2020 The TensorFlow Probability Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

spinoffs/inference_gym/inference_gym/backends/numpy_integration_test.py

@@ -1,4 +1,3 @@
-# Lint as: python3
 # Copyright 2020 The TensorFlow Probability Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

spinoffs/inference_gym/inference_gym/backends/rewrite.py

@@ -1,4 +1,3 @@
-# Lint as: python3
 # Copyright 2020 The TensorFlow Probability Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");