code: def _expected_exercise_fn(
      design, calibration_indices, continuation_value, exercise_value
  ):
      """Returns the expected continuation value for each path.
  
      Args:
        design: A real `Tensor` of shape `[batch_size, basis_size, num_samples]`.
        calibration_indices: A rank 1 integer `Tensor` denoting indices of samples
          used for regression.
        continuation_value: A `Tensor` of shape `[num_samples, batch_size]` and of
          the same dtype as `design`. The optimal value of the option conditional on
          not exercising now or earlier, taking future information into account.
        exercise_value: A `Tensor` of the same shape and dtype as
          `continuation_value`. Value of the option if exercised immideately at
          the current time
  
      Returns:
        A `Tensor` of the same shape and dtype as `continuation_value` whose
        `(n, v)`-th entry represents the expected continuation value of sample path
        `n` under the `v`-th payoff scheme.
      """
      # Zero out contributions from samples we'd never exercise at this point (i.e.,
      # these extra observations do not change the regression coefficients).
      mask = exercise_value > 0
      # Shape [batch_size, num_samples, basis_size]
      design_t = tf.transpose(design, [0, 2, 1])
      masked = tf.where(
          tf.expand_dims(tf.transpose(mask), axis=-1),
          design_t,
          tf.zeros_like(design_t),
      )
      if calibration_indices is None:
          submask = masked
          mask_cont_value = continuation_value
      else:
          # Shape [batch_size, num_calibration_samples, basis_size]
          submask = tf.gather(masked, calibration_indices, axis=1)
          # Shape [num_calibration_samples, batch_size]
          mask_cont_value = tf.gather(
              continuation_value, calibration_indices
          )
      # For design matrix X and response y, the coefficients beta of the best linear
      # unbiased estimate are contained in the equation X'X beta = X'y. Here `lhs`
      # is X'X and `rhs` is X'y, or rather a tensor of such left hand and right hand
      # sides, one for each payoff dimension.
      # Shape [batch_size, basis_size, basis_size]
      lhs = tf.matmul(submask, submask, transpose_a=True)
      # Use pseudo inverse for the regression matrix to ensure stability of the
      # algorithm.
      lhs_pinv = tf.linalg.pinv(lhs)
      # Shape [batch_size, basis_size, 1]
      rhs = tf.matmul(
          submask,
          tf.expand_dims(tf.transpose(mask_cont_value), axis=-1),
          transpose_a=True,
      )
      # Shape [batch_size, basis_size, 1]
      beta = tf.matmul(lhs_pinv, rhs)
      # Shape [batch_size, num_samples, 1]
      continuation = tf.matmul(design_t, beta)
      # Shape [num_samples, batch_size]
      return tf.nn.relu(tf.transpose(tf.squeeze(continuation, axis=-1)))
  
file_path: tf-quant-finance/tf_quant_finance/models/longstaff_schwartz/lsm.py
imports: from typing import Callable
  import tensorflow.compat.v2
  from tf_quant_finance import types
  from tf_quant_finance import utils
label_desc: Use image node_id 5 for calling a global function with example usage: _expected_exercise_fn(design, calibration_indices, continuation_value, exercise_value) and returns: tf