code: def _prepare_brent_args(
      objective_fn,
      left_bracket,
      right_bracket,
      value_at_left_bracket,
      value_at_right_bracket,
      absolute_root_tolerance=2e-7,
      relative_root_tolerance=None,
      function_tolerance=2e-7,
      max_iterations=100,
      stopping_policy_fn=None,
  ):
      r"""Prepares arguments for root search using Brent's method.
  
      Args:
        objective_fn: Python callable for which roots are searched. It must be a
          callable of a single `Tensor` parameter and return a `Tensor` of the same
          shape and dtype as `left_bracket`.
        left_bracket: `Tensor` or Python float representing the first starting
          points. The function will search for roots between each pair of points
          defined by `left_bracket` and `right_bracket`. The shape of `left_bracket`
          should match that of the input to `objective_fn`.
        right_bracket: `Tensor` of the same shape and dtype as `left_bracket` or
          Python float representing the second starting points. The function will
          search for roots between each pair of points defined by `left_bracket` and
          `right_bracket`. This argument must have the same shape as `left_bracket`.
        value_at_left_bracket: Optional `Tensor` or Python float representing the
          value of `objective_fn` at `left_bracket`. If specified, this argument
          must have the same shape as `left_bracket`. If not specified, the value
          will be evaluated during the search.
          Default value: None.
        value_at_right_bracket: Optional `Tensor` or Python float representing the
          value of `objective_fn` at `right_bracket`. If specified, this argument
          must have the same shape as `right_bracket`. If not specified, the value
          will be evaluated during the search.
          Default value: None.
        absolute_root_tolerance: Optional `Tensor` representing the absolute
          tolerance for estimated roots, with the total tolerance being calculated
          as `(absolute_root_tolerance + relative_root_tolerance * |root|) / 2`. If
          specified, this argument must be positive, broadcast with the shape of
          `left_bracket` and have the same dtype.
          Default value: `2e-7`.
        relative_root_tolerance: Optional `Tensor` representing the relative
          tolerance for estimated roots, with the total tolerance being calculated
          as `(absolute_root_tolerance + relative_root_tolerance * |root|) / 2`. If
          specified, this argument must be positive, broadcast with the shape of
          `left_bracket` and have the same dtype.
          Default value: `None` which translates to `4 *
            numpy.finfo(left_bracket.dtype.as_numpy_dtype).eps`.
        function_tolerance: Optional `Tensor` representing the tolerance used to
          check for roots. If the absolute value of `objective_fn` is smaller than
          or equal to `function_tolerance` at a given estimate, then that estimate
          is considered a root for the function. If specified, this argument must
          broadcast with the shape of `left_bracket` and have the same dtype. Set to
          zero to match Brent's original algorithm and to continue the search until
          an exact root is found.
          Default value: `2e-7`.
        max_iterations: Optional `Tensor` of an integral dtype or Python integer
          specifying the maximum number of steps to perform for each initial point.
          Must broadcast with the shape of `left_bracket`. If an element is set to
          zero, the function will not search for any root for the corresponding
          points in `left_bracket` and `right_bracket`. Instead, it will return the
          best estimate from the inputs.
          Default value: `100`.
        stopping_policy_fn: Python `callable` controlling the algorithm termination.
          It must be a callable accepting a `Tensor` of booleans with the shape of
          `left_bracket` (each denoting whether the search is finished for each
          starting point), and returning a scalar boolean `Tensor` (indicating
          whether the overall search should stop). Typical values are
          `tf.reduce_all` (which returns only when the search is finished for all
          pairs of points), and `tf.reduce_any` (which returns as soon as the search
          is finished for any pair of points).
          Default value: `None` which translates to `tf.reduce_all`.
  
      Returns:
        A tuple of 3 Python objects containing the state, parameters, and constants
        to use for the search.
      """
      stopping_policy_fn = stopping_policy_fn or tf.reduce_all
      if not callable(stopping_policy_fn):
          raise ValueError("stopping_policy_fn must be callable")
  
      left_bracket = tf.convert_to_tensor(
          left_bracket, name="left_bracket"
      )
      right_bracket = tf.convert_to_tensor(
          right_bracket, name="right_bracket", dtype=left_bracket.dtype
      )
  
      if value_at_left_bracket is None:
          value_at_left_bracket = objective_fn(left_bracket)
      if value_at_right_bracket is None:
          value_at_right_bracket = objective_fn(right_bracket)
  
      value_at_left_bracket = tf.convert_to_tensor(
          value_at_left_bracket,
          name="value_at_left_bracket",
          dtype=left_bracket.dtype.base_dtype,
      )
      value_at_right_bracket = tf.convert_to_tensor(
          value_at_right_bracket,
          name="value_at_right_bracket",
          dtype=left_bracket.dtype.base_dtype,
      )
  
      if relative_root_tolerance is None:
          relative_root_tolerance = (
              utils.default_relative_root_tolerance(
                  left_bracket.dtype.base_dtype
              )
          )
  
      absolute_root_tolerance = tf.convert_to_tensor(
          absolute_root_tolerance,
          name="absolute_root_tolerance",
          dtype=left_bracket.dtype,
      )
      relative_root_tolerance = tf.convert_to_tensor(
          relative_root_tolerance,
          name="relative_root_tolerance",
          dtype=left_bracket.dtype,
      )
      function_tolerance = tf.convert_to_tensor(
          function_tolerance,
          name="function_tolerance",
          dtype=left_bracket.dtype,
      )
  
      max_iterations = tf.broadcast_to(
          tf.convert_to_tensor(max_iterations),
          name="max_iterations",
          shape=left_bracket.shape,
      )
      num_iterations = tf.zeros_like(max_iterations)
  
      false = tf.constant(False, shape=left_bracket.shape)
  
      zero = tf.zeros_like(left_bracket)
      contrapoint = zero
      step_to_last_estimate = zero
      step_to_best_estimate = zero
  
      zero_value = tf.zeros_like(value_at_left_bracket)
      value_at_contrapoint = zero_value
  
      # Select the best root estimates from the inputs.
      # If no search is performed (e.g. `max_iterations` is `zero`), the estimate
      # computed this way will be returned. This differs slightly from the SciPy
      # implementation which always returns the `right_bracket`.
      swap_positions = tf.math.abs(value_at_left_bracket) < tf.math.abs(
          value_at_right_bracket
      )
      best_estimate, last_estimate = _swap_where(
          swap_positions, right_bracket, left_bracket
      )
      value_at_best_estimate, value_at_last_estimate = _swap_where(
          swap_positions, value_at_right_bracket, value_at_left_bracket
      )
  
      # Check if the current root estimate is good enough.
      # Using zero as `function_tolerance` will check for exact roots and match both
      # Brent's original algorithm and the SciPy implementation.
      finished = (
          (num_iterations >= max_iterations)
          | (~tf.math.is_finite(value_at_last_estimate))
          | (~tf.math.is_finite(value_at_best_estimate))
          | (tf.math.abs(value_at_best_estimate) <= function_tolerance)
      )
  
      return (
          _BrentSearchState(
              best_estimate=best_estimate,
              last_estimate=last_estimate,
              contrapoint=contrapoint,
              value_at_best_estimate=value_at_best_estimate,
              value_at_last_estimate=value_at_last_estimate,
              value_at_contrapoint=value_at_contrapoint,
              step_to_best_estimate=step_to_best_estimate,
              step_to_last_estimate=step_to_last_estimate,
              num_iterations=num_iterations,
              finished=finished,
          ),
          _BrentSearchParams(
              objective_fn=objective_fn,
              max_iterations=max_iterations,
              absolute_root_tolerance=absolute_root_tolerance,
              relative_root_tolerance=relative_root_tolerance,
              function_tolerance=function_tolerance,
              stopping_policy_fn=stopping_policy_fn,
          ),
          _BrentSearchConstants(
              false=false, zero=zero, zero_value=zero_value
          ),
      )
  
file_path: tf-quant-finance/tf_quant_finance/math/root_search/brent.py
imports: from typing import Callable
  import tensorflow.compat.v2
  from tf_quant_finance import types
  from tf_quant_finance import utils
  from tf_quant_finance.math.root_search import utils
label_desc: Use image node_id 6 for calling a global function with example usage: _prepare_brent_args(objective_fn, left_bracket, right_bracket, value_at_left_bracket, value_at_right_bracket, absolute_root_tolerance, relative_root_tolerance, function_tolerance, max_iterations, stopping_policy_fn) and returns: _BrentSearchState,  _BrentSearchParams,  _BrentSearchConstants