code: def _update_variance(
      mean_reversion,
      theta,
      volvol,
      rho,
      current_vol,
      time_step,
      normals,
      psi_c=1.5,
  ):
      """Updates variance value."""
      del rho
      psi_c = tf.convert_to_tensor(psi_c, dtype=mean_reversion.dtype)
      scaled_time = tf.exp(-mean_reversion * time_step)
      volvol_squared = volvol**2
      m = theta + (current_vol - theta) * scaled_time
      s_squared = (
          current_vol
          * volvol_squared
          * scaled_time
          / mean_reversion
          * (1 - scaled_time)
          + theta
          * volvol_squared
          / 2
          / mean_reversion
          * (1 - scaled_time) ** 2
      )
      psi = s_squared / m**2
      uniforms = 0.5 * (1 + tf.math.erf(normals / _SQRT_2))
      cond = psi < psi_c
      # Result where `cond` is true
      psi_inv = 2 / psi
      b_squared = psi_inv - 1 + tf.sqrt(psi_inv * (psi_inv - 1))
  
      a = m / (1 + b_squared)
      next_var_true = (
          a * (tf.sqrt(b_squared) + tf.squeeze(normals)) ** 2
      )
      # Result where `cond` is false
      p = (psi - 1) / (psi + 1)
      beta = (1 - p) / m
      next_var_false = (
          tf.where(
              uniforms > p,
              tf.math.log(1 - p) - tf.math.log(1 - uniforms),
              tf.zeros_like(uniforms),
          )
          / beta
      )
      next_var = tf.where(cond, next_var_true, next_var_false)
      return next_var
  
file_path: tf-quant-finance/tf_quant_finance/models/heston/heston_model.py
imports: from typing import Union, Callable, Optional
  import numpy
  import tensorflow.compat.v2
  from tf_quant_finance import types
  from tf_quant_finance import utils
  from tf_quant_finance.math import piecewise
  from tf_quant_finance.math import random_ops
  from tf_quant_finance.models import generic_ito_process
  from tf_quant_finance.models import utils
label_desc: Use image node_id 2 for calling a global function with example usage: _update_variance(mean_reversion, theta, volvol, rho, current_vol, time_step, normals, psi_c) and returns: next_var