code: def test_inhomogeneous_poisson_process_example(self):
      # Toy 1D data.
      index_points = (
          np.array([-10.0, -7.2, -4.0, -0.1, 0.1, 4.0, 6.2, 9.0])
          .reshape([-1, 1])
          .astype(np.float32)
      )
      observed_counts = np.array(
          [100, 90, 60, 13, 18, 37, 55, 42]
      ).astype(np.float32)
  
      # Trainable GP hyperparameters.
      kernel_log_amplitude = tf.Variable(
          0.0, name="kernel_log_amplitude"
      )
      kernel_log_lengthscale = tf.Variable(
          0.0, name="kernel_log_lengthscale"
      )
      observation_noise_log_scale = tf.Variable(
          0.0, name="observation_noise_log_scale"
      )
  
      # Generative model.
      def model_fn():
          kernel = exponentiated_quadratic.ExponentiatedQuadratic(
              amplitude=tf.exp(kernel_log_amplitude),
              length_scale=tf.exp(kernel_log_lengthscale),
          )
          latent_log_rates = yield jdc.JointDistributionCoroutine.Root(
              gaussian_process.GaussianProcess(
                  kernel,
                  index_points=index_points,
                  observation_noise_variance=tf.exp(
                      observation_noise_log_scale
                  ),
                  name="latent_log_rates",
              )
          )
          yield independent.Independent(
              poisson.Poisson(log_rate=latent_log_rates),
              reinterpreted_batch_ndims=1,
              name="y",
          )
  
      model = jdc.JointDistributionCoroutine(model_fn, name="model")
  
      # Variational model.
      logit_locs = tf.Variable(tf.zeros(observed_counts.shape))
      logit_softplus_scales = tf.Variable(
          tf.ones(observed_counts.shape) * -1
      )
  
      def variational_model_fn():
          _ = yield jdc.JointDistributionCoroutine.Root(
              independent.Independent(
                  normal.Normal(
                      loc=logit_locs,
                      scale=tf.nn.softplus(logit_softplus_scales),
                  ),
                  reinterpreted_batch_ndims=1,
              )
          )
          _ = yield deterministic.VectorDeterministic(observed_counts)
  
      q = jdc.JointDistributionCoroutine(
          variational_model_fn, name="variational_model"
      )
  
      losses, sample_path = optimization.fit_surrogate_posterior(
          target_log_prob_fn=lambda *args: model.log_prob(args),
          surrogate_posterior=q,
          optimizer=tf_keras.optimizers.Adam(learning_rate=0.1),
          num_steps=100,
          seed=test_util.test_seed(),
          sample_size=1,
          trace_fn=lambda t: (t.loss, q.sample(seed=42)[0]),
      )
  
      self.evaluate(tf1.global_variables_initializer())
      losses_, sample_path_ = self.evaluate((losses, sample_path))
      self.assertLess(losses_[-1], 80.0)  # Optimal loss is roughly 40.
      # Optimal latent logits are approximately the log observed counts.
      self.assertAllClose(
          sample_path_[-1], np.log(observed_counts), atol=1.0
      )
  
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import fill_scale_tril
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import deterministic
  from tensorflow_probability.python.distributions import gaussian_process
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_auto_batched
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import lognormal
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import poisson
  from tensorflow_probability.python.experimental.util import trainable
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.internal import tf_keras
  from tensorflow_probability.python.math.psd_kernels import exponentiated_quadratic
  from tensorflow_probability.python.util import deferred_tensor
  from tensorflow_probability.python.vi import optimization
label_desc: Use image node_id 4 for calling the OptimizationTests obj's underlying member method code with example usage: obj.test_inhomogeneous_poisson_process_example() without return types