class_docstr: Quasi-Gaussian HJM model for term-structure modeling.
  
  Heath-Jarrow-Morton (HJM) model for the interest rate term-structre
  modelling specifies the dynamics of the instantaneus forward rate `f(t,T)`
  with maturity `T` at time `t` as follows:
  
  ```None
    df(t,T) = mu(t,T) dt + sum_i sigma_i(t,  T) * dW_i(t),
    1 <= i <= n,
  ```
  where `mu(t,T)` and `sigma_i(t,T)` denote the drift and volatility
  for the forward rate and `W_i` are Brownian motions with instantaneous
  correlation `Rho`. The model above represents an `n-factor` HJM model. Under
  the risk-neutral measure, the drift `mu(t,T)` is computed as
  
  ```
    mu(t,T) = sum_i sigma_i(t,T)  int_t^T sigma_(t,u) du
  ```
  Using the separability condition, the HJM model above can be formulated as
  the following Markovian model:
  
  ```None
    sigma(t,T) = sigma(t) * h(T)    (Separability condition)
  ```
  A common choice for the function h(t) is `h(t) = exp(-kt)`. Using the above
  parameterization of sigma(t,T), we obtain the following Markovian
  formulation of the HJM model [1]:
  
  ```None
    HJM Model
    dx_i(t) = (sum_j [y_ij(t)] - k_i * x_i(t)) dt + sigma_i(t) dW_i
    dy_ij(t) = (rho_ij * sigma_i(t)*sigma_j(t) - (k_i + k_j) * y_ij(t)) dt
    r(t) = sum_i x_i(t) + f(0, t)
  ```
  where `x` is an `n`-dimensional vector and `y` is an `nxn` dimensional
  matrix. The HJM class implements the model outlined above by jointly
  simulating the state [x_t, y_t].
  
  The price at time `t` of a zero-coupon bond maturing at `T` is given by
  (Ref. [1]):
  
  ```None
  P(t,T) = P(0,T) / P(0,t) *
           exp(-x(t) * G(t,T) - 0.5 * y(t) * G(t,T)^2)
  ```
  
  The HJM model implentation supports constant mean-reversion rate `k` and
  `sigma(t)` can be an arbitrary function of `t` and `r`. The implementation
  uses Euler discretization to simulate the HJM model.
  
  #### Example. Simulate a 4-factor HJM process.
  
  ```python
  import numpy as np
  import tensorflow.compat.v2 as tf
  import tf_quant_finance as tff
  
  dtype = tf.float64
  def discount_fn(x):
    return 0.01 * tf.ones_like(x, dtype=dtype)
  
  process = tff.models.hjm.QuasiGaussianHJM(
      dim=4,
      mean_reversion=[0.03, 0.01, 0.02, 0.005],  # constant mean-reversion
      volatility=[0.01, 0.011, 0.015, 0.008],  # constant volatility
      initial_discount_rate_fn=discount_fn,
      dtype=dtype)
  times = np.array([0.1, 1.0, 2.0, 3.0])
  short_rate_paths, discount_paths, _, _ = process.sample_paths(
      times,
      num_samples=100000,
      time_step=0.1,
      random_type=tff.math.random.RandomType.STATELESS_ANTITHETIC,
      seed=[1, 2],
      skip=1000000)
  ```
  
  #### References:
    [1]: Leif B. G. Andersen and Vladimir V. Piterbarg. Interest Rate Modeling.
    Volume II: Term Structure Models.
docstr: Initializes a batch of HJM models.
  
  Args:
    dim: A Python scalar which corresponds to the number of factors
      comprising the model.
    mean_reversion: A real positive `Tensor` of shape `batch_shape + [dim]`.
      `batch_shape` denotes the shape of independent HJM models within the
      batch. Corresponds to the mean reversion rate of each factor.
    volatility: A real positive `Tensor` of the same `dtype` and shape as
      `mean_reversion` or a callable with the following properties:
      (a)  The callable should accept a scalar `Tensor` `t` and a `Tensor`
      `r(t)` of shape `batch_shape + [num_samples]` and returns a `Tensor` of
      shape compatible with `batch_shape + [num_samples, dim]`. The variable
      `t`  stands for time and `r(t)` is the short rate at time `t`. The
      function returns instantaneous volatility `sigma(t) = sigma(t, r(t))`.
      When `volatility` is specified as a real `Tensor`, each factor is
      assumed to have a constant instantaneous volatility  and the  model is
      effectively a Gaussian HJM model.
      Corresponds to the instantaneous volatility of each factor.
    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 shape `batch_shape + input_shape`.
      Corresponds to the zero coupon bond yield at the present time for the
      input expiry time.
    corr_matrix: A `Tensor` of shape `batch_shape + [dim, dim]` and the same
      `dtype` as `mean_reversion`.
      Corresponds to the correlation matrix `Rho`.
    validate_args: Optional boolean flag to enable validation of the input
      correlation matrix. If the flag is enabled and the input correlation
      matrix is not positive semidefinite, an error is raised.
      Default value: False.
    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
      `quasi_gaussian_hjm_model`.
