code: def _test_log_prob_gradient_is_bounded(
      self, dist_cls, values, params=()
  ):
      x = tf.constant(values)
      with tf.GradientTape(persistent=True) as g:
          df = dist_cls(**dict(params))
          g.watch(x)
          p = df.log_prob(x)
          # When this bound triggers, we expect the gradient to be zero.
          idx = p < -32.0
          p = tf.maximum(p, -32.0)
      dx = g.gradient(p, x)
      self.assertAllEqual(dx[idx], np.zeros_like(dx[idx]))
      self.assertTrue(
          np.all(np.isfinite(dx)), f"dx has non-finite value: {dx}"
      )
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow
  import tensorflow_probability
  from tensorflow_compression.python.distributions import deep_factorized
  from tensorflow_compression.python.distributions import round_adapters
label_desc: Use image node_id 1 for calling a global function with example usage: _test_log_prob_gradient_is_bounded(dist_cls, values, params) without return types