class_docstr: Multivariate Geometric Brownian Motion.
  
  Represents a d-dimensional Ito process:
  
  ```None
    dX_i(t) = means_i * X_i(t) * dt + volatilities_i * X_i(t) * dW_i(t),
    1 <= i <= d
  ```
  
  where `W(t) = (W_1(t), .., W_d(t))` is a d-dimensional Brownian motion with
  a correlation matrix `corr_matrix`, `means` and `volatilities` are `Tensor`s
  that correspond to mean and volatility of a Geometric Brownian Motion `X_i`
  
  ## Example
  
  ```python
  import tensorflow as tf
  import tf_quant_finance as tff
  corr_matrix = [[1, 0.1], [0.1, 1]]
  process = tff.models.MultivariateGeometricBrownianMotion(
      dim=2,
      means=1, volatilities=[0.1, 0.2],
      corr_matrix=corr_matrix,
      dtype=tf.float64)
  times = [0.1, 0.2, 1.0]
  initial_state=[1.0, 2.0]
  # Use SOBOL sequence to draw trajectories
  samples_sobol = process.sample_paths(
      times=times,
      initial_state=initial_state,
      random_type=tff.math.random.RandomType.SOBOL,
      num_samples=100000)
  
  # You can also supply the random normal draws directly to the sampler
  normal_draws = tf.random.stateless_normal(
      [100000, 3, 2], seed=[4, 2], dtype=tf.float64)
  samples_custom = process.sample_paths(
      times=times,
      initial_state=initial_state,
      normal_draws=normal_draws)
  ```
