code: def _prepare_args_with_initial_vertex(
      objective_function,
      initial_vertex,
      step_sizes,
      objective_at_initial_vertex,
      batch_evaluate_objective,
  ):
      """Constructs a standard axes aligned simplex."""
      dim = ps.size(initial_vertex)
      # tf.eye complains about np.array(.., np.int32) num_rows, only welcomes numpy
      # scalars. TODO(b/162529062): Remove the following line.
      dim = dim if tf.is_tensor(dim) else int(dim)
      num_vertices = dim + 1
      unit_vectors_along_axes = tf.reshape(
          tf.eye(
              dim,
              dim,
              dtype=dtype_util.base_dtype(initial_vertex.dtype),
          ),
          ps.concat([[dim], ps.shape(initial_vertex)], axis=0),
      )
  
      # If step_sizes does not broadcast to initial_vertex, the multiplication
      # in the second term will fail.
      simplex_face = (
          initial_vertex + step_sizes * unit_vectors_along_axes
      )
      simplex = tf.concat(
          [tf.expand_dims(initial_vertex, axis=0), simplex_face], axis=0
      )
      # Evaluate the objective function at the simplex vertices.
      if objective_at_initial_vertex is None:
          (
              objective_at_simplex,
              num_evaluations,
          ) = _evaluate_objective_multiple(
              objective_function, simplex, batch_evaluate_objective
          )
      else:
          (
              objective_at_simplex_face,
              num_evaluations,
          ) = _evaluate_objective_multiple(
              objective_function, simplex_face, batch_evaluate_objective
          )
          objective_at_simplex = tf.concat(
              [
                  tf.expand_dims(objective_at_initial_vertex, axis=0),
                  objective_at_simplex_face,
              ],
              axis=0,
          )
  
      return (
          dim,
          num_vertices,
          simplex,
          objective_at_simplex,
          num_evaluations,
      )
  
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 13 for calling a global function with example usage: _prepare_args_with_initial_vertex(objective_function, initial_vertex, step_sizes, objective_at_initial_vertex, batch_evaluate_objective) and returns: dim,  num_vertices,  simplex,  objective_at_simplex,  num_evaluations