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"]
code: def _matmul(self, x, adjoint=False, adjoint_arg=False):
      # The product is: r @ base_mat @ r^T @ x
      #
      # Where:
      #
      # x:     [B] + [R, C]
      # r:     [B] + [R, M^n]
      # r^T:   [B] + [M^n, R]
      # base_mat:   [B] + [M^n, M^n]
      #
      # We do the products right to left.
      if adjoint or adjoint_arg:
          raise NotImplementedError(
              "`LinearOperatorInterpolatedPSDKernel.adjoint` not implemented."
          )
      x = tf.convert_to_tensor(x, self.dtype)
      if tf.get_static_value(ps.shape_slice(x, np.s_[-2:])) is None:
          raise ValueError(
              "The last two dimensions of `x` must be statically "
              f"known. Saw shape: {x.shape}"
          )
      x1, x2 = self._x1_x2()
  
      batch_shape = ps.broadcast_shape(
          ps.shape_slice(x, np.s_[:-2]), ps.shape_slice(x2, np.s_[:-2])
      )
      base_mat = self._base_mat()
  
      # y = r^T @ x
      # [B] + [*M^n, C]
      # We implement the transpose using autodiff, to make it easier to use custom
      # interpolation functions. TODO(siege): Maybe adjust the interpolation_fn
      # API to somehow allow the user to write a custom transpose.
      _, y = gradient.value_and_gradient(
          lambda inp: self._interpolate(  # pylint: disable=g-long-lambda
              x2, inp, axis=-1 - ps.dimension_size(x1, -1)
          ),
          tf.ones(
              ps.concat(
                  [
                      batch_shape,
                      ps.stack(
                          [self._num_interpolation_points]
                          * self._bounds_min.shape[-1]
                          + [ps.dimension_size(x, -1)]
                      ),
                  ],
                  -1,
              ),
              x.dtype,
          ),
          output_gradients=tf.broadcast_to(
              x,
              ps.concat(
                  [batch_shape, ps.shape_slice(x, np.s_[-2:])], -1
              ),
          ),
      )
      # [B] + [M^n, C]
      y = tf.reshape(
          y,
          ps.concat(
              [
                  ps.shape_slice(
                      y, np.s_[: -1 - ps.dimension_size(x1, -1)]
                  ),
                  ps.stack(
                      [
                          ps.dimension_size(base_mat, -1),
                          ps.dimension_size(y, -1),
                      ]
                  ),
              ],
              -1,
          ),
      )
  
      # [B] + [M^n, C]
      # y = base_mat @ r^T @ x
      y = base_mat @ y
  
      # [B] + [M..., C]
      y = tf.reshape(
          y,
          ps.concat(
              [
                  ps.shape_slice(y, np.s_[:-2]),
                  ps.stack(
                      [self._num_interpolation_points]
                      * self._bounds_min.shape[-1]
                      + [-1]
                  ),
              ],
              -1,
          ),
      )
      # [B] + [R, C]
      # y = r @ mat @ r^T @ x
      res = self._interpolate(
          x1, y, axis=-1 - ps.dimension_size(x1, -1)
      )
  
      if self.is_self_adjoint and USE_EXACT_DIAG:
          res += (self._exact_diag_part() - self._approx_diag_part())[
              ..., tf.newaxis
          ] * x
  
      if self._diag_shift is not None:
          res += self._diag_shift[..., tf.newaxis, tf.newaxis] * x
      return res
  
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 12 for calling the LinearOperatorInterpolatedPSDKernel obj's underlying member method code with example usage: obj._matmul(x, adjoint, adjoint_arg) and returns: res