code: def test_expected_grads_gives_gradients(self):
      priors = {
          "deep_factorized": deep_factorized.NoisyDeepFactorized(
              batch_shape=(16,)
          ),
          "normal": uniform_noise.NoisyNormal(
              loc=tf.range(16.0), scale=1.0
          ),
      }
      prior = priors["deep_factorized"]
      em = universal.UniversalBatchedEntropyModel(
          prior, coding_rank=2, compression=True, expected_grads=True
      )
      self.assertTrue(em._expected_grads)
      num_symbols = 1000
      # Source distribution is fixed as gaussian.
      source = priors["normal"].base
      x = source.sample((3, num_symbols), seed=0)
      with tf.GradientTape(persistent=True) as g:
          g.watch(x)
          x2, bits = em(x, training=True)
      self.assertIsInstance(g.gradient(x2, x), tf.Tensor)
      self.assertIsInstance(g.gradient(bits, x), tf.Tensor)
      for variable in em.trainable_variables:
          self.assertIsInstance(g.gradient(bits, variable), tf.Tensor)
  
imports: from absl.testing import parameterized
  import scipy.stats
  import tensorflow
  import tensorflow_probability
  from tensorflow_compression.python.distributions import deep_factorized
  from tensorflow_compression.python.distributions import uniform_noise
  from tensorflow_compression.python.entropy_models import universal
label_desc: Use image node_id 6 for calling the UniversalBatchedEntropyModelTest obj's underlying member method code with example usage: obj.test_expected_grads_gives_gradients() without return types