code: def test_zinb_is_trainable(self):
      init_fn, apply_fn = trainable.make_trainable_stateless(
          inflated.ZeroInflatedNegativeBinomial,
          batch_and_event_shape=[5],
          parameter_dtype=tf.float32,
      )
      init_obj = init_fn(seed=test_util.test_seed())
      # ZeroInflatedNegativeBinomial should have three parameters per scalar
      # distribution:  two from NegativeBinomial, and one for the mixture weight.
      self.assertEqual(3, len(init_obj))
      init_dist = apply_fn(init_obj)
  
      lprob = init_dist.log_prob(
          np.array([0, 1, 2, 3, 4], dtype=np.float32)
      )
      self.assertEqual((5,), lprob.shape)
  
      samples = init_dist.sample(seed=test_util.test_seed())
      self.assertEqual((5,), samples.shape)
  
file_path: probability/tensorflow_probability/python/distributions/inflated_test.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import gamma
  from tensorflow_probability.python.distributions import inflated
  from tensorflow_probability.python.distributions import lognormal
  from tensorflow_probability.python.distributions import negative_binomial
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.experimental.util import trainable
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
label_desc: Use image node_id 6 for calling the DistributionsTest obj's underlying member method code with example usage: obj.test_zinb_is_trainable() without return types