Remove unused function

Add test cases

Author: nicolaspi

tensorflow_probability/python/stats/sample_stats.py

@@ -925,13 +925,6 @@ def _prepare_window_args(x, low_indices=None, high_indices=None, axis=0):
   return x, indices, axis
 
 
-def _safe_average(totals, counts):
-  # This tf.where protects `totals` from getting a gradient signal
-  # when `counts` is 0.
-  safe_totals = tf.where(~tf.equal(counts, 0), totals, 0)
-  return tf.where(~tf.equal(counts, 0), safe_totals / counts, 0)
-
-
 def log_average_probs(logits, sample_axis=0, event_axis=None, keepdims=False,
                       validate_args=False, name=None):
   """Computes `log(average(to_probs(logits)))` in a numerically stable manner.

tensorflow_probability/python/stats/sample_stats_test.py

@@ -731,13 +731,16 @@ def check_windowed(self, func, numpy_func):
     check_fn((64, 4, 8), (32, 1, 1), axis=0)
     check_fn((64, 4, 8), (32, 4, 1), axis=0)
     check_fn((64, 4, 8), (32, 4, 8), axis=0)
+    check_fn((64, 4, 8), (64, 4, 8), axis=0)
     check_fn((64, 4, 8), (64, 64, 1), axis=1)
     check_fn((64, 4, 8), (1, 64, 1), axis=1)
     check_fn((64, 4, 8), (64, 2, 8), axis=1)
+    check_fn((64, 4, 8), (64, 4, 8), axis=1)
     check_fn((64, 4, 8), (64, 4, 64), axis=2)
     check_fn((64, 4, 8), (1, 1, 64), axis=2)
     check_fn((64, 4, 8), (64, 4, 4), axis=2)
     check_fn((64, 4, 8), (1, 1, 4), axis=2)
+    check_fn((64, 4, 8), (64, 4, 8), axis=2)
 
     with self.assertRaises(Exception):
       # Non broadcastable shapes