code: def test_batch_inputs_produce_batch_results(self):
      batch_shape = [3]
      num_timesteps = 50
      # Build synthetic loss curves: a quadratic decrease, scaled by random
      # positive values.
      losses = tf.range(
          num_timesteps, 0, -1, dtype=tf.float32
      ) ** 2 * tf.exp(
          1e-2
          * tf.random.normal(
              batch_shape + [num_timesteps], seed=test_util.test_seed()
          )
      )
  
      # Test that a convergence criterion with batched params matches the
      # results of separately constructed criteria.
      window_sizes = [5, 6, 4]
      atols = [0.1, 15.0, 40.0]
      rtols = [0.4, 1e-2, 1e-6]
      min_num_steps = [6, num_timesteps - 1, 10]
      batch_criterion = loss_not_decreasing.LossNotDecreasing(
          window_size=window_sizes,
          atol=atols,
          rtol=rtols,
          min_num_steps=min_num_steps,
      )
  
      component_criteria = [
          loss_not_decreasing.LossNotDecreasing(
              window_size=ws, atol=at, rtol=rt, min_num_steps=mns
          )
          for (ws, at, rt, mns) in zip(
              window_sizes, atols, rtols, min_num_steps
          )
      ]
  
      # Initialize the auxiliary states for the batch and component criteria.
      batch_auxiliary_state = batch_criterion.bootstrap(
          losses[..., 0], [], []
      )
      component_auxiliary_state = [
          c.bootstrap(losses[i, 0], [], [])
          for (i, c) in enumerate(component_criteria)
      ]
  
      # Run all criteria forward and check that the results match.
      batch_has_converged = []
      batch_auxiliary_states = [batch_auxiliary_state]
      component_has_converged = []
      component_auxiliary_states = [component_auxiliary_state]
      for step in range(1, num_timesteps):
          has_converged, auxiliary_state = batch_criterion.one_step(
              step,
              losses[..., step],
              [],
              [],
              batch_auxiliary_states[-1],
          )
          batch_has_converged.append(has_converged)
          batch_auxiliary_states.append(auxiliary_state)
  
          has_converged, auxiliary_state = zip(
              *[
                  c.one_step(step, losses[i, step], [], [], cs)
                  for i, (c, cs) in enumerate(
                      zip(
                          component_criteria,
                          component_auxiliary_states[-1],
                      )
                  )
              ]
          )
          component_has_converged.append(has_converged)
          component_auxiliary_states.append(auxiliary_state)
  
      (
          batch_has_converged_,
          component_has_converged_,
          batch_auxiliary_states_,
          component_auxiliary_states_,
      ) = self.evaluate(
          [
              batch_has_converged,
              component_has_converged,
              batch_auxiliary_states,
              component_auxiliary_states,
          ]
      )
  
      for i in range(len(component_criteria)):
          self.assertAllEqual(
              [x[i] for x in component_has_converged_],
              [x[i] for x in batch_has_converged_],
          )
          self.assertAllEqual(
              [x[i] for x in component_auxiliary_states_],
              [
                  tf.nest.map_structure(lambda x: x[i], batch_state)
                  for batch_state in batch_auxiliary_states_
              ],
          )
  
          # Assert that we've set parameters so that the test is interesting.
          self.assertEqual(component_has_converged_[-1][i], True)
  
file_path: probability/tensorflow_probability/python/optimizer/convergence_criteria/loss_not_decreasing_test.py
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.optimizer.convergence_criteria import loss_not_decreasing
label_desc: Use image node_id 2 for calling the LossNotDecreasingTests obj's underlying member method code with example usage: obj.test_batch_inputs_produce_batch_results() without return types