code: def test_fitting_example(self):
      if not JAX_MODE:
          self.skipTest("Requires JAX with optax.")
      import optax  # pylint: disable=g-import-not-at-top
  
      build_trainable_normal_stateless = (
          trainable_state_util.as_stateless_builder(normal_generator)
      )
      init_fn, apply_fn = build_trainable_normal_stateless(shape=[])
  
      # Find the maximum likelihood distribution given observed data.
      x_observed = [3.0, -2.0, 1.7]
      mle_parameters, _ = minimize_stateless(
          loss_fn=lambda *params: -apply_fn(*params).log_prob(
              x_observed
          ),
          init=init_fn(
              seed=test_util.test_seed(sampler_type="stateless")
          ),
          optimizer=optax.adam(1.0),
          num_steps=400,
      )
      mle_dist = apply_fn(mle_parameters)
      self.assertAllClose(
          mle_dist.mean(), np.mean(x_observed), atol=0.1
      )
      self.assertAllClose(
          mle_dist.stddev(), np.std(x_observed), atol=0.1
      )
  
imports: import functools
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.bijectors import joint_map
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import deterministic
  from tensorflow_probability.python.distributions import joint_distribution_auto_batched
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import lognormal
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import prefer_static
  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.internal import trainable_state_util
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.math.minimize import minimize
  from tensorflow_probability.python.math.minimize import minimize_stateless
label_desc: Use image node_id 13 for calling the TestWrapGeneratorAsStateless obj's underlying member method code with example usage: obj.test_fitting_example() without return types