class_docstr: Structured interpolation to approximate a large matrix.
  
  This implements a component of the Structured Kernel Interpolation [1]
  algorithm. We approximate the pairwise `kernel` values for two inputs `x1`,
  `x2` with the following product:
  
  ```none
  k(x1, x2) = r(x1) @ k(u, u) @ r(x2)^T
  ```
  
  where `u` is a set of points regularly spaced on grid and `r` is an
  interpolation matrix. In short, instead of evaluating `kernel` on `x1` and
  `x2`, we evaluate it on `u`, which is chosen to contain fewer points than `x1`
  and `x2` and the interpolate using the interpolation matrix.
  
  This construction lets us compute matrix products efficiently. If `x1`/`x2`
  are of shape `O(n)` and `u` is of shape `O(m)`, this reduces the compute cost
  and memory to `O(n + m^2)` from `O(n^2)`.
  
  In practice, the interpolation matrix is implicitly defined using `interp_fn`,
  of which `tfp.math.batch_interp_regular_nd_grid` linear interpolation is a
  prototypical example. When `x1 == x2` we can relatively cheaply compute the
  diagonal of that matrix exactly, to preserve positive-semi-definiteness.
  
  Since we rely on a dense grid `u`, this method works best when the the kernel
  inputs are low dimensional (2 or 3).
  
  #### References
  
  [1]: Wilson, A. G., & Nickisch, H. Kernel Interpolation for Scalable
       Structured Gaussian Processes (KISS-GP). 2005. arXiv.
       http://arxiv.org/abs/1503.01057
class_docstr_tok: ["Structured","interpolation","to","approximate","a","large","matrix",".","This","implements","a","component","of","the","Structured","Kernel","Interpolation","[","1","]","algorithm",".","We","approximate","the","pairwise","`","kernel","`","values","for","two","inputs","`","x1","`",",","`","x2","`","with","the","following","product",":","``","`","none","k","(","x1",",","x2",")","=","r","(","x1",")","@","k","(","u",",","u",")","@","r","(","x2",")","^T","``","`","where","`","u","`","is","a","set","of","points","regularly","spaced","on","grid","and","`","r","`","is","an","interpolation","matrix",".","In","short",",","instead","of","evaluating","`","kernel","`","on","`","x1","`","and","`","x2","`",",","we","evaluate","it","on","`","u","`",",","which","is","chosen","to","contain","fewer","points","than","`","x1","`","and","`","x2","`","and","the","interpolate","using","the","interpolation","matrix",".","This","construction","lets","us","compute","matrix","products","efficiently",".","If","`","x1","`","\/","`","x2","`","are","of","shape","`","O","(","n",")","`","and","`","u","`","is","of","shape","`","O","(","m",")","`",",","this","reduces","the","compute","cost","and","memory","to","`","O","(","n","+","m^2",")","`","from","`","O","(","n^2",")","`",".","In","practice",",","the","interpolation","matrix","is","implicitly","defined","using","`","interp_fn","`",",","of","which","`","tfp.math.batch_interp_regular_nd_grid","`","linear","interpolation","is","a","prototypical","example",".","When","`","x1","==","x2","`","we","can","relatively","cheaply","compute","the","diagonal","of","that","matrix","exactly",",","to","preserve","positive-semi-definiteness",".","Since","we","rely","on","a","dense","grid","`","u","`",",","this","method","works","best","when","the","the","kernel","inputs","are","low","dimensional","(","2","or","3",")",".","#","#","#","#","References","[","1","]",":","Wilson",",","A.","G.",",","&","Nickisch",",","H.","Kernel","Interpolation","for","Scalable","Structured","Gaussian","Processes","(","KISS-GP",")",".","2005.","arXiv",".","http",":","\/\/arxiv.org\/abs\/1503.01057"]
docstr: Gets a row from the dense operator.
  
  Args:
    index: The index (indices) of the row[s] to get, may be scalar or up to
      batch shape.
  
  Returns:
    rows: Row[s] of the matrix, with shape `(...batch_shape..., num_cols)`.
      Effectively the same as `operator.to_dense()[..., index, :]` for a
      scalar `index`, analogous to gather for non-scalar.
