code: def nelder_mead_one_step(
      current_simplex,
      current_objective_values,
      objective_function=None,
      dim=None,
      func_tolerance=None,
      position_tolerance=None,
      batch_evaluate_objective=False,
      reflection=None,
      expansion=None,
      contraction=None,
      shrinkage=None,
      name=None,
  ):
      """A single iteration of the Nelder Mead algorithm."""
      with tf.name_scope(name or "nelder_mead_one_step"):
          domain_dtype = dtype_util.base_dtype(current_simplex.dtype)
          order = tf.argsort(
              current_objective_values,
              direction="ASCENDING",
              stable=True,
          )
          (best_index, worst_index, second_worst_index) = (
              order[0],
              order[-1],
              order[-2],
          )
  
          worst_vertex = current_simplex[worst_index]
  
          (
              best_objective_value,
              worst_objective_value,
              second_worst_objective_value,
          ) = (
              current_objective_values[best_index],
              current_objective_values[worst_index],
              current_objective_values[second_worst_index],
          )
  
          # Compute the centroid of the face opposite the worst vertex.
          face_centroid = (
              tf.reduce_sum(current_simplex, axis=0) - worst_vertex
          )
          face_centroid /= tf.cast(dim, domain_dtype)
  
          # Reflect the worst vertex through the opposite face.
          reflected = face_centroid + reflection * (
              face_centroid - worst_vertex
          )
          objective_at_reflected = objective_function(reflected)
  
          num_evaluations = 1
          has_converged = _check_convergence(
              current_simplex,
              current_simplex[best_index],
              best_objective_value,
              worst_objective_value,
              func_tolerance,
              position_tolerance,
          )
  
          def _converged_fn():
              return (
                  True,
                  current_simplex,
                  current_objective_values,
                  np.int32(0),
              )
  
          case0 = has_converged, _converged_fn
          accept_reflected = (
              objective_at_reflected < second_worst_objective_value
          ) & (objective_at_reflected >= best_objective_value)
          accept_reflected_fn = _accept_reflected_fn(
              current_simplex,
              current_objective_values,
              worst_index,
              reflected,
              objective_at_reflected,
          )
          case1 = accept_reflected, accept_reflected_fn
          do_expansion = objective_at_reflected < best_objective_value
          expansion_fn = _expansion_fn(
              objective_function,
              current_simplex,
              current_objective_values,
              worst_index,
              reflected,
              objective_at_reflected,
              face_centroid,
              expansion,
          )
          case2 = do_expansion, expansion_fn
          do_outside_contraction = (
              objective_at_reflected < worst_objective_value
          ) & (objective_at_reflected >= second_worst_objective_value)
          outside_contraction_fn = _outside_contraction_fn(
              objective_function,
              current_simplex,
              current_objective_values,
              face_centroid,
              best_index,
              worst_index,
              reflected,
              objective_at_reflected,
              contraction,
              shrinkage,
              batch_evaluate_objective,
          )
          case3 = do_outside_contraction, outside_contraction_fn
          default_fn = _inside_contraction_fn(
              objective_function,
              current_simplex,
              current_objective_values,
              face_centroid,
              best_index,
              worst_index,
              worst_objective_value,
              contraction,
              shrinkage,
              batch_evaluate_objective,
          )
          (
              converged,
              next_simplex,
              next_objective_at_simplex,
              case_evals,
          ) = ps.case(
              [case0, case1, case2, case3],
              default=default_fn,
              exclusive=False,
          )
          tensorshape_util.set_shape(
              next_simplex, current_simplex.shape
          )
          tensorshape_util.set_shape(
              next_objective_at_simplex, current_objective_values.shape
          )
          return (
              converged,
              next_simplex,
              next_objective_at_simplex,
              num_evaluations + case_evals,
          )
  
imports: import collections
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 2 for calling a global function with example usage: nelder_mead_one_step(current_simplex, current_objective_values, objective_function, dim, func_tolerance, position_tolerance, batch_evaluate_objective, reflection, expansion, contraction, shrinkage, name) and returns: converged,  next_simplex,  next_objective_at_simplex,  unknown, True,  current_simplex,  current_objective_values,  np