code: def test_inhomogeneous_poisson_process_example(self):
      opt_seed, eval_seed = samplers.split_seed(
          test_util.test_seed(sampler_type="stateless"), n=2
      )
  
      # 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)
  
      # Generative model.
      def model_fn():
          kernel_amplitude = yield lognormal.LogNormal(
              loc=0.0, scale=1.0, name="kernel_amplitude"
          )
          kernel_lengthscale = yield lognormal.LogNormal(
              loc=0.0, scale=1.0, name="kernel_lengthscale"
          )
          observation_noise_scale = yield lognormal.LogNormal(
              loc=0.0, scale=1.0, name="observation_noise_scale"
          )
          kernel = exponentiated_quadratic.ExponentiatedQuadratic(
              amplitude=kernel_amplitude,
              length_scale=kernel_lengthscale,
          )
          latent_log_rates = yield gaussian_process.GaussianProcess(
              kernel,
              index_points=index_points,
              observation_noise_variance=observation_noise_scale,
              name="latent_log_rates",
          )
          yield independent.Independent(
              poisson.Poisson(log_rate=latent_log_rates),
              reinterpreted_batch_ndims=1,
              name="y",
          )
  
      model = jdab.JointDistributionCoroutineAutoBatched(model_fn)
      pinned = model.experimental_pin(y=observed_counts)
  
      initial_parameters = (
          0.0,
          0.0,
          0.0,  # Raw kernel parameters.
          tf.zeros_like(observed_counts),  # `logit_locs`
          tf.zeros_like(observed_counts),
      )  # `logit_raw_scales`
  
      def build_surrogate_posterior_fn(
          raw_kernel_amplitude,
          raw_kernel_lengthscale,
          raw_observation_noise_scale,
          logit_locs,
          logit_raw_scales,
      ):
          def variational_model_fn():
              # Fit the kernel parameters as point masses.
              yield deterministic.Deterministic(
                  tf.nn.softplus(raw_kernel_amplitude),
                  name="kernel_amplitude",
              )
              yield deterministic.Deterministic(
                  tf.nn.softplus(raw_kernel_lengthscale),
                  name="kernel_lengthscale",
              )
              yield deterministic.Deterministic(
                  tf.nn.softplus(raw_observation_noise_scale),
                  name="kernel_observation_noise_scale",
              )
              # Factored normal posterior over the GP logits.
              yield independent.Independent(
                  normal.Normal(
                      loc=logit_locs,
                      scale=tf.nn.softplus(logit_raw_scales),
                  ),
                  reinterpreted_batch_ndims=1,
                  name="latent_log_rates",
              )
  
          return jdab.JointDistributionCoroutineAutoBatched(
              variational_model_fn
          )
  
      if not JAX_MODE:
          return
      import optax  # pylint: disable=g-import-not-at-top
  
      def seeded_target_log_prob_fn(*xs, seed=None):
          # Add a tiny amount of noise to the target log-prob to see if it works.
          ret = pinned.unnormalized_log_prob(xs)
          return ret + samplers.normal(
              ret.shape, stddev=0.01, seed=seed
          )
  
      [
          optimized_parameters,
          (losses, _, sample_path),
      ] = optimization.fit_surrogate_posterior_stateless(
          target_log_prob_fn=seeded_target_log_prob_fn,
          build_surrogate_posterior_fn=build_surrogate_posterior_fn,
          initial_parameters=initial_parameters,
          optimizer=optax.adam(learning_rate=0.1),
          sample_size=1,
          num_steps=500,
          trace_fn=lambda traceable_quantities: (  # pylint: disable=g-long-lambda
              traceable_quantities.loss,
              tf.nn.softplus(traceable_quantities.parameters[0]),
              build_surrogate_posterior_fn(
                  *traceable_quantities.parameters
              ).sample(seed=eval_seed)[-1],
          ),
          seed=opt_seed,
      )
      surrogate_posterior = build_surrogate_posterior_fn(
          *optimized_parameters
      )
      surrogate_posterior.sample(seed=eval_seed)
  
      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 2 for calling the StatelessOptimizationTests obj's underlying member method code with example usage: obj.test_inhomogeneous_poisson_process_example() and returns: jdab, unknown