s/total_shape/concentration_shape/

PiperOrigin-RevId: 613268316
tensorflow · Mar 6, 2024 · 3adb838 · 3adb838
1 parent e3cf1f4
commit 3adb838
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tensorflow_probability/python/distributions/hypothesis_testlib.py b/tensorflow_probability/python/distributions/hypothesis_testlib.py
@@ -274,7 +274,7 @@ def fix_bates(d):
     'tailweight': tfp_hps.softplus_plus_eps(),
     'temperature': tfp_hps.softplus_plus_eps(),
     'total_count': lambda x: tf.floor(tf.sigmoid(x / 100) * 100) + 1,
-    'total_shape': tfp_hps.shapes(min_ndims=1, min_lastdimsize=2),
+    'concentration_shape': tfp_hps.shapes(min_ndims=1, min_lastdimsize=2),
     'Bates': fix_bates,
     'Bernoulli': lambda d: dict(d, dtype=tf.float32),
     'CholeskyLKJ': fix_lkj,