Inference Gym: Fix ground truth for eight schools.

The Stan model had the incorrect prior compared to the TFP code.

The old tests passed presumably because their ESS was so bad that the z-test had no power whatsoever. I boosted the number of leapfrog steps to increase it substantially.

PiperOrigin-RevId: 384281260

Author: SiegeLordEx

spinoffs/inference_gym/inference_gym/targets/eight_schools_test.py

@@ -46,7 +46,7 @@ def testEightSchoolsHMC(self):
         model,
         num_chains=4,
         num_steps=4000,
-        num_leapfrog_steps=3,
+        num_leapfrog_steps=10,
         step_size=0.4,
     )
 

spinoffs/inference_gym/inference_gym/targets/ground_truth/eight_schools.py

@@ -27,58 +27,58 @@
 import numpy as np
 
 IDENTITY_AVG_EFFECT_MEAN = np.array([
-    3.2137590385402826,
+    5.75975286026516,
 ]).reshape(())
 
 IDENTITY_AVG_EFFECT_MEAN_STANDARD_ERROR = np.array([
-    0.011278555132469106,
+    0.01786934334960865,
 ]).reshape(())
 
 IDENTITY_AVG_EFFECT_STANDARD_DEVIATION = np.array([
-    3.972567360644285,
+    5.46540575237222,
 ]).reshape(())
 
 IDENTITY_LOG_STDDEV_MEAN = np.array([
-    2.4598093936288152,
+    2.45326548536025,
 ]).reshape(())
 
 IDENTITY_LOG_STDDEV_MEAN_STANDARD_ERROR = np.array([
-    0.0018535105349442114,
+    0.0019127444310499044,
 ]).reshape(())
 
 IDENTITY_LOG_STDDEV_STANDARD_DEVIATION = np.array([
-    0.5110436981514841,
+    0.5146604914005865,
 ]).reshape(())
 
 IDENTITY_SCHOOL_EFFECTS_MEAN = np.array([
-    13.355582058626903,
-    6.102211570244123,
-    1.024737484826771,
-    5.369929612269985,
-    0.7926315948389041,
-    2.244649709577478,
-    11.763530998898284,
-    6.261180511721017,
+    14.76492662128668,
+    7.1559568069514174,
+    2.5889680578429823,
+    6.556760136283709,
+    1.8189620982794488,
+    3.3973176011560953,
+    12.791827392809507,
+    7.94913889128149,
 ]).reshape((8,))
 
 IDENTITY_SCHOOL_EFFECTS_MEAN_STANDARD_ERROR = np.array([
-    0.02544285870818445,
-    0.014957463084240069,
-    0.021316316118379103,
-    0.01575696229480018,
-    0.01373758544135346,
-    0.01550588664268995,
-    0.017107240447608994,
-    0.02209825465307125,
+    0.024677770813437375,
+    0.014741663125129497,
+    0.02289079528770855,
+    0.015999446293377167,
+    0.014481515793592113,
+    0.016516334290635457,
+    0.01649735808006444,
+    0.02296023745612875,
 ]).reshape((8,))
 
 IDENTITY_SCHOOL_EFFECTS_STANDARD_DEVIATION = np.array([
-    10.807906054798417,
-    7.766699608891697,
-    10.244640099796499,
-    8.219198474822052,
-    7.303418026739246,
-    8.258823240336538,
-    8.19256497826584,
-    10.673579208829045,
+    10.796464460046044,
+    7.813843501513539,
+    10.473283439128995,
+    8.314419652710829,
+    7.45775465347688,
+    8.454005086842871,
+    8.16339794409753,
+    10.913796824041532,
 ]).reshape((8,))

spinoffs/inference_gym/inference_gym/tools/get_ground_truth.py

@@ -23,7 +23,7 @@
 venv=$(mktemp -d)
 virtualenv -p python3.6 $venv
 source $venv/bin/activate
-pip install cmdstanpy==0.9 pandas numpy tf-nightly tfp-nightly tfds-nightly
+pip install 'cmdstanpy>=0.9.0' pandas numpy tf-nightly tfp-nightly tfds-nightly
 install_cmdstan
 
 python -m inference_gym.tools.get_ground_truth \

spinoffs/inference_gym/inference_gym/tools/stan/eight_schools.py

@@ -49,7 +49,7 @@ def eight_schools():
     school_effects <- std_school_effects * exp(log_stddev) + avg_effect;
   }
   model {
-    avg_effect ~ normal(0, 5);
+    avg_effect ~ normal(0, 10);
     log_stddev ~ normal(5, 1);
     std_school_effects ~ normal(0, 1);
     treatment_effects ~ normal(school_effects, treatment_stddevs);