Merge pull request tensorflow#1573 from RenuPatelGoogle:patch-1

PiperOrigin-RevId: 453217973
xgxg1314 · Jun 6, 2022 · 6685d04 · 6685d04
2 parents b689df6 + 9cdc4ee
commit 6685d04
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/tensorflow_probability/python/mcmc/diagnostic.py b/tensorflow_probability/python/mcmc/diagnostic.py
@@ -160,13 +160,14 @@ def effective_sample_size(states,
         target_log_prob_fn=target.log_prob,
         step_size=0.05,
         num_leapfrog_steps=20))
-  states.shape
+  print(states.shape)
   ==> (1000, 2)
 
   ess = effective_sample_size(states, filter_beyond_positive_pairs=True)
-  ==> Shape (2,) Tensor
+  print(ess.shape)
+  ==> (2,)
 
-  mean, variance = tf.nn.moments(states, axis=0)
+  mean, variance = tf.nn.moments(states, axes=0)
   standard_error = tf.sqrt(variance / ess)
   ```