code: def _forward_matmul_one_part(
      kernel, x1, x2, x, part_size, part_index, remainder_part_size=None
  ):
      """Applies a single chunk of matmul split along the axis defined by `x2`."""
      x2_ax = tf.rank(x2) - kernel.feature_ndims - 1
      x2_ax_selector = tf.equal(tf.range(tf.rank(x2)), x2_ax)
      begin_x2 = tf.where(x2_ax_selector, part_size * part_index, 0)
      size_x2 = tf.where(
          x2_ax_selector,
          part_size
          if remainder_part_size is None
          else remainder_part_size,
          tf.shape(x2),
      )
  
      x_ax = tf.rank(x) - 2
      x_ax_selector = tf.equal(tf.range(tf.rank(x)), x_ax)
      begin_x = tf.where(x_ax_selector, part_size * part_index, 0)
      size_x = tf.where(
          x_ax_selector,
          part_size
          if remainder_part_size is None
          else remainder_part_size,
          tf.shape(x),
      )
  
      return tf.matmul(
          kernel.matrix(x1, _slice(x2, begin_x2, size_x2, x2_ax)),
          _slice(x, begin_x, size_x, x_ax),
      )
  
imports: import functools
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  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.python.ops import control_flow_util
label_desc: Use image node_id 2 for calling a global function with example usage: _forward_matmul_one_part(kernel, x1, x2, x, part_size, part_index, remainder_part_size) and returns: tf