Use jit_compile= instead of experimental_compile= in tf.function.

PiperOrigin-RevId: 380008718

Author: jburnim

tensorflow_probability/examples/jupyter_notebooks/Bayesian_Switchpoint_Analysis.ipynb

@@ -320,7 +320,7 @@
         "num_results = 10000\n",
         "num_burnin_steps = 3000\n",
         "\n",
-        "@tf.function(autograph=False, experimental_compile=True)\n",
+        "@tf.function(autograph=False, jit_compile=True)\n",
         "def make_chain(target_log_prob_fn):\n",
         "   kernel = tfp.mcmc.TransformedTransitionKernel(\n",
         "       inner_kernel=tfp.mcmc.HamiltonianMonteCarlo(\n",

tensorflow_probability/examples/jupyter_notebooks/Eight_Schools.ipynb

@@ -274,7 +274,7 @@
         "\n",
         "# Improve performance by tracing the sampler using `tf.function`\n",
         "# and compiling it using XLA.\n",
-        "@tf.function(autograph=False, experimental_compile=True)\n",
+        "@tf.function(autograph=False, jit_compile=True)\n",
         "def do_sampling():\n",
         "  return tfp.mcmc.sample_chain(\n",
         "      num_results=num_results,\n",

tensorflow_probability/examples/jupyter_notebooks/Gaussian_Process_Latent_Variable_Model.ipynb

@@ -347,7 +347,7 @@
         "\n",
         "optimizer = tf.optimizers.Adam(learning_rate=1.0)\n",
         "\n",
-        "@tf.function(autograph=False, experimental_compile=True)\n",
+        "@tf.function(autograph=False, jit_compile=True)\n",
         "def train_model():\n",
         "  with tf.GradientTape() as tape:\n",
         "    loss_value = loss_fn()\n",

tensorflow_probability/examples/jupyter_notebooks/Gaussian_Process_Regression_In_TFP.ipynb

@@ -542,7 +542,7 @@
       "outputs": [],
       "source": [
         "# Use `tf.function` to trace the loss for more efficient evaluation.\n",
-        "@tf.function(autograph=False, experimental_compile=False)\n",
+        "@tf.function(autograph=False, jit_compile=False)\n",
         "def target_log_prob(amplitude, length_scale, observation_noise_variance):\n",
         "  return gp_joint_model.log_prob({\n",
         "      'amplitude': amplitude,\n",
@@ -790,7 +790,7 @@
       ],
       "source": [
         "# Speed up sampling by tracing with `tf.function`.\n",
-        "@tf.function(autograph=False, experimental_compile=False)\n",
+        "@tf.function(autograph=False, jit_compile=False)\n",
         "def do_sampling():\n",
         "  return tfp.mcmc.sample_chain(\n",
         "      kernel=adaptive_sampler,\n",

tensorflow_probability/examples/jupyter_notebooks/Linear_Mixed_Effects_Models.ipynb

@@ -721,7 +721,7 @@
         "    num_leapfrog_steps=3)\n",
         "kernel_results = hmc.bootstrap_results(current_state)\n",
         "\n",
-        "@tf.function(autograph=False, experimental_compile=True)\n",
+        "@tf.function(autograph=False, jit_compile=True)\n",
         "def one_e_step(current_state, kernel_results):\n",
         "  next_state, next_kernel_results = hmc.one_step(\n",
         "      current_state=current_state,\n",
@@ -731,7 +731,7 @@
         "optimizer = tf.optimizers.Adam(learning_rate=.01)\n",
         "\n",
         "# Set up M-step (gradient descent).\n",
-        "@tf.function(autograph=False, experimental_compile=True)\n",
+        "@tf.function(autograph=False, jit_compile=True)\n",
         "def one_m_step(current_state):\n",
         "  with tf.GradientTape() as tape:\n",
         "    loss = -target_log_prob_fn(*current_state)\n",
@@ -843,7 +843,7 @@
       },
       "outputs": [],
       "source": [
-        "@tf.function(autograph=False, experimental_compile=True)\n",
+        "@tf.function(autograph=False, jit_compile=True)\n",
         "def run_k_e_steps(k, current_state, kernel_results):\n",
         "  _, next_state, next_kernel_results = tf.while_loop(\n",
         "      cond=lambda i, state, pkr: i \u003c k,\n",

tensorflow_probability/examples/jupyter_notebooks/STS_approximate_inference_for_models_with_non_Gaussian_observations.ipynb

@@ -359,7 +359,7 @@
    ],
    "source": [
     "# Speed up sampling by tracing with `tf.function`.\n",
-    "@tf.function(autograph=False, experimental_compile=True)\n",
+    "@tf.function(autograph=False, jit_compile=True)\n",
     "def do_sampling():\n",
     "  return tfp.mcmc.sample_chain(\n",
     "      kernel=adaptive_sampler,\n",

tensorflow_probability/examples/jupyter_notebooks/Structural_Time_Series_Modeling_Case_Studies_Atmospheric_CO2_and_Electricity_Demand.ipynb

@@ -535,7 +535,7 @@
         "\n",
         "optimizer = tf.optimizers.Adam(learning_rate=.1)\n",
         "# Using fit_surrogate_posterior to build and optimize the variational loss function.\n",
-        "@tf.function(experimental_compile=True)\n",
+        "@tf.function(jit_compile=True)\n",
         "def train():\n",
         "  elbo_loss_curve = tfp.vi.fit_surrogate_posterior(\n",
         "    target_log_prob_fn=co2_model.joint_log_prob(\n",
@@ -980,7 +980,7 @@
         "\n",
         "optimizer = tf.optimizers.Adam(learning_rate=.1)\n",
         "# Using fit_surrogate_posterior to build and optimize the variational loss function.\n",
-        "@tf.function(experimental_compile=True)\n",
+        "@tf.function(jit_compile=True)\n",
         "def train():\n",
         "  elbo_loss_curve = tfp.vi.fit_surrogate_posterior(\n",
         "    target_log_prob_fn=demand_model.joint_log_prob(\n",

tensorflow_probability/python/experimental/sts_gibbs/benchmarks_test.py

@@ -84,7 +84,7 @@ def test_benchmark_sampling_with_xla(self):
     samples = tf.function(
         gibbs_sampler.fit_with_gibbs_sampling,
         autograph=False,
-        experimental_compile=True)(
+        jit_compile=True)(
             model,
             tfp.sts.MaskedTimeSeries(observed_time_series[..., tf.newaxis],
                                      is_missing),