class_docstr: Ensemble of correlated Hull-White Models.
  
  Represents the Ito process:
  
  ```None
    dr_i(t) = (theta_i(t) - a_i(t) * r_i(t)) dt + sigma_i(t) * dW_{r_i}(t),
    1 <= i <= n,
  ```
  where `W_{r_i}` are 1D Brownian motions with a correlation matrix `Rho(t)`.
  For each `i`, `r_i` is the Hull-White process.
  `theta_i`, `a_i`, `sigma_i`, `Rho` are positive functions of time.
  `a_i` correspond to the mean-reversion rate, `sigma_i` is the volatility of
  the process, `theta_i(t)` is the function that determines long run behaviour
  of the process `r(t) = (r_1(t), ..., r_n(t))`
  and is computed to match the initial (at t=0) discount curve:
  
  ```None
  \theta_i = df_i(t) / dt + a_i * f_i(t) + 0.5 * sigma_i**2 / a_i
             * (1 - exp(-2 * a_i *t)), 1 <= i <= n
  ```
  where `f_i(t)` is the instantaneous forward rate at time `0` for a maturity
  `t` and `df_i(t)/dt` is the gradient of `f_i` with respect to the maturity.
  See Section 3.3.1 of [1] for details.
  
  The price at time `t` of a zero-coupon bond maturing at `T` is given by
  (Ref. [2]):
  
  ```None
  P(t,T) = P(0,T) / P(0,t) *
           exp(-(r(t) - f(0,t)) * G(t,T) - 0.5 * y(t) * G(t,T)^2)
  
  y(t) = int_0^t [exp(-2 int_u^t (a(s) ds)) sigma(u)^2 du]
  
  G(t,T) = int_t^T [exp(-int_t^u a(s) ds) du]
  ```
  
  If mean-reversion, `a_i`, is constant and the volatility (`sigma_i`), and
  correlation (`Rho`) are piecewise constant functions, the process is sampled
  exactly. Otherwise, Euler sampling is used.
  
  For `n=1` this class represents Hull-White Model (see
  tff.models.hull_white.HullWhiteModel1F).
  
  #### Example. Two correlated Hull-White processes.
  
  ```python
  import numpy as np
  import tensorflow as tf
  import tf_quant_finance as tff
  
  dtype = tf.float64
  # Mean-reversion is constant for the two processes. `mean_reversion(t)`
  # has shape `[dim] + t.shape`.
  mean_reversion = [0.03, 0.02]
  # Volatility is a piecewise constant function with jumps at the same locations
  # for both Hull-White processes. `volatility(t)` has shape `[dim] + t.shape`.
  volatility = tff.math.piecewise.PiecewiseConstantFunc(
      jump_locations=[[0.1, 2.], [0.1, 2.]],
      values=[[0.01, 0.02, 0.01], [0.01, 0.015, 0.01]],
      dtype=dtype)
  # Correlation matrix is constant
  corr_matrix = [[1., 0.1], [0.1, 1.]]
  initial_discount_rate_fn = lambda *args: [0.01, 0.015]
  process = tff.models.hull_white.VectorHullWhiteModel(
      dim=2, mean_reversion=mean_reversion,
      volatility=volatility,
      initial_discount_rate_fn=initial_discount_rate_fn,
      corr_matrix=corr_matrix,
      dtype=dtype)
  # Sample 10000 paths using Sobol numbers as a random type.
  times = np.linspace(0., 1.0, 10)
  num_samples = 10000  # number of trajectories
  paths = process.sample_paths(
      times,
      num_samples=num_samples,
      random_type=tff.math.random.RandomType.SOBOL,
      seed=[4, 2])
  # Compute mean for each Hull-White process at the terminal value
  tf.math.reduce_mean(paths[:, -1, :], axis=0)
  # Expected: [0.01070005, 0.01540942]
  
  # User can supply time discretization grid along with the underlying
  # normal draws directly to the sampler.
  times_grid = tf.sort(tf.concat([np.linspace(0., 1.0, 10), [0.1]], axis=-1))
  normal_draws = tf.random.stateless_normal(
      [num_samples // 2, 10, 2], seed=[4, 2], dtype=dtype)
  normal_draws = tf.concat([normal_draws, -normal_draws], axis=0)
  paths = process.sample_paths(
      times,
      num_samples=num_samples,
      times_grid = times_grid,
      normal_draws=normal_draws)
  # Compute mean for each Hull-White process at the terminal value
  tf.math.reduce_mean(paths[:, -1, :], axis=0)
  # Expected: [0.0101668 , 0.01509873]
  ```
  
  #### References:
    [1]: D. Brigo, F. Mercurio. Interest Rate Models. 2007.
    [2]: Leif B. G. Andersen and Vladimir V. Piterbarg. Interest Rate Modeling.
    Volume II: Term Structure Models.
