code: def test_criteria_can_run_under_xla_with_static_shape(self):
      if not tf.config.experimental_functions_run_eagerly():
          self.skipTest(
              "XLA test does not make sense without tf.function"
          )
  
      init_x = np.array([0.0, 0.0]).astype(np.float32)
      target_x = np.array([3.0, 4.0]).astype(np.float32)
  
      x = tf.Variable(init_x)
      loss_fn = lambda: tf.reduce_sum((x - target_x) ** 2)
      opt = _get_adam_optimizer(0.1)
      num_steps = 100
  
      # This test verifies that it works to compile the entire optimization loop,
      # as opposed to the `jit_compile` argument to `minimize`, which only
      # compiles an optimization step.
      @tf.function(jit_compile=True)
      def do_minimization(return_full_length_trace):
          return minimize(
              loss_fn=loss_fn,
              num_steps=num_steps,
              optimizer=opt,
              trace_fn=lambda ms: (ms.loss, ms.has_converged),
              convergence_criterion=(
                  optimizer.convergence_criteria.LossNotDecreasing(
                      atol=0.1
                  )
              ),
              return_full_length_trace=return_full_length_trace,
          )
  
      trace = do_minimization(return_full_length_trace=True)
      self.evaluate(tf1.global_variables_initializer())
      losses, has_converged = self.evaluate(trace)
      self.assertEqual(num_steps, losses.shape[0])
      self.assertEqual(num_steps, has_converged.shape[0])
  
      # Verify that the test is interesting, i.e., that we actually converged
      # before the end.
      self.assertTrue(has_converged[-2])
  
      # Verify that the final loss is tiled up to the end of the array.
      converged_at_step = np.argmax(has_converged)
      self.assertTrue(
          np.all(
              losses[converged_at_step + 1 :]
              == losses[converged_at_step]
          )
      )
  
imports: import collections
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python import optimizer
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.internal import tf_keras
  from tensorflow_probability.python.math.minimize import minimize
  from tensorflow_probability.python.math.minimize import minimize_stateless
label_desc: Use image node_id 9 for calling the MinimizeTests obj's underlying member method code with example usage: obj.test_criteria_can_run_under_xla_with_static_shape() and returns: minimize