docstr_tok: ["Initializes","a","batch","of","HJM","models",".","Args",":","dim",":","A","Python","scalar","which","corresponds","to","the","number","of","factors","comprising","the","model",".","mean_reversion",":","A","real","positive","`","Tensor","`","of","shape","`","batch_shape","+","[","dim","]","`",".","`","batch_shape","`","denotes","the","shape","of","independent","HJM","models","within","the","batch",".","Corresponds","to","the","mean","reversion","rate","of","each","factor",".","volatility",":","A","real","positive","`","Tensor","`","of","the","same","`","dtype","`","and","shape","as","`","mean_reversion","`","or","a","callable","with","the","following","properties",":","(","a",")","The","callable","should","accept","a","scalar","`","Tensor","`","`","t","`","and","a","`","Tensor","`","`","r","(","t",")","`","of","shape","`","batch_shape","+","[","num_samples","]","`","and","returns","a","`","Tensor","`","of","shape","compatible","with","`","batch_shape","+","[","num_samples",",","dim","]","`",".","The","variable","`","t","`","stands","for","time","and","`","r","(","t",")","`","is","the","short","rate","at","time","`","t","`",".","The","function","returns","instantaneous","volatility","`","sigma","(","t",")","=","sigma","(","t",",","r","(","t",")",")","`",".","When","`","volatility","`","is","specified","as","a","real","`","Tensor","`",",","each","factor","is","assumed","to","have","a","constant","instantaneous","volatility","and","the","model","is","effectively","a","Gaussian","HJM","model",".","Corresponds","to","the","instantaneous","volatility","of","each","factor",".","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","shape","`","batch_shape","+","input_shape","`",".","Corresponds","to","the","zero","coupon","bond","yield","at","the","present","time","for","the","input","expiry","time",".","corr_matrix",":","A","`","Tensor","`","of","shape","`","batch_shape","+","[","dim",",","dim","]","`","and","the","same","`","dtype","`","as","`","mean_reversion","`",".","Corresponds","to","the","correlation","matrix","`","Rho","`",".","validate_args",":","Optional","boolean","flag","to","enable","validation","of","the","input","correlation","matrix",".","If","the","flag","is","enabled","and","the","input","correlation","matrix","is","not","positive","semidefinite",",","an","error","is","raised",".","Default","value",":","False",".","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","`","quasi_gaussian_hjm_model","`","."]
code: def __init__(
      self,
      dim: int,
      mean_reversion: types.RealTensor,
      volatility: Union[
          types.RealTensor, Callable[..., types.RealTensor]
      ],
      initial_discount_rate_fn: Callable[..., types.RealTensor],
      corr_matrix: types.RealTensor = None,
      validate_args: bool = False,
      dtype: tf.DType = None,
      name: str = None,
  ):
      """Initializes a batch of HJM models.
  
      Args:
        dim: A Python scalar which corresponds to the number of factors
          comprising the model.
        mean_reversion: A real positive `Tensor` of shape `batch_shape + [dim]`.
          `batch_shape` denotes the shape of independent HJM models within the
          batch. Corresponds to the mean reversion rate of each factor.
        volatility: A real positive `Tensor` of the same `dtype` and shape as
          `mean_reversion` or a callable with the following properties:
          (a)  The callable should accept a scalar `Tensor` `t` and a `Tensor`
          `r(t)` of shape `batch_shape + [num_samples]` and returns a `Tensor` of
          shape compatible with `batch_shape + [num_samples, dim]`. The variable
          `t`  stands for time and `r(t)` is the short rate at time `t`. The
          function returns instantaneous volatility `sigma(t) = sigma(t, r(t))`.
          When `volatility` is specified as a real `Tensor`, each factor is
          assumed to have a constant instantaneous volatility  and the  model is
          effectively a Gaussian HJM model.
          Corresponds to the instantaneous volatility of each factor.
        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 shape `batch_shape + input_shape`.
          Corresponds to the zero coupon bond yield at the present time for the
          input expiry time.
        corr_matrix: A `Tensor` of shape `batch_shape + [dim, dim]` and the same
          `dtype` as `mean_reversion`.
          Corresponds to the correlation matrix `Rho`.
        validate_args: Optional boolean flag to enable validation of the input
          correlation matrix. If the flag is enabled and the input correlation
          matrix is not positive semidefinite, an error is raised.
          Default value: False.
        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
          `quasi_gaussian_hjm_model`.
      """
      self._name = name or "quasi_gaussian_hjm_model"
      with tf.name_scope(self._name):
          self._dtype = dtype or tf.float32
          # x has dimensionality of `dim` and y `dim * dim`
          self._dim = dim + dim**2
          self._factors = dim
  
          def _instant_forward_rate_fn(t):
              t = tf.convert_to_tensor(t, dtype=self._dtype)
  
              def _log_zero_coupon_bond(x):
                  # Shape `x.shape`
                  r = tf.convert_to_tensor(
                      initial_discount_rate_fn(x), dtype=self._dtype
                  )
                  # Shape `x.shape`
                  return -r * x
  
              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):
              return tf.convert_to_tensor(
                  initial_discount_rate_fn(t), dtype=self._dtype
              )
  
          self._instant_forward_rate_fn = _instant_forward_rate_fn
          self._initial_discount_rate_fn = _initial_discount_rate_fn
          mean_reversion = tf.convert_to_tensor(
              mean_reversion, dtype=dtype, name="mean_reversion"
          )
  
          def _infer_batch_shape():
              zero = tf.constant([0], dtype=self._dtype)
              return _initial_discount_rate_fn(zero).shape.as_list()[
                  :-1
              ]
  
          self._batch_shape = _infer_batch_shape()
          self._batch_rank = len(self._batch_shape)
          self._mean_reversion = mean_reversion
  
          # Setup volatility
          if callable(volatility):
              self._volatility = volatility
          else:
              volatility = tf.convert_to_tensor(volatility, dtype=dtype)
              # Add a dimension corresponding to `num_samples` during simulations
              if self._batch_rank > 0:
                  volatility = tf.expand_dims(
                      volatility, axis=self._batch_rank
                  )
  
              def _tensor_to_volatility_fn(t, r):
                  del t, r
                  return volatility
  
              self._volatility = _tensor_to_volatility_fn
  
          if corr_matrix is None:
              corr_matrix = tf.eye(
                  dim,
                  dim,
                  batch_shape=self._batch_shape,
                  dtype=self._dtype,
              )
          self._rho = tf.convert_to_tensor(
              corr_matrix, dtype=self._dtype, name="rho"
          )
          if validate_args:
              try:
                  self._sqrt_rho = tf.linalg.cholesky(self._rho)
              except:
                  raise ValueError(
                      "The input correlation matrix is not "
                      "positive semidefinite."
                  )
          else:
              self._sqrt_rho = _get_valid_sqrt_matrix(self._rho)
  
      # Volatility function
      def _vol_fn(t, state):
          """Volatility function of qG-HJM."""
          # Get parameter values at time `t`
          x = state[..., : self._factors]
          batch_shape_x = x.shape.as_list()[:-1]
          r_t = self._instant_forward_rate_fn(t) + tf.reduce_sum(
              x, axis=-1, keepdims=True
          )
          volatility = self._volatility(t, r_t)
          volatility = tf.expand_dims(volatility, axis=-1)
  
          diffusion_x = tf.broadcast_to(
              tf.expand_dims(self._sqrt_rho, axis=self._batch_rank)
              * volatility,
              batch_shape_x + [self._factors, self._factors],
          )
          paddings = tf.constant(
              [[0, 0]] * len(batch_shape_x)
              + [[0, self._factors**2], [0, self._factors**2]],
              dtype=tf.int32,
          )
          diffusion = tf.pad(diffusion_x, paddings)
  
          return diffusion
  
      # Drift function
      def _drift_fn(t, state):
          """Drift function of qG-HJM."""
          x = state[..., : self._factors]
          y = state[..., self._factors :]
  
          batch_shape_x = x.shape.as_list()[:-1]
          y = tf.reshape(
              y, batch_shape_x + [self._factors, self._factors]
          )
          r_t = self._instant_forward_rate_fn(t) + tf.reduce_sum(
              x, axis=-1, keepdims=True
          )
          volatility = self._volatility(t, r_t)
          volatility = tf.expand_dims(volatility, axis=-1)
          # create matrix v(i,j) = vol(i)*vol(j)
          volatility_squared = tf.linalg.matmul(
              volatility, volatility, transpose_b=True
          )
          # create a matrix k2(i,j) = k(i) + k(j)
          mr2 = tf.expand_dims(self._mean_reversion, axis=-1)
          perm = list(range(self._batch_rank)) + [
              self._batch_rank + 1,
              self._batch_rank,
          ]
          mr2 = mr2 + tf.transpose(mr2, perm=perm)
          # Add a dimension corresponding to `num_samples`
          mr2 = tf.expand_dims(mr2, axis=self._batch_rank)
  
          mr = self._mean_reversion
          if self._batch_rank > 0:
              mr = tf.expand_dims(self._mean_reversion, axis=1)
          drift_x = tf.math.reduce_sum(y, axis=-1) - mr * x
  
          drift_y = (
              tf.expand_dims(self._rho, axis=self._batch_rank)
              * volatility_squared
              - mr2 * y
          )
          drift_y = tf.reshape(
              drift_y, batch_shape_x + [self._factors * self._factors]
          )
          drift = tf.concat([drift_x, drift_y], axis=-1)
          return drift
  
      super(QuasiGaussianHJM, self).__init__(
          self._dim, _drift_fn, _vol_fn, self._dtype, self._name
      )
  
file_path: tf-quant-finance/tf_quant_finance/models/hjm/quasi_gaussian_hjm.py
imports: from typing import Callable, Union, Tuple
  import tensorflow.compat.v2
  from tf_quant_finance import types
  from tf_quant_finance.math import gradient
  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 QuasiGaussianHJM object from inherited base classes: generic_ito_process with example: obj = QuasiGaussianHJM(dim, mean_reversion, volatility, initial_discount_rate_fn, corr_matrix, validate_args, dtype, name)