class_docstr_tok: ["Ensemble","of","correlated","Hull-White","Models",".","Represents","the","Ito","process",":","``","`","None","dr_i","(","t",")","=","(","theta_i","(","t",")","-","a_i","(","t",")","*","r_i","(","t",")",")","dt","+","sigma_i","(","t",")","*","dW_","{","r_i","}","(","t",")",",","1","<","=","i","<","=","n",",","``","`","where","`","W_","{","r_i","}","`","are","1D","Brownian","motions","with","a","correlation","matrix","`","Rho","(","t",")","`",".","For","each","`","i","`",",","`","r_i","`","is","the","Hull-White","process",".","`","theta_i","`",",","`","a_i","`",",","`","sigma_i","`",",","`","Rho","`","are","positive","functions","of","time",".","`","a_i","`","correspond","to","the","mean-reversion","rate",",","`","sigma_i","`","is","the","volatility","of","the","process",",","`","theta_i","(","t",")","`","is","the","function","that","determines","long","run","behaviour","of","the","process","`","r","(","t",")","=","(","r_1","(","t",")",",","...",",","r_n","(","t",")",")","`","and","is","computed","to","match","the","initial","(","at","t=0",")","discount","curve",":","``","`","None","\\theta_i","=","df_i","(","t",")","\/","dt","+","a_i","*","f_i","(","t",")","+","0.5","*","sigma_i","*","*","2","\/","a_i","*","(","1","-","exp","(","-2","*","a_i","*","t",")",")",",","1","<","=","i","<","=","n","``","`","where","`","f_i","(","t",")","`","is","the","instantaneous","forward","rate","at","time","`","0","`","for","a","maturity","`","t","`","and","`","df_i","(","t",")","\/dt","`","is","the","gradient","of","`","f_i","`","with","respect","to","the","maturity",".","See","Section","3.3.1","of","[","1","]","for","details",".","The","price","at","time","`","t","`","of","a","zero-coupon","bond","maturing","at","`","T","`","is","given","by","(","Ref",".","[","2","]",")",":","``","`","None","P","(","t",",","T",")","=","P","(","0",",","T",")","\/","P","(","0",",","t",")","*","exp","(","-","(","r","(","t",")","-","f","(","0",",","t",")",")","*","G","(","t",",","T",")","-","0.5","*","y","(","t",")","*","G","(","t",",","T",")","^2",")","y","(","t",")","=","int_0^t","[","exp","(","-2","int_u^t","(","a","(","s",")","ds",")",")","sigma","(","u",")","^2","du","]","G","(","t",",","T",")","=","int_t^T","[","exp","(","-int_t^u","a","(","s",")","ds",")","du","]","``","`","If","mean-reversion",",","`","a_i","`",",","is","constant","and","the","volatility","(","`","sigma_i","`",")",",","and","correlation","(","`","Rho","`",")","are","piecewise","constant","functions",",","the","process","is","sampled","exactly",".","Otherwise",",","Euler","sampling","is","used",".","For","`","n=1","`","this","class","represents","Hull-White","Model","(","see","tff.models.hull_white.HullWhiteModel1F",")",".","#","#","#","#","Example",".","Two","correlated","Hull-White","processes",".","``","`","python","import","numpy","as","np","import","tensorflow","as","tf","import","tf_quant_finance","as","tff","dtype","=","tf.float64","#","Mean-reversion","is","constant","for","the","two","processes",".","`","mean_reversion","(","t",")","`","#","has","shape","`","[","dim","]","+","t.shape","`",".","mean_reversion","=","[","0.03",",","0.02","]","#","Volatility","is","a","piecewise","constant","function","with","jumps","at","the","same","locations","#","for","both","Hull-White","processes",".","`","volatility","(","t",")","`","has","shape","`","[","dim","]","+","t.shape","`",".","volatility","=","tff.math.piecewise.PiecewiseConstantFunc","(","jump_locations=","[","[","0.1",",","2",".","]",",","[","0.1",",","2",".","]","]",",","values=","[","[","0.01",",","0.02",",","0.01","]",",","[","0.01",",","0.015",",","0.01","]","]",",","dtype=dtype",")","#","Correlation","matrix","is","constant","corr_matrix","=","[","[","1.",",","0.1","]",",","[","0.1",",","1",".","]","]","initial_discount_rate_fn","=","lambda","*","args",":","[","0.01",",","0.015","]","process","=","tff.models.hull_white.VectorHullWhiteModel","(","dim=2",",","mean_reversion=mean_reversion",",","volatility=volatility",",","initial_discount_rate_fn=initial_discount_rate_fn",",","corr_matrix=corr_matrix",",","dtype=dtype",")","#","Sample","10000","paths","using","Sobol","numbers","as","a","random","type",".","times","=","np.linspace","(","0.",",","1.0",",","10",")","num_samples","=","10000","#","number","of","trajectories","paths","=","process.sample_paths","(","times",",","num_samples=num_samples",",","random_type=tff.math.random.RandomType.SOBOL",",","seed=","[","4",",","2","]",")","#","Compute","mean","for","each","Hull-White","process","at","the","terminal","value","tf.math.reduce_mean","(","paths","[",":",",","-1",",",":","]",",","axis=0",")","#","Expected",":","[","0.01070005",",","0.01540942","]","#","User","can","supply","time","discretization","grid","along","with","the","underlying","#","normal","draws","directly","to","the","sampler",".","times_grid","=","tf.sort","(","tf.concat","(","[","np.linspace","(","0.",",","1.0",",","10",")",",","[","0.1","]","]",",","axis=-1",")",")","normal_draws","=","tf.random.stateless_normal","(","[","num_samples","\/\/","2",",","10",",","2","]",",","seed=","[","4",",","2","]",",","dtype=dtype",")","normal_draws","=","tf.concat","(","[","normal_draws",",","-normal_draws","]",",","axis=0",")","paths","=","process.sample_paths","(","times",",","num_samples=num_samples",",","times_grid","=","times_grid",",","normal_draws=normal_draws",")","#","Compute","mean","for","each","Hull-White","process","at","the","terminal","value","tf.math.reduce_mean","(","paths","[",":",",","-1",",",":","]",",","axis=0",")","#","Expected",":","[","0.0101668",",","0.01509873","]","``","`","#","#","#","#","References",":","[","1","]",":","D.","Brigo",",","F.","Mercurio",".","Interest","Rate","Models",".","2007",".","[","2","]",":","Leif","B.","G.","Andersen","and","Vladimir","V.","Piterbarg",".","Interest","Rate","Modeling",".","Volume","II",":","Term","Structure","Models","."]
docstr: Initializes the Correlated Hull-White Model.
  
  Args:
    dim: A Python scalar which corresponds to the number of correlated
      Hull-White Models.
    mean_reversion: A real positive `Tensor` of shape `[dim]` or a Python
      callable. The callable can be one of the following:
        (a) A left-continuous piecewise constant object (e.g.,
        `tff.math.piecewise.PiecewiseConstantFunc`) that has a property
        `is_piecewise_constant` set to `True`. In this case the object
        should have a method `jump_locations(self)` that returns a
        `Tensor` of shape `[dim, num_jumps]`. `mean_reversion(t)` should
        return a `Tensor` of shape `[dim, num_points]` where `t` is either a
        rank 1 `Tensor` of shape [num_points] or a `Tensor` of shape
        `[dim, num_points]`. Here `num_points` is arbitrary number of points.
        See example in the class docstring.
       (b) A callable that accepts scalars (stands for time `t`) and returns a
       `Tensor` of shape `[dim]`.
      Corresponds to the mean reversion rate.
    volatility: A real positive `Tensor` of the same `dtype` as
      `mean_reversion` or a callable with the same specs as above.
      Corresponds to the lond run price variance.
    initial_discount_rate_fn: A Python callable that accepts expiry time as
      a real `Tensor` of the same `dtype` as `mean_reversion` and returns
      a `Tensor` of either shape `input_shape` or `input_shape + [dim]`.
      Corresponds to the zero coupon bond yield at the present time for the
      input expiry time.
    corr_matrix: A `Tensor` of shape `[dim, dim]` and the same `dtype` as
      `mean_reversion` or a Python callable. The callable can be one of
      the following:
        (a) A left-continuous piecewise constant object (e.g.,
        `tff.math.piecewise.PiecewiseConstantFunc`) that has a property
        `is_piecewise_constant` set to `True`. In this case the object
        should have a method `jump_locations(self)` that returns a
        `Tensor` of shape `[num_jumps]`. `corr_matrix(t)` should return a
        `Tensor` of shape `t.shape + [dim, dim]`, where `t` is a rank 1
        `Tensor` of the same `dtype` as the output.
       (b) A callable that accepts scalars (stands for time `t`) and returns a
       `Tensor` of shape `[dim, dim]`.
      Corresponds to the correlation matrix `Rho`.
    dtype: The default dtype to use when converting values to `Tensor`s.
      Default value: `None` which maps to `tf.float32`.
    name: Python string. The name to give to the ops created by this class.
      Default value: `None` which maps to the default name `hull_white_model`.
  
  Raises:
    ValueError:
      (a) If either `mean_reversion`, `volatility`, or `corr_matrix` is
        a piecewise constant function where `jump_locations` have batch shape
        of rank > 1.
      (b): If batch rank of the `jump_locations` is `[n]` with `n` different
        from `dim`.