docstr_tok: ["Gets","a","row","from","the","dense","operator",".","Args",":","index",":","The","index","(","indices",")","of","the","row","[","s","]","to","get",",","may","be","scalar","or","up","to","batch","shape",".","Returns",":","rows",":","Row","[","s","]","of","the","matrix",",","with","shape","`","(","...","batch_shape","...",",","num_cols",")","`",".","Effectively","the","same","as","`","operator.to_dense","(",")","[","...",",","index",",",":","]","`","for","a","scalar","`","index","`",",","analogous","to","gather","for","non-scalar","."]
code: def row(self, index):
      """Gets a row from the dense operator.
  
      Args:
        index: The index (indices) of the row[s] to get, may be scalar or up to
          batch shape.
  
      Returns:
        rows: Row[s] of the matrix, with shape `(...batch_shape..., num_cols)`.
          Effectively the same as `operator.to_dense()[..., index, :]` for a
          scalar `index`, analogous to gather for non-scalar.
      """
      index = tf.convert_to_tensor(index)
      x1, x2 = self._x1_x2()
      index_rank = tf.get_static_value(ps.rank(index))
      if index_rank is None:
          raise ValueError(
              "Rank of `index` must be known statically. Saw shape: "
              f"{index.shape}"
          )
      if index_rank == 0:
          scalar_index = True
          index = index[tf.newaxis]
      else:
          scalar_index = False
      x1_rows = tf.gather(x1, index, axis=-2)
  
      base_mat = self._base_mat()
      # [B] + [*M^n, *M^n]
      base_mat = tf.reshape(
          base_mat,
          ps.concat(
              [
                  ps.shape(base_mat)[:-2],
                  ps.repeat(
                      (self._num_interpolation_points,),
                      2 * ps.shape(self._bounds_min)[-1],
                  ),
              ],
              0,
          ),
      )
  
      # [B] + [R', *M^n]
      row = self._interpolate(
          x1_rows, base_mat, axis=-2 * int(x1.shape[-1])
      )
      # [B] + [R, R']
      row = self._interpolate(
          x2,
          distribution_util.move_dimension(
              row, -1 - int(x1.shape[-1]), -1
          ),
          axis=-1 - int(x1.shape[-1]),
      )
      # [B] + [R', R]
      row = distribution_util.move_dimension(row, -1, -2)
  
      if self.is_self_adjoint and USE_EXACT_DIAG:
          exact_val = self._exact_diag_part(x1_rows)
          approx_val = self._approx_diag_part(x1_rows)
          # [B] + [R']
          corr = exact_val - approx_val
          corr_shape = ps.shape(corr)
          batch_shape = corr_shape[:-1]
          corr_ndims = ps.size(corr_shape)
  
          # [R', R] + [B]
          zeros = tf.zeros(
              ps.concat(
                  [ps.shape(row)[-2:], ps.shape(corr)[:-1]], axis=-1
              ),
              self.dtype,
          )
          diag_corrections = tf.tensor_scatter_nd_update(
              zeros,
              ps.stack([index, ps.range(ps.size(index))], axis=-1),
              tf.transpose(
                  corr,
                  ps.concat(
                      [
                          [corr_ndims - 1],
                          ps.range(ps.size(batch_shape)),
                      ],
                      axis=-1,
                  ),
              ),
          )
          # [B] + [R', R]
          diag_corrections = tf.transpose(
              diag_corrections,
              ps.concat(
                  [2 + ps.range(ps.size(batch_shape)), [0, 1]], axis=-1
              ),
          )
          row += diag_corrections
  
      if self._diag_shift is not None:
          # Normally we'd use one_hot, but it doesn't understand batching.
          # tf.eye, on the other hand, doesn't understand numpy arrays... but we can
          # fix this.
          num_rows = ps.size(index)
          num_cols = ps.shape(row)[-1]
          num_rows_ = tf.get_static_value(num_rows)
          num_cols_ = tf.get_static_value(num_cols)
          if num_rows_ is not None:
              num_rows = int(num_rows_)
          if num_cols_ is not None:
              num_cols = int(num_cols_)
          row += (
              tf.eye(num_rows, num_cols, dtype=row.dtype)
              * self._diag_shift[..., tf.newaxis, tf.newaxis]
          )
  
      if scalar_index:
          row = row[..., 0, :]
      return row
  
imports: import functools
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.math import interpolation
label_desc: Use image node_id 18 for calling the LinearOperatorInterpolatedPSDKernel obj's underlying member method code with example usage: obj.row(index) and returns: row