code: def _solve(
      time_direction_fn,
      start_time,
      end_time,
      coord_grid,
      values_grid,
      num_steps=None,
      start_step_count=0,
      time_step=None,
      one_step_fn=None,
      boundary_conditions=None,
      values_transform_fn=None,
      second_order_coeff_fn=None,
      first_order_coeff_fn=None,
      zeroth_order_coeff_fn=None,
      inner_second_order_coeff_fn=None,
      inner_first_order_coeff_fn=None,
      maximum_steps=None,
      swap_memory=True,
      name=None,
  ):
      """Common code for solve_backward and solve_forward."""
      if (num_steps is None) == (time_step is None):
          raise ValueError(
              "Exactly one of num_steps or time_step"
              " should be supplied."
          )
      coord_grid = [
          tf.convert_to_tensor(dim_grid, dtype=values_grid.dtype)
          for dim_grid in coord_grid
      ]
      # Try broadcasting batch_shapes of coord_grid
      coord_grid = list(utils.broadcast_common_batch_shape(*coord_grid))
      n_dims = len(coord_grid)
      if one_step_fn is None:
          if n_dims == 1:
              one_step_fn = oscillation_damped_crank_nicolson_step()
          else:
              one_step_fn = douglas_adi_step(theta=0.5)
  
      if boundary_conditions is None:
  
          def zero_dirichlet(t, grid):
              del t, grid
              return 1, None, tf.constant(0, dtype=values_grid.dtype)
  
          boundary_conditions = [
              (zero_dirichlet, zero_dirichlet)
          ] * n_dims
  
      with tf.compat.v1.name_scope(
          name,
          default_name="solve",
          values=[
              start_time,
              end_time,
              coord_grid,
              values_grid,
              num_steps,
              time_step,
          ],
      ):
          time_step_fn, est_max_steps = _get_time_steps_info(
              start_time,
              end_time,
              num_steps,
              time_step,
              time_direction_fn,
          )
          if est_max_steps is None and maximum_steps is not None:
              est_max_steps = maximum_steps
  
          def loop_cond(
              should_stop, time, x_grid, f_grid, steps_performed
          ):
              del time, x_grid, f_grid, steps_performed
              return tf.logical_not(should_stop)
  
          def loop_body(
              should_stop, time, x_grid, f_grid, steps_performed
          ):
              """Propagates the grid in time."""
              del should_stop
              next_should_stop, t_next = time_step_fn(time)
              next_xs, next_fs = one_step_fn(
                  time=time,
                  next_time=t_next,
                  coord_grid=x_grid,
                  value_grid=f_grid,
                  boundary_conditions=boundary_conditions,
                  second_order_coeff_fn=second_order_coeff_fn,
                  first_order_coeff_fn=first_order_coeff_fn,
                  zeroth_order_coeff_fn=zeroth_order_coeff_fn,
                  inner_second_order_coeff_fn=inner_second_order_coeff_fn,
                  inner_first_order_coeff_fn=inner_first_order_coeff_fn,
                  num_steps_performed=steps_performed,
              )
  
              if values_transform_fn is not None:
                  next_xs, next_fs = values_transform_fn(
                      t_next, next_xs, next_fs
                  )
              return (
                  next_should_stop,
                  t_next,
                  next_xs,
                  next_fs,
                  steps_performed + 1,
              )
  
          # If the start time is already equal to end time, no stepping is needed.
          # solve_backward, solve_forward already took care of the case when end_time
          # is on the "wrong side" of start_time.
          should_already_stop = start_time == end_time
          initial_args = (
              should_already_stop,
              start_time,
              coord_grid,
              values_grid,
              start_step_count,
          )
          (
              _,
              final_time,
              final_coords,
              final_values,
              steps_performed,
          ) = tf.while_loop(
              loop_cond,
              loop_body,
              initial_args,
              swap_memory=swap_memory,
              maximum_iterations=est_max_steps,
          )
          return final_values, final_coords, final_time, steps_performed
  
file_path: tf-quant-finance/tf_quant_finance/math/pde/fd_solvers.py
imports: import tensorflow.compat.v2
  from tf_quant_finance import utils
  from tf_quant_finance.math.pde.steppers.douglas_adi import douglas_adi_step
  from tf_quant_finance.math.pde.steppers.oscillation_damped_crank_nicolson import oscillation_damped_crank_nicolson_step
label_desc: Use image node_id 3 for calling a global function with example usage: _solve(time_direction_fn, start_time, end_time, coord_grid, values_grid, num_steps, start_step_count, time_step, one_step_fn, boundary_conditions, values_transform_fn, second_order_coeff_fn, first_order_coeff_fn, zeroth_order_coeff_fn, inner_second_order_coeff_fn, inner_first_order_coeff_fn, maximum_steps, swap_memory, name) and returns: final_values,  final_coords,  final_time,  steps_performed, int,  None,  tf, tf, next_should_stop,  t_next,  next_xs,  next_fs,  unknown