docstr_tok: ["Initializes","the","Correlated","Hull-White","Model",".","Args",":","dim",":","A","Python","scalar","which","corresponds","to","the","number","of","correlated","Hull-White","Models",".","mean_reversion",":","A","real","positive","`","Tensor","`","of","shape","`","[","dim","]","`","or","a","Python","callable",".","The","callable","can","be","one","of","the","following",":","(","a",")","A","left-continuous","piecewise","constant","object","(","e.g.",",","`","tff.math.piecewise.PiecewiseConstantFunc","`",")","that","has","a","property","`","is_piecewise_constant","`","set","to","`","True","`",".","In","this","case","the","object","should","have","a","method","`","jump_locations","(","self",")","`","that","returns","a","`","Tensor","`","of","shape","`","[","dim",",","num_jumps","]","`",".","`","mean_reversion","(","t",")","`","should","return","a","`","Tensor","`","of","shape","`","[","dim",",","num_points","]","`","where","`","t","`","is","either","a","rank","1","`","Tensor","`","of","shape","[","num_points","]","or","a","`","Tensor","`","of","shape","`","[","dim",",","num_points","]","`",".","Here","`","num_points","`","is","arbitrary","number","of","points",".","See","example","in","the","class","docstring",".","(","b",")","A","callable","that","accepts","scalars","(","stands","for","time","`","t","`",")","and","returns","a","`","Tensor","`","of","shape","`","[","dim","]","`",".","Corresponds","to","the","mean","reversion","rate",".","volatility",":","A","real","positive","`","Tensor","`","of","the","same","`","dtype","`","as","`","mean_reversion","`","or","a","callable","with","the","same","specs","as","above",".","Corresponds","to","the","lond","run","price","variance",".","initial_discount_rate_fn",":","A","Python","callable","that","accepts","expiry","time","as","a","real","`","Tensor","`","of","the","same","`","dtype","`","as","`","mean_reversion","`","and","returns","a","`","Tensor","`","of","either","shape","`","input_shape","`","or","`","input_shape","+","[","dim","]","`",".","Corresponds","to","the","zero","coupon","bond","yield","at","the","present","time","for","the","input","expiry","time",".","corr_matrix",":","A","`","Tensor","`","of","shape","`","[","dim",",","dim","]","`","and","the","same","`","dtype","`","as","`","mean_reversion","`","or","a","Python","callable",".","The","callable","can","be","one","of","the","following",":","(","a",")","A","left-continuous","piecewise","constant","object","(","e.g.",",","`","tff.math.piecewise.PiecewiseConstantFunc","`",")","that","has","a","property","`","is_piecewise_constant","`","set","to","`","True","`",".","In","this","case","the","object","should","have","a","method","`","jump_locations","(","self",")","`","that","returns","a","`","Tensor","`","of","shape","`","[","num_jumps","]","`",".","`","corr_matrix","(","t",")","`","should","return","a","`","Tensor","`","of","shape","`","t.shape","+","[","dim",",","dim","]","`",",","where","`","t","`","is","a","rank","1","`","Tensor","`","of","the","same","`","dtype","`","as","the","output",".","(","b",")","A","callable","that","accepts","scalars","(","stands","for","time","`","t","`",")","and","returns","a","`","Tensor","`","of","shape","`","[","dim",",","dim","]","`",".","Corresponds","to","the","correlation","matrix","`","Rho","`",".","dtype",":","The","default","dtype","to","use","when","converting","values","to","`","Tensor","`","s",".","Default","value",":","`","None","`","which","maps","to","`","tf.float32","`",".","name",":","Python","string",".","The","name","to","give","to","the","ops","created","by","this","class",".","Default","value",":","`","None","`","which","maps","to","the","default","name","`","hull_white_model","`",".","Raises",":","ValueError",":","(","a",")","If","either","`","mean_reversion","`",",","`","volatility","`",",","or","`","corr_matrix","`","is","a","piecewise","constant","function","where","`","jump_locations","`","have","batch","shape","of","rank",">","1",".","(","b",")",":","If","batch","rank","of","the","`","jump_locations","`","is","`","[","n","]","`","with","`","n","`","different","from","`","dim","`","."]
code: def __init__(
      self,
      dim: int,
      mean_reversion: Union[
          types.RealTensor, Callable[..., types.RealTensor]
      ],
      volatility: Union[
          types.RealTensor, Callable[..., types.RealTensor]
      ],
      initial_discount_rate_fn: Callable[..., types.RealTensor],
      corr_matrix: Optional[types.RealTensor] = None,
      dtype: Optional[tf.DType] = None,
      name: Optional[str] = None,
  ):
      """Initializes the Correlated Hull-White Model.
  
      Args:
        dim: A Python scalar which corresponds to the number of correlated
          Hull-White Models.
        mean_reversion: A real positive `Tensor` of shape `[dim]` or a Python
          callable. The callable can be one of the following:
            (a) A left-continuous piecewise constant object (e.g.,
            `tff.math.piecewise.PiecewiseConstantFunc`) that has a property
            `is_piecewise_constant` set to `True`. In this case the object
            should have a method `jump_locations(self)` that returns a
            `Tensor` of shape `[dim, num_jumps]`. `mean_reversion(t)` should
            return a `Tensor` of shape `[dim, num_points]` where `t` is either a
            rank 1 `Tensor` of shape [num_points] or a `Tensor` of shape
            `[dim, num_points]`. Here `num_points` is arbitrary number of points.
            See example in the class docstring.
           (b) A callable that accepts scalars (stands for time `t`) and returns a
           `Tensor` of shape `[dim]`.
          Corresponds to the mean reversion rate.
        volatility: A real positive `Tensor` of the same `dtype` as
          `mean_reversion` or a callable with the same specs as above.
          Corresponds to the lond run price variance.
        initial_discount_rate_fn: A Python callable that accepts expiry time as
          a real `Tensor` of the same `dtype` as `mean_reversion` and returns
          a `Tensor` of either shape `input_shape` or `input_shape + [dim]`.
          Corresponds to the zero coupon bond yield at the present time for the
          input expiry time.
        corr_matrix: A `Tensor` of shape `[dim, dim]` and the same `dtype` as
          `mean_reversion` or a Python callable. The callable can be one of
          the following:
            (a) A left-continuous piecewise constant object (e.g.,
            `tff.math.piecewise.PiecewiseConstantFunc`) that has a property
            `is_piecewise_constant` set to `True`. In this case the object
            should have a method `jump_locations(self)` that returns a
            `Tensor` of shape `[num_jumps]`. `corr_matrix(t)` should return a
            `Tensor` of shape `t.shape + [dim, dim]`, where `t` is a rank 1
            `Tensor` of the same `dtype` as the output.
           (b) A callable that accepts scalars (stands for time `t`) and returns a
           `Tensor` of shape `[dim, dim]`.
          Corresponds to the correlation matrix `Rho`.
        dtype: The default dtype to use when converting values to `Tensor`s.
          Default value: `None` which maps to `tf.float32`.
        name: Python string. The name to give to the ops created by this class.
          Default value: `None` which maps to the default name `hull_white_model`.
  
      Raises:
        ValueError:
          (a) If either `mean_reversion`, `volatility`, or `corr_matrix` is
            a piecewise constant function where `jump_locations` have batch shape
            of rank > 1.
          (b): If batch rank of the `jump_locations` is `[n]` with `n` different
            from `dim`.
      """
      self._name = name or "hull_white_model"
      with tf.name_scope(self._name):
          self._dtype = dtype or tf.float32
          # If the parameter is callable but not a piecewise constant use
          # generic sampling method (e.g., Euler).
          self._sample_with_generic = False
          # A flag to identify whether all parameters are piecewise constants
          self._is_piecewise_constant = True
  
          def _instant_forward_rate_fn(t):
              t = tf.convert_to_tensor(t, dtype=self._dtype)
  
              def _log_zero_coupon_bond(x):
                  r = tf.convert_to_tensor(
                      initial_discount_rate_fn(x), dtype=self._dtype
                  )
                  # If initial_discount_rate_fn returns a Tensor of the same input,
                  # expand the output dimension to `input_shape + [1]`. Otherwise,
                  # it is expected that `r` is of shape `input_shape + [dim]`.
                  if r.shape.rank == x.shape.rank:
                      r = tf.expand_dims(r, axis=-1)
                  # Shape `x.shape + [dim]`
                  return -r * tf.expand_dims(x, axis=-1)
  
              rate = -gradient.fwd_gradient(
                  _log_zero_coupon_bond,
                  t,
                  use_gradient_tape=True,
                  unconnected_gradients=tf.UnconnectedGradients.ZERO,
              )
              return rate
  
          def _initial_discount_rate_fn(t):
              r = tf.convert_to_tensor(
                  initial_discount_rate_fn(t), dtype=self._dtype
              )
              # If initial_discount_rate_fn returns a Tensor of the same input,
              # expand the output dimension to `input_shape + [1]`. Otherwise,
              # it is expected that `r` is of shape `input_shape + [dim]`.
              if r.shape.rank == t.shape.rank:
                  r = tf.expand_dims(r, axis=-1)
              return r
  
          self._instant_forward_rate_fn = _instant_forward_rate_fn
          self._initial_discount_rate_fn = _initial_discount_rate_fn
          (
              self._mean_reversion,
              sample_with_generic,
              is_piecewise_constant,
          ) = _input_type(
              mean_reversion,
              dim=dim,
              dtype=dtype,
              name="mean_reversion",
          )
  
          # Update flag to whether to sample with a generic sampler.
          self._sample_with_generic |= sample_with_generic
          # Update flag to whether all parameters are piecewise constant.
          self._is_piecewise_constant &= is_piecewise_constant
          # Get the volatility type
          (
              self._volatility,
              sample_with_generic,
              is_piecewise_constant,
          ) = _input_type(
              volatility, dim=dim, dtype=dtype, name="volatility"
          )
  
          # Update flag to whether to sample with a generic sampler.
          self._sample_with_generic |= sample_with_generic
          # Update flag to whether all parameters are piecewise constant.
          self._is_piecewise_constant &= is_piecewise_constant
          if corr_matrix is not None:
              # Get correlation matrix type
              (
                  self._corr_matrix,
                  sample_with_generic,
                  is_piecewise_constant,
              ) = _input_type(
                  corr_matrix, dim=dim, dtype=dtype, name="corr_matrix"
              )
              # Update flag to whether to sample with a generic sampler.
              self._sample_with_generic |= sample_with_generic
              # Update flag to whether all parameters are piecewise constant.
              self._is_piecewise_constant &= is_piecewise_constant
          else:
              self._corr_matrix = None
  
          if self._is_piecewise_constant:
              self._exact_discretization_setup(dim)
  
      # Volatility function
      def _vol_fn(t, x):
          """Volatility function of correlated Hull-White."""
          # Get parameter values at time `t`
          volatility = _get_parameters(
              tf.expand_dims(t, -1), self._volatility
          )[0]
          volatility = tf.transpose(volatility)
          if self._corr_matrix is not None:
              corr_matrix = _get_parameters(
                  tf.expand_dims(t, -1), self._corr_matrix
              )
              corr_matrix = corr_matrix[0]
              corr_matrix = tf.linalg.cholesky(corr_matrix)
          else:
              corr_matrix = tf.eye(self._dim, dtype=volatility.dtype)
  
          return volatility * corr_matrix + tf.zeros(
              x.shape.as_list()[:-1] + [self._dim, self._dim],
              dtype=volatility.dtype,
          )
  
      # Drift function
      def _drift_fn(t, x):
          """Drift function of correlated Hull-White."""
          # Get parameter values at time `t`
          (
              mean_reversion,
              volatility,
          ) = _get_parameters(  # pylint: disable=unbalanced-tuple-unpacking
              tf.expand_dims(t, -1),
              self._mean_reversion,
              self._volatility,
          )
          fwd_rates = self._instant_forward_rate_fn(t)
          fwd_rates_grad = gradient.fwd_gradient(
              self._instant_forward_rate_fn,
              t,
              use_gradient_tape=True,
              unconnected_gradients=tf.UnconnectedGradients.ZERO,
          )
          drift = fwd_rates_grad + mean_reversion * fwd_rates
          drift += (
              volatility**2
              / 2
              / mean_reversion
              * (1 - tf.math.exp(-2 * mean_reversion * t))
              - mean_reversion * x
          )
          return drift
  
      super(VectorHullWhiteModel, self).__init__(
          dim, _drift_fn, _vol_fn, self._dtype, name
      )
  
imports: from typing import Callable, Union, Tuple, Optional
  import tensorflow.compat.v2
  from tf_quant_finance import types
  from tf_quant_finance import utils
  from tf_quant_finance.math import gradient
  from tf_quant_finance.math import piecewise
  from tf_quant_finance.math import random
  from tf_quant_finance.models import euler_sampling
  from tf_quant_finance.models import generic_ito_process
  from tf_quant_finance.models import utils
label_desc: Use image node_id 1 to create a new VectorHullWhiteModel object from inherited base classes: generic_ito_process with example: obj = VectorHullWhiteModel(dim, mean_reversion, volatility, initial_discount_rate_fn, corr_matrix, dtype, name)