Export default tracing functions in mcmc module

PiperOrigin-RevId: 721125228
tensorflow · Jan 29, 2025 · 3ea485c · 3ea485c
1 parent 21e1c4c
commit 3ea485c
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/tensorflow_probability/python/experimental/mcmc/__init__.py b/tensorflow_probability/python/experimental/mcmc/__init__.py
@@ -70,6 +70,8 @@
 from tensorflow_probability.python.experimental.mcmc.weighted_resampling import resample_independent
 from tensorflow_probability.python.experimental.mcmc.weighted_resampling import resample_stratified
 from tensorflow_probability.python.experimental.mcmc.weighted_resampling import resample_systematic
+from tensorflow_probability.python.experimental.mcmc.windowed_sampling import default_hmc_trace_fn
+from tensorflow_probability.python.experimental.mcmc.windowed_sampling import default_nuts_trace_fn
 from tensorflow_probability.python.experimental.mcmc.windowed_sampling import windowed_adaptive_hmc
 from tensorflow_probability.python.experimental.mcmc.windowed_sampling import windowed_adaptive_nuts
 from tensorflow_probability.python.experimental.mcmc.with_reductions import WithReductions
@@ -135,6 +137,8 @@
     'simple_heuristic_tuning',
     'snaper_criterion',
     'step_kernel',
+    'default_hmc_trace_fn',
+    'default_nuts_trace_fn',
     'windowed_adaptive_hmc',
-    'windowed_adaptive_nuts'
-    ]
+    'windowed_adaptive_nuts',
+]