code: def _construct_tridiagonal_matrix(
      value_grid,
      second_order_coeff,
      first_order_coeff,
      inner_second_order_coeff,
      inner_first_order_coeff,
      delta,
      dim,
      batch_rank,
      lower_trim_indices,
      upper_trim_indices,
      has_default_lower_boundary,
      has_default_upper_boundary,
      n_dims,
  ):
      """Constructs contributions of first and non-mixed second order terms."""
  
      # If the boundary condition for dimension `dim` is default, we need to
      # update trimming for that dimension to remove the boundary, and apply the
      # default BC separately.
      # Note that for all other dimensions, the space-discretization tridiagonal
      # matrix is trimmed according to the trimming indices.
      (
          trimmed_lower_indices,
          trimmed_upper_indices,
          _,
      ) = _remove_default_boundary(
          lower_trim_indices,
          upper_trim_indices,
          has_default_lower_boundary,
          has_default_upper_boundary,
          batch_rank,
          (dim,),
          n_dims,
      )
  
      zeros = tf.zeros_like(value_grid)
      zeros = _trim_boundaries(
          zeros,
          from_dim=batch_rank,
          lower_trim_indices=trimmed_lower_indices,
          upper_trim_indices=trimmed_upper_indices,
      )
  
      def create_trimming_shifts(dim_shift):
          # See _trim_boundaries. We need to apply shift only to the dimension `dim`.
          shifts = [0] * n_dims
          shifts[dim] = dim_shift
          return shifts
  
      # Discretize first-order term.
      if first_order_coeff is None and inner_first_order_coeff is None:
          # No first-order term.
          superdiag_first_order = zeros
          diag_first_order = zeros
          subdiag_first_order = zeros
      else:
          superdiag_first_order = -1 / (2 * delta)
          subdiag_first_order = 1 / (2 * delta)
          diag_first_order = (
              -superdiag_first_order - subdiag_first_order
          )
          if first_order_coeff is not None:
              first_order_coeff = _trim_boundaries(
                  first_order_coeff,
                  from_dim=batch_rank,
                  lower_trim_indices=trimmed_lower_indices,
                  upper_trim_indices=trimmed_upper_indices,
              )
              superdiag_first_order *= first_order_coeff
              subdiag_first_order *= first_order_coeff
              diag_first_order *= first_order_coeff
          if inner_first_order_coeff is not None:
              superdiag_first_order *= _trim_boundaries(
                  inner_first_order_coeff,
                  from_dim=batch_rank,
                  lower_trim_indices=trimmed_lower_indices,
                  upper_trim_indices=trimmed_upper_indices,
                  shifts=create_trimming_shifts(1),
              )
              subdiag_first_order *= _trim_boundaries(
                  inner_first_order_coeff,
                  from_dim=batch_rank,
                  lower_trim_indices=trimmed_lower_indices,
                  upper_trim_indices=trimmed_upper_indices,
                  shifts=create_trimming_shifts(-1),
              )
              diag_first_order *= _trim_boundaries(
                  inner_first_order_coeff,
                  from_dim=batch_rank,
                  lower_trim_indices=trimmed_lower_indices,
                  upper_trim_indices=trimmed_upper_indices,
              )
  
      # Discretize second-order term.
      if (
          second_order_coeff is None
          and inner_second_order_coeff is None
      ):
          # No second-order term.
          superdiag_second_order = zeros
          diag_second_order = zeros
          subdiag_second_order = zeros
      else:
          superdiag_second_order = -1 / (delta * delta)
          subdiag_second_order = -1 / (delta * delta)
          diag_second_order = (
              -superdiag_second_order - subdiag_second_order
          )
          if second_order_coeff is not None:
              second_order_coeff = _trim_boundaries(
                  second_order_coeff,
                  from_dim=batch_rank,
                  lower_trim_indices=trimmed_lower_indices,
                  upper_trim_indices=trimmed_upper_indices,
              )
              superdiag_second_order *= second_order_coeff
              subdiag_second_order *= second_order_coeff
              diag_second_order *= second_order_coeff
          if inner_second_order_coeff is not None:
              superdiag_second_order *= _trim_boundaries(
                  inner_second_order_coeff,
                  from_dim=batch_rank,
                  lower_trim_indices=trimmed_lower_indices,
                  upper_trim_indices=trimmed_upper_indices,
                  shifts=create_trimming_shifts(1),
              )
              subdiag_second_order *= _trim_boundaries(
                  inner_second_order_coeff,
                  from_dim=batch_rank,
                  lower_trim_indices=trimmed_lower_indices,
                  upper_trim_indices=trimmed_upper_indices,
                  shifts=create_trimming_shifts(-1),
              )
              diag_second_order *= _trim_boundaries(
                  inner_second_order_coeff,
                  from_dim=batch_rank,
                  lower_trim_indices=trimmed_lower_indices,
                  upper_trim_indices=trimmed_upper_indices,
              )
  
      superdiag = superdiag_first_order + superdiag_second_order
      subdiag = subdiag_first_order + subdiag_second_order
      diag = diag_first_order + diag_second_order
      return superdiag, diag, subdiag
  
imports: import tensorflow.compat.v2
  from tf_quant_finance import utils
label_desc: Use image node_id 3 for calling a global function with example usage: _construct_tridiagonal_matrix(value_grid, second_order_coeff, first_order_coeff, inner_second_order_coeff, inner_first_order_coeff, delta, dim, batch_rank, lower_trim_indices, upper_trim_indices, has_default_lower_boundary, has_default_upper_boundary, n_dims) and returns: superdiag,  diag,  subdiag, shifts