class_docstr_tok: ["Multivariate","Geometric","Brownian","Motion",".","Represents","a","d-dimensional","Ito","process",":","``","`","None","dX_i","(","t",")","=","means_i","*","X_i","(","t",")","*","dt","+","volatilities_i","*","X_i","(","t",")","*","dW_i","(","t",")",",","1","<","=","i","<","=","d","``","`","where","`","W","(","t",")","=","(","W_1","(","t",")",",","..",",","W_d","(","t",")",")","`","is","a","d-dimensional","Brownian","motion","with","a","correlation","matrix","`","corr_matrix","`",",","`","means","`","and","`","volatilities","`","are","`","Tensor","`","s","that","correspond","to","mean","and","volatility","of","a","Geometric","Brownian","Motion","`","X_i","`","#","#","Example","``","`","python","import","tensorflow","as","tf","import","tf_quant_finance","as","tff","corr_matrix","=","[","[","1",",","0.1","]",",","[","0.1",",","1","]","]","process","=","tff.models.MultivariateGeometricBrownianMotion","(","dim=2",",","means=1",",","volatilities=","[","0.1",",","0.2","]",",","corr_matrix=corr_matrix",",","dtype=tf.float64",")","times","=","[","0.1",",","0.2",",","1.0","]","initial_state=","[","1.0",",","2.0","]","#","Use","SOBOL","sequence","to","draw","trajectories","samples_sobol","=","process.sample_paths","(","times=times",",","initial_state=initial_state",",","random_type=tff.math.random.RandomType.SOBOL",",","num_samples=100000",")","#","You","can","also","supply","the","random","normal","draws","directly","to","the","sampler","normal_draws","=","tf.random.stateless_normal","(","[","100000",",","3",",","2","]",",","seed=","[","4",",","2","]",",","dtype=tf.float64",")","samples_custom","=","process.sample_paths","(","times=times",",","initial_state=initial_state",",","normal_draws=normal_draws",")","``","`"]
docstr: Returns a solver for the Fokker Plank equation of this process.
  
  The Fokker Plank equation (also known as the Kolmogorov Forward equation)
  associated to this Ito process is given by:
  
  ```None
    V_t + Sum[(mu_i(t, x) V)_i, 1<=i<=n]
      - (1/2) Sum[ (D_{ij} V)_{ij}, 1 <= i,j <= n] = 0
  ```
  
  with the initial value condition $$V(0, x) = u(x)$$.
  
  This method evolves a spatially discretized solution of the above PDE from
  time `t0` to time `t1 > t0` (i.e. forwards in time).
  The solution `V(t,x)` is assumed to be discretized on an `n`-dimensional
  rectangular grid. A rectangular grid, G, in n-dimensions may be described
  by specifying the coordinates of the points along each axis. For example,
  a 2 x 4 grid in two dimensions can be specified by taking the cartesian
  product of [1, 3] and [5, 6, 7, 8] to yield the grid points with
  coordinates: `[(1, 5), (1, 6), (1, 7), (1, 8), (3, 5) ... (3, 8)]`.
  
  Batching of solutions is supported. In this context, batching means
  the ability to represent and evolve multiple independent functions `V`
  (e.g. V1, V2 ...) simultaneously. A single discretized solution is specified
  by stating its values at each grid point. This can be represented as a
  `Tensor` of shape [d1, d2, ... dn] where di is the grid size along the `i`th
  axis. A batch of such solutions is represented by a `Tensor` of shape:
  [K, d1, d2, ... dn] where `K` is the batch size. This method only requires
  that the input parameter `values_grid` be broadcastable with shape
  [K, d1, ... dn].
  
  The evolution of the solution from `t0` to `t1` is often done by
  discretizing the differential equation to a difference equation along
  the spatial and temporal axes. The temporal discretization is given by a
  (sequence of) time steps [dt_1, dt_2, ... dt_k] such that the sum of the
  time steps is equal to the total time step `t1 - t0`. If a uniform time
  step is used, it may equivalently be specified by stating the number of
  steps (n_steps) to take. This method provides both options via the
  `time_step` and `num_steps` parameters. However, not all methods need
  discretization along time direction (e.g. method of lines) so this argument
  may not be applicable to some implementations.
  
  The workhorse of this method is the `one_step_fn`. For the commonly used
  methods, see functions in `math.pde.steppers` module.
  
  The mapping between the arguments of this method and the above
  equation are described in the Args section below.
  
  For a simple instructive example of implementation of this method, see
  `models.GenericItoProcess.fd_solver_forward`.
  
  # TODO(b/142309558): Complete documentation.
  
  Args:
    start_time: Real positive scalar `Tensor`. The start time of the grid.
      Corresponds to time `t0` above.
    end_time: Real scalar `Tensor` smaller than the `start_time` and greater
      than zero. The time to step back to. Corresponds to time `t1` above.
    coord_grid: List of `n` rank 1 real `Tensor`s. `n` is the dimension of the
      domain. The i-th `Tensor` has shape, `[d_i]` where `d_i` is the size of
      the grid along axis `i`. The coordinates of the grid points. Corresponds
      to the spatial grid `G` above.
    values_grid: Real `Tensor` containing the function values at time
      `start_time` which have to be stepped back to time `end_time`. The shape
      of the `Tensor` must broadcast with `[K, d_1, d_2, ..., d_n]`. The first
      axis of size `K` is the values batch dimension and allows multiple
      functions (with potentially different boundary/final conditions) to be
      stepped back simultaneously.
    one_step_fn: The transition kernel. A callable that consumes the following
      arguments by keyword:
        1. 'time': Current time
        2. 'next_time': The next time to step to. For the backwards in time
          evolution, this time will be smaller than the current time.
        3. 'coord_grid': The coordinate grid.
        4. 'values_grid': The values grid.
        5. 'quadratic_coeff': A callable returning the quadratic coefficients
          of the PDE (i.e. `(1/2)D_{ij}(t, x)` above). The callable accepts
          the time and  coordinate grid as keyword arguments and returns a
          `Tensor` with shape that broadcasts with `[dim, dim]`.
        6. 'linear_coeff': A callable returning the linear coefficients of the
          PDE (i.e. `mu_i(t, x)` above). Accepts time and coordinate grid as
          keyword arguments and returns a `Tensor` with shape that broadcasts
          with `[dim]`.
        7. 'constant_coeff': A callable returning the coefficient of the
          linear homogeneous term (i.e. `r(t,x)` above). Same spec as above.
          The `one_step_fn` callable returns a 2-tuple containing the next
          coordinate grid, next values grid.
    boundary_conditions: A list of size `dim` containing boundary conditions.
      The i'th element of the list is a 2-tuple containing the lower and upper
      boundary condition for the boundary along the i`th axis.
    start_step_count: Scalar integer `Tensor`. Initial value for the number of
      time steps performed.
      Default value: 0 (i.e. no previous steps performed).
    num_steps: Positive int scalar `Tensor`. The number of time steps to take
      when moving from `start_time` to `end_time`. Either this argument or the
      `time_step` argument must be supplied (but not both). If num steps is
      `k>=1`, uniform time steps of size `(t0 - t1)/k` are taken to evolve the
      solution from `t0` to `t1`. Corresponds to the `n_steps` parameter
      above.
    time_step: The time step to take. Either this argument or the `num_steps`
      argument must be supplied (but not both). The type of this argument may
      be one of the following (in order of generality): (a) None in which case
        `num_steps` must be supplied. (b) A positive real scalar `Tensor`. The
        maximum time step to take. If the value of this argument is `dt`, then
        the total number of steps taken is N = (t1 - t0) / dt rounded up to
        the nearest integer. The first N-1 steps are of size dt and the last
        step is of size `t1 - t0 - (N-1) * dt`. (c) A callable accepting the
        current time and returning the size of the step to take. The input and
        the output are real scalar `Tensor`s.
    values_transform_fn: An optional callable applied to transform the
      solution values at each time step. The callable is invoked after the
      time step has been performed. The callable should accept the time of the
      grid, the coordinate grid and the values grid and should return the
      values grid. All input arguments to be passed by keyword.
    dtype: The dtype to use.
    name: The name to give to the ops.
      Default value: None which means `solve_forward` is used.
    **kwargs: Additional keyword args:
      (1) pde_solver_fn: Function to solve the PDE that accepts all the above
        arguments by name and returns the same tuple object as required below.
        Defaults to `tff.math.pde.fd_solvers.solve_forward`.
  
  Returns:
    A tuple object containing at least the following attributes:
      final_values_grid: A `Tensor` of same shape and dtype as `values_grid`.
        Contains the final state of the values grid at time `end_time`.
      final_coord_grid: A list of `Tensor`s of the same specification as
        the input `coord_grid`. Final state of the coordinate grid at time
        `end_time`.
      step_count: The total step count (i.e. the sum of the `start_step_count`
        and the number of steps performed in this call.).
      final_time: The final time at which the evolution stopped. This value
        is given by `max(min(end_time, start_time), 0)`.
docstr_tok: ["Returns","a","solver","for","the","Fokker","Plank","equation","of","this","process",".","The","Fokker","Plank","equation","(","also","known","as","the","Kolmogorov","Forward","equation",")","associated","to","this","Ito","process","is","given","by",":","``","`","None","V_t","+","Sum","[","(","mu_i","(","t",",","x",")","V",")","_i",",","1","<","=i","<","=n","]","-","(","1\/2",")","Sum","[","(","D_","{","ij","}","V",")","_","{","ij","}",",","1","<","=","i",",","j","<","=","n","]","=","0","``","`","with","the","initial","value","condition","$","$","V","(","0",",","x",")","=","u","(","x",")","$","$",".","This","method","evolves","a","spatially","discretized","solution","of","the","above","PDE","from","time","`","t0","`","to","time","`","t1",">","t0","`","(","i.e",".","forwards","in","time",")",".","The","solution","`","V","(","t",",","x",")","`","is","assumed","to","be","discretized","on","an","`","n","`","-dimensional","rectangular","grid",".","A","rectangular","grid",",","G",",","in","n-dimensions","may","be","described","by","specifying","the","coordinates","of","the","points","along","each","axis",".","For","example",",","a","2","x","4","grid","in","two","dimensions","can","be","specified","by","taking","the","cartesian","product","of","[","1",",","3","]","and","[","5",",","6",",","7",",","8","]","to","yield","the","grid","points","with","coordinates",":","`","[","(","1",",","5",")",",","(","1",",","6",")",",","(","1",",","7",")",",","(","1",",","8",")",",","(","3",",","5",")","...","(","3",",","8",")","]","`",".","Batching","of","solutions","is","supported",".","In","this","context",",","batching","means","the","ability","to","represent","and","evolve","multiple","independent","functions","`","V","`","(","e.g",".","V1",",","V2","...",")","simultaneously",".","A","single","discretized","solution","is","specified","by","stating","its","values","at","each","grid","point",".","This","can","be","represented","as","a","`","Tensor","`","of","shape","[","d1",",","d2",",","...","dn","]","where","di","is","the","grid","size","along","the","`","i","`","th","axis",".","A","batch","of","such","solutions","is","represented","by","a","`","Tensor","`","of","shape",":","[","K",",","d1",",","d2",",","...","dn","]","where","`","K","`","is","the","batch","size",".","This","method","only","requires","that","the","input","parameter","`","values_grid","`","be","broadcastable","with","shape","[","K",",","d1",",","...","dn","]",".","The","evolution","of","the","solution","from","`","t0","`","to","`","t1","`","is","often","done","by","discretizing","the","differential","equation","to","a","difference","equation","along","the","spatial","and","temporal","axes",".","The","temporal","discretization","is","given","by","a","(","sequence","of",")","time","steps","[","dt_1",",","dt_2",",","...","dt_k","]","such","that","the","sum","of","the","time","steps","is","equal","to","the","total","time","step","`","t1","-","t0","`",".","If","a","uniform","time","step","is","used",",","it","may","equivalently","be","specified","by","stating","the","number","of","steps","(","n_steps",")","to","take",".","This","method","provides","both","options","via","the","`","time_step","`","and","`","num_steps","`","parameters",".","However",",","not","all","methods","need","discretization","along","time","direction","(","e.g",".","method","of","lines",")","so","this","argument","may","not","be","applicable","to","some","implementations",".","The","workhorse","of","this","method","is","the","`","one_step_fn","`",".","For","the","commonly","used","methods",",","see","functions","in","`","math.pde.steppers","`","module",".","The","mapping","between","the","arguments","of","this","method","and","the","above","equation","are","described","in","the","Args","section","below",".","For","a","simple","instructive","example","of","implementation","of","this","method",",","see","`","models.GenericItoProcess.fd_solver_forward","`",".","#","TODO","(","b\/142309558",")",":","Complete","documentation",".","Args",":","start_time",":","Real","positive","scalar","`","Tensor","`",".","The","start","time","of","the","grid",".","Corresponds","to","time","`","t0","`","above",".","end_time",":","Real","scalar","`","Tensor","`","smaller","than","the","`","start_time","`","and","greater","than","zero",".","The","time","to","step","back","to",".","Corresponds","to","time","`","t1","`","above",".","coord_grid",":","List","of","`","n","`","rank","1","real","`","Tensor","`","s",".","`","n","`","is","the","dimension","of","the","domain",".","The","i-th","`","Tensor","`","has","shape",",","`","[","d_i","]","`","where","`","d_i","`","is","the","size","of","the","grid","along","axis","`","i","`",".","The","coordinates","of","the","grid","points",".","Corresponds","to","the","spatial","grid","`","G","`","above",".","values_grid",":","Real","`","Tensor","`","containing","the","function","values","at","time","`","start_time","`","which","have","to","be","stepped","back","to","time","`","end_time","`",".","The","shape","of","the","`","Tensor","`","must","broadcast","with","`","[","K",",","d_1",",","d_2",",","...",",","d_n","]","`",".","The","first","axis","of","size","`","K","`","is","the","values","batch","dimension","and","allows","multiple","functions","(","with","potentially","different","boundary\/final","conditions",")","to","be","stepped","back","simultaneously",".","one_step_fn",":","The","transition","kernel",".","A","callable","that","consumes","the","following","arguments","by","keyword",":","1",".","'time","'",":","Current","time","2",".","'next_time","'",":","The","next","time","to","step","to",".","For","the","backwards","in","time","evolution",",","this","time","will","be","smaller","than","the","current","time",".","3",".","'coord_grid","'",":","The","coordinate","grid",".","4",".","'values_grid","'",":","The","values","grid",".","5",".","'quadratic_coeff","'",":","A","callable","returning","the","quadratic","coefficients","of","the","PDE","(","i.e",".","`","(","1\/2",")","D_","{","ij","}","(","t",",","x",")","`","above",")",".","The","callable","accepts","the","time","and","coordinate","grid","as","keyword","arguments","and","returns","a","`","Tensor","`","with","shape","that","broadcasts","with","`","[","dim",",","dim","]","`",".","6",".","'linear_coeff","'",":","A","callable","returning","the","linear","coefficients","of","the","PDE","(","i.e",".","`","mu_i","(","t",",","x",")","`","above",")",".","Accepts","time","and","coordinate","grid","as","keyword","arguments","and","returns","a","`","Tensor","`","with","shape","that","broadcasts","with","`","[","dim","]","`",".","7",".","'constant_coeff","'",":","A","callable","returning","the","coefficient","of","the","linear","homogeneous","term","(","i.e",".","`","r","(","t",",","x",")","`","above",")",".","Same","spec","as","above",".","The","`","one_step_fn","`","callable","returns","a","2-tuple","containing","the","next","coordinate","grid",",","next","values","grid",".","boundary_conditions",":","A","list","of","size","`","dim","`","containing","boundary","conditions",".","The","i'th","element","of","the","list","is","a","2-tuple","containing","the","lower","and","upper","boundary","condition","for","the","boundary","along","the","i","`","th","axis",".","start_step_count",":","Scalar","integer","`","Tensor","`",".","Initial","value","for","the","number","of","time","steps","performed",".","Default","value",":","0","(","i.e",".","no","previous","steps","performed",")",".","num_steps",":","Positive","int","scalar","`","Tensor","`",".","The","number","of","time","steps","to","take","when","moving","from","`","start_time","`","to","`","end_time","`",".","Either","this","argument","or","the","`","time_step","`","argument","must","be","supplied","(","but","not","both",")",".","If","num","steps","is","`","k",">","=1","`",",","uniform","time","steps","of","size","`","(","t0","-","t1",")","\/k","`","are","taken","to","evolve","the","solution","from","`","t0","`","to","`","t1","`",".","Corresponds","to","the","`","n_steps","`","parameter","above",".","time_step",":","The","time","step","to","take",".","Either","this","argument","or","the","`","num_steps","`","argument","must","be","supplied","(","but","not","both",")",".","The","type","of","this","argument","may","be","one","of","the","following","(","in","order","of","generality",")",":","(","a",")","None","in","which","case","`","num_steps","`","must","be","supplied",".","(","b",")","A","positive","real","scalar","`","Tensor","`",".","The","maximum","time","step","to","take",".","If","the","value","of","this","argument","is","`","dt","`",",","then","the","total","number","of","steps","taken","is","N","=","(","t1","-","t0",")","\/","dt","rounded","up","to","the","nearest","integer",".","The","first","N-1","steps","are","of","size","dt","and","the","last","step","is","of","size","`","t1","-","t0","-","(","N-1",")","*","dt","`",".","(","c",")","A","callable","accepting","the","current","time","and","returning","the","size","of","the","step","to","take",".","The","input","and","the","output","are","real","scalar","`","Tensor","`","s",".","values_transform_fn",":","An","optional","callable","applied","to","transform","the","solution","values","at","each","time","step",".","The","callable","is","invoked","after","the","time","step","has","been","performed",".","The","callable","should","accept","the","time","of","the","grid",",","the","coordinate","grid","and","the","values","grid","and","should","return","the","values","grid",".","All","input","arguments","to","be","passed","by","keyword",".","dtype",":","The","dtype","to","use",".","name",":","The","name","to","give","to","the","ops",".","Default","value",":","None","which","means","`","solve_forward","`","is","used",".","*","*","kwargs",":","Additional","keyword","args",":","(","1",")","pde_solver_fn",":","Function","to","solve","the","PDE","that","accepts","all","the","above","arguments","by","name","and","returns","the","same","tuple","object","as","required","below",".","Defaults","to","`","tff.math.pde.fd_solvers.solve_forward","`",".","Returns",":","A","tuple","object","containing","at","least","the","following","attributes",":","final_values_grid",":","A","`","Tensor","`","of","same","shape","and","dtype","as","`","values_grid","`",".","Contains","the","final","state","of","the","values","grid","at","time","`","end_time","`",".","final_coord_grid",":","A","list","of","`","Tensor","`","s","of","the","same","specification","as","the","input","`","coord_grid","`",".","Final","state","of","the","coordinate","grid","at","time","`","end_time","`",".","step_count",":","The","total","step","count","(","i.e",".","the","sum","of","the","`","start_step_count","`","and","the","number","of","steps","performed","in","this","call.",")",".","final_time",":","The","final","time","at","which","the","evolution","stopped",".","This","value","is","given","by","`","max","(","min","(","end_time",",","start_time",")",",","0",")","`","."]
code: def fd_solver_forward(
      self,
      start_time,
      end_time,
      coord_grid,
      values_grid,
      one_step_fn=None,
      boundary_conditions=None,
      start_step_count=0,
      num_steps=None,
      time_step=None,
      values_transform_fn=None,
      dtype=None,
      name=None,
      **kwargs,
  ):
      r"""Returns a solver for the Fokker Plank equation of this process.
  
      The Fokker Plank equation (also known as the Kolmogorov Forward equation)
      associated to this Ito process is given by:
  
      ```None
        V_t + Sum[(mu_i(t, x) V)_i, 1<=i<=n]
          - (1/2) Sum[ (D_{ij} V)_{ij}, 1 <= i,j <= n] = 0
      ```
  
      with the initial value condition $$V(0, x) = u(x)$$.
  
      This method evolves a spatially discretized solution of the above PDE from
      time `t0` to time `t1 > t0` (i.e. forwards in time).
      The solution `V(t,x)` is assumed to be discretized on an `n`-dimensional
      rectangular grid. A rectangular grid, G, in n-dimensions may be described
      by specifying the coordinates of the points along each axis. For example,
      a 2 x 4 grid in two dimensions can be specified by taking the cartesian
      product of [1, 3] and [5, 6, 7, 8] to yield the grid points with
      coordinates: `[(1, 5), (1, 6), (1, 7), (1, 8), (3, 5) ... (3, 8)]`.
  
      Batching of solutions is supported. In this context, batching means
      the ability to represent and evolve multiple independent functions `V`
      (e.g. V1, V2 ...) simultaneously. A single discretized solution is specified
      by stating its values at each grid point. This can be represented as a
      `Tensor` of shape [d1, d2, ... dn] where di is the grid size along the `i`th
      axis. A batch of such solutions is represented by a `Tensor` of shape:
      [K, d1, d2, ... dn] where `K` is the batch size. This method only requires
      that the input parameter `values_grid` be broadcastable with shape
      [K, d1, ... dn].
  
      The evolution of the solution from `t0` to `t1` is often done by
      discretizing the differential equation to a difference equation along
      the spatial and temporal axes. The temporal discretization is given by a
      (sequence of) time steps [dt_1, dt_2, ... dt_k] such that the sum of the
      time steps is equal to the total time step `t1 - t0`. If a uniform time
      step is used, it may equivalently be specified by stating the number of
      steps (n_steps) to take. This method provides both options via the
      `time_step` and `num_steps` parameters. However, not all methods need
      discretization along time direction (e.g. method of lines) so this argument
      may not be applicable to some implementations.
  
      The workhorse of this method is the `one_step_fn`. For the commonly used
      methods, see functions in `math.pde.steppers` module.
  
      The mapping between the arguments of this method and the above
      equation are described in the Args section below.
  
      For a simple instructive example of implementation of this method, see
      `models.GenericItoProcess.fd_solver_forward`.
  
      # TODO(b/142309558): Complete documentation.
  
      Args:
        start_time: Real positive scalar `Tensor`. The start time of the grid.
          Corresponds to time `t0` above.
        end_time: Real scalar `Tensor` smaller than the `start_time` and greater
          than zero. The time to step back to. Corresponds to time `t1` above.
        coord_grid: List of `n` rank 1 real `Tensor`s. `n` is the dimension of the
          domain. The i-th `Tensor` has shape, `[d_i]` where `d_i` is the size of
          the grid along axis `i`. The coordinates of the grid points. Corresponds
          to the spatial grid `G` above.
        values_grid: Real `Tensor` containing the function values at time
          `start_time` which have to be stepped back to time `end_time`. The shape
          of the `Tensor` must broadcast with `[K, d_1, d_2, ..., d_n]`. The first
          axis of size `K` is the values batch dimension and allows multiple
          functions (with potentially different boundary/final conditions) to be
          stepped back simultaneously.
        one_step_fn: The transition kernel. A callable that consumes the following
          arguments by keyword:
            1. 'time': Current time
            2. 'next_time': The next time to step to. For the backwards in time
              evolution, this time will be smaller than the current time.
            3. 'coord_grid': The coordinate grid.
            4. 'values_grid': The values grid.
            5. 'quadratic_coeff': A callable returning the quadratic coefficients
              of the PDE (i.e. `(1/2)D_{ij}(t, x)` above). The callable accepts
              the time and  coordinate grid as keyword arguments and returns a
              `Tensor` with shape that broadcasts with `[dim, dim]`.
            6. 'linear_coeff': A callable returning the linear coefficients of the
              PDE (i.e. `mu_i(t, x)` above). Accepts time and coordinate grid as
              keyword arguments and returns a `Tensor` with shape that broadcasts
              with `[dim]`.
            7. 'constant_coeff': A callable returning the coefficient of the
              linear homogeneous term (i.e. `r(t,x)` above). Same spec as above.
              The `one_step_fn` callable returns a 2-tuple containing the next
              coordinate grid, next values grid.
        boundary_conditions: A list of size `dim` containing boundary conditions.
          The i'th element of the list is a 2-tuple containing the lower and upper
          boundary condition for the boundary along the i`th axis.
        start_step_count: Scalar integer `Tensor`. Initial value for the number of
          time steps performed.
          Default value: 0 (i.e. no previous steps performed).
        num_steps: Positive int scalar `Tensor`. The number of time steps to take
          when moving from `start_time` to `end_time`. Either this argument or the
          `time_step` argument must be supplied (but not both). If num steps is
          `k>=1`, uniform time steps of size `(t0 - t1)/k` are taken to evolve the
          solution from `t0` to `t1`. Corresponds to the `n_steps` parameter
          above.
        time_step: The time step to take. Either this argument or the `num_steps`
          argument must be supplied (but not both). The type of this argument may
          be one of the following (in order of generality): (a) None in which case
            `num_steps` must be supplied. (b) A positive real scalar `Tensor`. The
            maximum time step to take. If the value of this argument is `dt`, then
            the total number of steps taken is N = (t1 - t0) / dt rounded up to
            the nearest integer. The first N-1 steps are of size dt and the last
            step is of size `t1 - t0 - (N-1) * dt`. (c) A callable accepting the
            current time and returning the size of the step to take. The input and
            the output are real scalar `Tensor`s.
        values_transform_fn: An optional callable applied to transform the
          solution values at each time step. The callable is invoked after the
          time step has been performed. The callable should accept the time of the
          grid, the coordinate grid and the values grid and should return the
          values grid. All input arguments to be passed by keyword.
        dtype: The dtype to use.
        name: The name to give to the ops.
          Default value: None which means `solve_forward` is used.
        **kwargs: Additional keyword args:
          (1) pde_solver_fn: Function to solve the PDE that accepts all the above
            arguments by name and returns the same tuple object as required below.
            Defaults to `tff.math.pde.fd_solvers.solve_forward`.
  
      Returns:
        A tuple object containing at least the following attributes:
          final_values_grid: A `Tensor` of same shape and dtype as `values_grid`.
            Contains the final state of the values grid at time `end_time`.
          final_coord_grid: A list of `Tensor`s of the same specification as
            the input `coord_grid`. Final state of the coordinate grid at time
            `end_time`.
          step_count: The total step count (i.e. the sum of the `start_step_count`
            and the number of steps performed in this call.).
          final_time: The final time at which the evolution stopped. This value
            is given by `max(min(end_time, start_time), 0)`.
      """
      pde_solver_fn = kwargs.get(
          "pde_solver_fn", fd_solvers.solve_forward
      )
  
      (
          backward_second_order,
          backward_first_order,
          backward_zeroth_order,
      ) = _backward_pde_coeffs(
          self._drift_fn, self._volatility_fn, discounting=None
      )
  
      # Transform backward to forward equation.
      inner_second_order_coeff_fn = lambda t, x: -backward_second_order(
          t, x
      )
      inner_first_order_coeff_fn = backward_first_order
      zeroth_order_coeff_fn = backward_zeroth_order
  
      return pde_solver_fn(
          start_time=start_time,
          end_time=end_time,
          coord_grid=coord_grid,
          values_grid=values_grid,
          num_steps=num_steps,
          start_step_count=start_step_count,
          time_step=time_step,
          one_step_fn=one_step_fn,
          boundary_conditions=boundary_conditions,
          values_transform_fn=values_transform_fn,
          inner_second_order_coeff_fn=inner_second_order_coeff_fn,
          inner_first_order_coeff_fn=inner_first_order_coeff_fn,
          zeroth_order_coeff_fn=zeroth_order_coeff_fn,
          dtype=dtype,
          name=name,
      )
  
file_path: tf-quant-finance/tf_quant_finance/models/geometric_brownian_motion/multivariate_geometric_brownian_motion.py
imports: import tensorflow.compat.v2
  from tf_quant_finance import utils
  from tf_quant_finance.math.pde import fd_solvers
  from tf_quant_finance.models import ito_process
  from tf_quant_finance.models import utils
label_desc: Use image node_id 10 for calling the MultivariateGeometricBrownianMotion obj's underlying member method code with example usage: obj.fd_solver_forward(start_time, end_time, coord_grid, values_grid, one_step_fn, boundary_conditions, start_step_count, num_steps, time_step, values_transform_fn, dtype, name) and returns: pde_solver_fn