Migrate experimental_relax_shapes to reduce_retracing

PiperOrigin-RevId: 437742142

Author: faizan-m

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/turnkey_inference_candidate/window_tune_nuts_sampling.py

@@ -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

tensorflow_probability/examples/jupyter_notebooks/TFP_Release_Notebook_0_13_0.ipynb

@@ -954,7 +954,7 @@
             "  c1=<tf.Tensor: shape=(), dtype=float32, numpy=1.7696666>,\n",
             "  counts=<tf.Tensor: shape=(10,), dtype=float32, numpy=array([ 6., 10., 23.,  7.,  2., 20., 14., 16., 22., 17.], dtype=float32)>\n",
             ")\n",
-            "WARNING:tensorflow:6 out of the last 6 calls to <function windowed_adaptive_hmc at 0x7fda42bed8c0> 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 <function windowed_adaptive_hmc at 0x7fda42bed8c0> 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",
             "StructTuple(\n",
             "  c0=<tf.Tensor: shape=(), dtype=float32, numpy=0.7161876>,\n",
             "  c1=<tf.Tensor: shape=(), dtype=float32, numpy=1.7696666>,\n",

tensorflow_probability/python/distributions/gamma.py

@@ -510,7 +510,7 @@ def _random_gamma_noncpu(
 # tf.function required to access Grappler's implementation_selector.
 @implementation_selection.never_runs_functions_eagerly
 # TODO(b/163029794): Shape relaxation breaks XLA.
-@tf.function(autograph=False, experimental_relax_shapes=False)
+@tf.function(autograph=False, reduce_retracing=False)
 def _random_gamma_no_gradient(
     shape, concentration, rate, log_rate, seed, log_space):
   """Sample a gamma, CPU specialized to stateless_gamma.

tensorflow_probability/python/internal/backend/numpy/v2.py

@@ -60,7 +60,7 @@
 
 def _function(func=None, input_signature=None, autograph=True,  # pylint: disable=unused-argument
               experimental_autograph_options=None,  # pylint: disable=unused-argument
-              experimental_relax_shapes=False, jit_compile=None):  # pylint: disable=unused-argument
+              reduce_retracing=False, jit_compile=None):  # pylint: disable=unused-argument
   """Like `tf.function`, for JAX."""
   transform = lambda fn: fn
   if jit_compile: