Fixes breakage involving passing a TF tensor to SciPy.

PiperOrigin-RevId: 502341769
Change-Id: I94439fc26f0567d3e700480f71f9a32dc2e8ed4c

Author: Johannes Ballé

tensorflow_compression/python/entropy_models/universal_test.py

@@ -371,7 +371,7 @@ def test_quantization_noise_is_uniform(self, scale):
     u = x - x_hat
     self.assertAllLessEqual(tf.abs(u), 0.5)
     # Check distribution has right statistics.
-    _, p = scipy.stats.kstest(u, "uniform", (-0.5, 1.0))
+    _, p = scipy.stats.kstest(tf.reshape(u, [-1]), "uniform", (-0.5, 1.0))
     self.assertGreater(p, 1e-6)
 
   def test_expected_grads_or_not_gives_same_bits(self):