code: def test_importance_sampling_example(self):
      if not JAX_MODE:
          self.skipTest("Requires optax.")
      import optax  # pylint: disable=g-import-not-at-top
  
      init_seed, opt_seed, eval_seed = samplers.split_seed(
          test_util.test_seed(sampler_type="stateless"), n=3
      )
  
      def log_prob(z, x):
          return normal.Normal(0.0, 1.0).log_prob(z) + normal.Normal(
              z, 1.0
          ).log_prob(x)
  
      conditioned_log_prob = lambda z: log_prob(z, x=5.0)
  
      init_normal, build_normal = trainable.make_trainable_stateless(
          normal.Normal
      )
      # Fit `q` with an importance-weighted variational loss.
      (
          optimized_parameters,
          _,
      ) = optimization.fit_surrogate_posterior_stateless(
          conditioned_log_prob,
          build_surrogate_posterior_fn=build_normal,
          initial_parameters=init_normal(seed=init_seed),
          importance_sample_size=10,
          optimizer=optax.adam(0.1),
          num_steps=200,
          seed=opt_seed,
      )
      q_z = build_normal(*optimized_parameters)
  
      # Estimate posterior statistics with importance sampling.
      zs, q_log_prob = self.evaluate(
          q_z.experimental_sample_and_log_prob(1000, seed=eval_seed)
      )
      self_normalized_log_weights = tf.nn.log_softmax(
          conditioned_log_prob(zs) - q_log_prob
      )
      posterior_mean = tf.reduce_sum(
          tf.exp(self_normalized_log_weights) * zs, axis=0
      )
      self.assertAllClose(posterior_mean, 2.5, atol=1e-1)
  
      posterior_variance = tf.reduce_sum(
          tf.exp(self_normalized_log_weights)
          * (zs - posterior_mean) ** 2,
          axis=0,
      )
      self.assertAllClose(posterior_variance, 0.5, atol=1e-1)
  
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 1 for calling the StatelessOptimizationTests obj's underlying member method code with example usage: obj.test_importance_sampling_example() and returns: unknown