test_rhat

PiperOrigin-RevId: 738878610

Author: lukmaz

meridian/model/posterior_sampler.py

@@ -20,6 +20,7 @@
 import arviz as az
 from meridian import constants
 import numpy as np
+import pandas as pd
 import tensorflow as tf
 import tensorflow_probability as tfp
 
@@ -463,7 +464,8 @@ def __call__(
 
     states = []
     traces = []
-    for n_chains_batch in n_chains_list:
+    print("\n\n. NEW SEED VERSION\n\n")
+    for num, n_chains_batch in enumerate(n_chains_list):
       try:
         mcmc = _xla_windowed_adaptive_nuts(
             n_draws=n_burnin + n_keep,
@@ -480,6 +482,7 @@ def __call__(
             seed=seed,
             **pins,
         )
+        seed = [x + 1 for x in (seed or [0, 0])]
       except tf.errors.ResourceExhaustedError as error:
         raise MCMCOOMError(
             "ERROR: Out of memory. Try reducing `n_keep` or pass a list of"
@@ -488,6 +491,37 @@ def __call__(
         ) from error
       states.append(mcmc.all_states._asdict())
       traces.append(mcmc.trace)
+      # TODO
+      print(f"\n\n\n BATCH {num}")
+      tmp_state = mcmc.all_states._asdict()
+      tmp_mcmc_states = {
+          k: tf.einsum(
+              "ij...->ji...",
+              tmp_state[k][n_burnin:, ...],
+          )
+          for k in tmp_state.keys()
+          if k not in constants.UNSAVED_PARAMETERS
+      }
+      rhat = pd.DataFrame()
+      vis_mcmc_states = {k: v.values for k, v in tmp_mcmc_states.items()}
+      for k, v in tfp.mcmc.potential_scale_reduction(
+          {k: tf.einsum("ij...->ji...", v) for k, v in vis_mcmc_states.items()}
+      ).items():
+        rhat_temp = v.numpy().flatten()
+        rhat = pd.concat([
+            rhat,
+            pd.DataFrame({
+                constants.PARAMETER: np.repeat(k, len(rhat_temp)),
+                constants.RHAT: rhat_temp,
+            }),
+        ])
+
+      # If the MCMC sampling fails, the r-hat value calculated will be very large.
+      if (rhat[constants.RHAT] > 1e10).any():
+        max_rhat = max(rhat[constants.RHAT])
+        print(f" MAX RHAT: {max_rhat}")
+      print(f" ALL RHATS: {rhat[constants.RHAT]}")
+      # TODO
 
     mcmc_states = {
         k: tf.einsum(