code: def _apply_robin_boundary_conditions(
      value_grid,
      dim,
      batch_rank,
      boundary_conditions,
      has_default_lower_boundary,
      has_default_upper_boundary,
      lower_trim_indices,
      upper_trim_indices,
      coord_grid,
      superdiag,
      diag,
      subdiag,
      delta,
      t,
  ):
      """Updates contributions according to boundary conditions."""
      # This is analogous to _apply_boundary_conditions_to_discretized_equation in
      # pde_kernels.py. The difference is that we work with the given spatial
      # dimension. In particular, in all the tensor slices we have to slice
      # into the dimension `batch_rank + dim` instead of the last dimension.
      # We do not perform the updates where the boundary condition is default
  
      # If both boundareis are default, there is no need to update the
      # space-discretization matrix
      if (
          has_default_lower_boundary[dim]
          and has_default_upper_boundary[dim]
      ):
          return (superdiag, diag, subdiag), tf.zeros_like(diag)
  
      # Remove current dimension to the boundary conditions
      lower_trim_indices_bc = (
          lower_trim_indices[:dim] + lower_trim_indices[dim + 1 :]
      )
      upper_trim_indices_bc = (
          upper_trim_indices[:dim] + upper_trim_indices[dim + 1 :]
      )
  
      def reshape_fn(bound_coeff):
          """Reshapes boundary coefficient."""
          # Say the grid shape is (b, nz, ny, nx), and dim = 1.
          # The boundary condition coefficients are expected to have shape
          # (b, nz, nx). We need to:
          # - Trim the boundaries: nz -> nz-2, nx -> nx-2, because we work with
          # the inner part of the grid here.
          # - Expand dimension batch_rank+dim=2, because broadcasting won't always
          # do this correctly in subsequent computations: if a has shape (5, 1) and
          # b has shape (5,) then a*b has shape (5, 5)!
          # Thus this function turns (b, nz, nx) into (b, nz-2, 1, nx-2).
          return _reshape_boundary_conds(
              bound_coeff,
              trim_from=batch_rank,
              expand_dim_at=batch_rank + dim,
              lower_trim_indices=lower_trim_indices_bc,
              upper_trim_indices=upper_trim_indices_bc,
          )
  
      # Retrieve the boundary conditions in the form alpha V + beta V' = gamma.
      if has_default_lower_boundary[dim]:
          # No need for the BC as default BC was applied
          alpha_l, beta_l, gamma_l = None, None, None
      else:
          alpha_l, beta_l, gamma_l = boundary_conditions[dim][0](
              t, coord_grid
          )
      if has_default_upper_boundary[dim]:
          # No need for the BC as default BC was applied
          alpha_u, beta_u, gamma_u = None, None, None
      else:
          alpha_u, beta_u, gamma_u = boundary_conditions[dim][1](
              t, coord_grid
          )
  
      alpha_l, beta_l, gamma_l, alpha_u, beta_u, gamma_u = (
          _prepare_boundary_conditions(b, value_grid, batch_rank, dim)
          for b in (alpha_l, beta_l, gamma_l, alpha_u, beta_u, gamma_u)
      )
      alpha_l, beta_l, gamma_l, alpha_u, beta_u, gamma_u = map(
          reshape_fn,
          (alpha_l, beta_l, gamma_l, alpha_u, beta_u, gamma_u),
      )
      slice_dim = dim + batch_rank
      subdiag_first = _slice(subdiag, slice_dim, 0, 1)
      superdiag_last = _slice(superdiag, slice_dim, -1, 0)
      diag_inner = _slice(diag, slice_dim, 1, -1)
  
      if beta_l is None and beta_u is None:
          # Dirichlet or default conditions on both boundaries. In this case there are
          # no corrections to the tridiagonal matrix, so we can take a shortcut.
          if has_default_lower_boundary[dim]:
              # Inhomgeneous term is zero for default BC
              first_inhomog_element = tf.zeros_like(subdiag_first)
          else:
              first_inhomog_element = subdiag_first * gamma_l / alpha_l
          if has_default_upper_boundary[dim]:
              # Inhomgeneous term is zero for default BC
              last_inhomog_element = tf.zeros_like(superdiag_last)
          else:
              last_inhomog_element = superdiag_last * gamma_u / alpha_u
          inhomog_term = _append_first_and_last(
              first_inhomog_element,
              tf.zeros_like(diag_inner),
              last_inhomog_element,
              axis=slice_dim,
          )
          return (superdiag, diag, subdiag), inhomog_term
  
      # A few more slices we're going to need.
      subdiag_last = _slice(subdiag, slice_dim, -1, 0)
      subdiag_except_last = _slice(subdiag, slice_dim, 0, -1)
      superdiag_first = _slice(superdiag, slice_dim, 0, 1)
      superdiag_except_first = _slice(superdiag, slice_dim, 1, 0)
      diag_first = _slice(diag, slice_dim, 0, 1)
      diag_last = _slice(diag, slice_dim, -1, 0)
  
      # Convert the boundary conditions into the form v0 = xi1 v1 + xi2 v2 + eta,
      # and calculate corrections to the tridiagonal matrix and the inhomogeneous
      # term.
      if has_default_lower_boundary[dim]:
          # No update for the default BC
          diag_first_correction = 0
          superdiag_correction = 0
          first_inhomog_element = tf.zeros_like(subdiag_first)
      else:
          xi1, xi2, eta = _discretize_boundary_conditions(
              delta, delta, alpha_l, beta_l, gamma_l
          )
          diag_first_correction = subdiag_first * xi1
          superdiag_correction = subdiag_first * xi2
          first_inhomog_element = subdiag_first * eta
      if has_default_upper_boundary[dim]:
          # Inhomgeneous term is zero for default BC
          diag_last_correction = 0
          subdiag_correction = 0
          last_inhomog_element = tf.zeros_like(superdiag_last)
      else:
          xi1, xi2, eta = _discretize_boundary_conditions(
              delta, delta, alpha_u, beta_u, gamma_u
          )
          diag_last_correction = superdiag_last * xi1
          subdiag_correction = superdiag_last * xi2
          last_inhomog_element = superdiag_last * eta
      diag = _append_first_and_last(
          diag_first + diag_first_correction,
          diag_inner,
          diag_last + diag_last_correction,
          axis=slice_dim,
      )
  
      superdiag = _append_first(
          superdiag_first + superdiag_correction,
          superdiag_except_first,
          axis=slice_dim,
      )
      subdiag = _append_last(
          subdiag_except_last,
          subdiag_last + subdiag_correction,
          axis=slice_dim,
      )
      inhomog_term = _append_first_and_last(
          first_inhomog_element,
          tf.zeros_like(diag_inner),
          last_inhomog_element,
          axis=slice_dim,
      )
      return (superdiag, diag, subdiag), inhomog_term
  
imports: import tensorflow.compat.v2
  from tf_quant_finance import utils
label_desc: Use image node_id 8 for calling a global function with example usage: _apply_robin_boundary_conditions(value_grid, dim, batch_rank, boundary_conditions, has_default_lower_boundary, has_default_upper_boundary, lower_trim_indices, upper_trim_indices, coord_grid, superdiag, diag, subdiag, delta, t) and returns: superdiag,  diag,  subdiag,  inhomog_term, superdiag,  diag,  subdiag,  tf, _reshape_boundary_conds, superdiag,  diag,  subdiag,  inhomog_term