code: def newton(
      backward_differences,
      max_num_iters,
      newton_coefficient,
      ode_fn_vec,
      order,
      step_size,
      time,
      tol,
      unitary,
      upper,
  ):
      """Runs Newton's method to solve the BDF equation."""
      initial_guess = tf.reduce_sum(
          tf1.where(
              tf.range(MAX_ORDER + 1) <= order,
              backward_differences[: MAX_ORDER + 1],
              tf.zeros_like(backward_differences)[: MAX_ORDER + 1],
          ),
          axis=0,
      )
  
      np_dtype = np_dtype = dtype_util.as_numpy_dtype(
          backward_differences.dtype
      )
  
      rhs_constant_term = newton_coefficient * tf.reduce_sum(
          tf1.where(
              tf.range(1, MAX_ORDER + 1) <= order,
              RECIPROCAL_SUMS[1:, np.newaxis].astype(np_dtype)
              * backward_differences[1 : MAX_ORDER + 1],
              tf.zeros_like(backward_differences)[1 : MAX_ORDER + 1],
          ),
          axis=0,
      )
  
      next_time = time + step_size
      step_size_cast = tf.cast(step_size, backward_differences.dtype)
      real_dtype = tf.abs(backward_differences).dtype
  
      def newton_body(iterand):
          """Performs one iteration of Newton's method."""
          next_backward_difference = iterand.next_backward_difference
          next_state_vec = iterand.next_state_vec
  
          rhs = (
              newton_coefficient
              * step_size_cast
              * ode_fn_vec(next_time, next_state_vec)
              - rhs_constant_term
              - next_backward_difference
          )
          delta = tf.squeeze(
              tf.linalg.triangular_solve(
                  upper,
                  tf.matmul(tf.transpose(unitary), rhs[:, tf.newaxis]),
                  lower=False,
              )
          )
          num_iters = iterand.num_iters + 1
  
          next_backward_difference += delta
          next_state_vec += delta
  
          delta_norm = tf.cast(tf.norm(delta), real_dtype)
          lipschitz_const = delta_norm / iterand.prev_delta_norm
  
          # Stop if method has converged.
          approx_dist_to_sol = (
              lipschitz_const / (1.0 - lipschitz_const) * delta_norm
          )
          close_to_sol = approx_dist_to_sol < tol
          delta_norm_is_zero = tf.equal(
              delta_norm, tf.constant(0.0, dtype=real_dtype)
          )
          converged = close_to_sol | delta_norm_is_zero
          finished = converged
  
          # Stop if any of the following conditions are met:
          # (A) We have hit the maximum number of iterations.
          # (B) The method is converging too slowly.
          # (C) The method is not expected to converge.
          too_slow = lipschitz_const > 1.0
          finished = finished | too_slow
          if max_num_iters is not None:
              too_many_iters = tf.equal(num_iters, max_num_iters)
              num_iters_left = max_num_iters - num_iters
              num_iters_left_cast = tf.cast(num_iters_left, real_dtype)
              wont_converge = (
                  approx_dist_to_sol
                  * lipschitz_const**num_iters_left_cast
                  > tol
              )
              finished = finished | too_many_iters | wont_converge
  
          return [
              _NewtonIterand(
                  converged=converged,
                  finished=finished,
                  next_backward_difference=next_backward_difference,
                  next_state_vec=next_state_vec,
                  num_iters=num_iters,
                  prev_delta_norm=delta_norm,
              )
          ]
  
      iterand = _NewtonIterand(
          converged=False,
          finished=False,
          next_backward_difference=tf.zeros_like(initial_guess),
          next_state_vec=tf.identity(initial_guess),
          num_iters=0,
          prev_delta_norm=tf.constant(np.array(-0.0), dtype=real_dtype),
      )
      [iterand] = tf.while_loop(
          lambda iterand: tf.logical_not(iterand.finished),
          newton_body,
          [iterand],
      )
      return (
          iterand.converged,
          iterand.next_backward_difference,
          iterand.next_state_vec,
          iterand.num_iters,
      )
  
file_path: probability/tensorflow_probability/python/math/ode/bdf_util.py
imports: import collections
  import numpy
  import tensorflow.compat.v1
  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 5 for calling a global function with example usage: newton(backward_differences, max_num_iters, newton_coefficient, ode_fn_vec, order, step_size, time, tol, unitary, upper) and returns: iterand,  iterand,  iterand,  iterand, list