code: def test_basic_variational_fitting_stateless(self):
      if not JAX_MODE:
          self.skipTest("Uses `optax` for stateless optimization")
      import optax  # pylint: disable=g-import-not-at-top
  
      batch_shape = [2, 3]
      num_timesteps = 5
      num_inits = 10
      observed_time_series = self._build_tensor(
          np.random.randn(*(batch_shape + [num_timesteps]))
      )
  
      model = self._build_model(observed_time_series)
      seed = test_util.test_seed(sampler_type="stateless")
      init_seed, fit_seed = samplers.split_seed(seed, n=2)
  
      (
          init_fn,
          build_surrogate_fn,
      ) = fitting.build_factored_surrogate_posterior_stateless(
          model, batch_shape=num_inits
      )
      jd = model.joint_distribution(
          observed_time_series=observed_time_series
      )
      _, loss_curve = optimization.fit_surrogate_posterior_stateless(
          jd.log_prob,
          build_surrogate_posterior_fn=build_surrogate_fn,
          initial_parameters=init_fn(init_seed),
          sample_size=3,
          num_steps=10,
          optimizer=optax.adam(1e-1),
          jit_compile=True,
          seed=fit_seed,
      )
      self.assertLess(np.mean(loss_curve[-1]), np.mean(loss_curve[0]))
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.sts import fitting
  from tensorflow_probability.python.sts.components import local_linear_trend
  from tensorflow_probability.python.sts.components import seasonal
  from tensorflow_probability.python.sts.components import sum
  from tensorflow_probability.python.sts.internal import missing_values_util
  from tensorflow_probability.python.vi import csiszar_divergence
  from tensorflow_probability.python.vi import optimization
label_desc: Use image node_id 3 for calling the _VariationalInferenceTests obj's underlying member method code with example usage: obj.test_basic_variational_fitting_stateless() without return types