code: def _backward_matmul_one_part(
      dcovx,
      kernel_fn,
      kernel_args,
      x1,
      x2,
      x,
      part_size,
      part_index,
      remainder_part_size=None,
  ):
      """Applies a single chunk of backprop split along the axis defined by `x1`."""
      # Assume `cov = kernel.matrix(x1, x2)` has shape (A,B), and `x` has shape
      # (B,C). Then `cov @ x` had shape (A,C), as will `dcovx`. Below, we refer to
      # the "part" size as P.
      dcovx_ax = tf.rank(dcovx) - 2
      dcovx_ax_selector = tf.equal(tf.range(tf.rank(dcovx)), dcovx_ax)
      begin_dcovx = tf.where(
          dcovx_ax_selector, part_size * part_index, 0
      )
      size_dcovx = tf.where(
          dcovx_ax_selector,
          part_size
          if remainder_part_size is None
          else remainder_part_size,
          tf.shape(dcovx),
      )
      dcovxpart = _slice(
          dcovx, begin_dcovx, size_dcovx, dcovx_ax
      )  # PxC
      dcovpart = tf.matmul(
          dcovxpart, x, transpose_b=True
      )  # PxC @ (BxC).T = PxB
      with tf.GradientTape(watch_accessed_variables=False) as tape:
          # Here begins the "recomputed" part of the gradient.
          tape.watch((x1, x2, kernel_args))
          kernel = kernel_fn(*kernel_args)
          x1_ax = tf.rank(x1) - kernel.feature_ndims - 1
          x1_ax_selector = tf.equal(tf.range(tf.rank(x1)), x1_ax)
          begin_x1 = tf.where(x1_ax_selector, part_size * part_index, 0)
          size_x1 = tf.where(
              x1_ax_selector,
              part_size
              if remainder_part_size is None
              else remainder_part_size,
              tf.shape(x1),
          )
          covpart = kernel.matrix(
              _slice(x1, begin_x1, size_x1, x1_ax), x2
          )  # PxB
      dx1part, dx2part, dkernel_args = tape.gradient(
          covpart, (x1, x2, kernel_args), output_gradients=dcovpart
      )
      dxpart = tf.matmul(
          covpart, dcovxpart, transpose_a=True
      )  # (PxB).T @ PxC
      dxpart = tf.reduce_sum(
          dxpart, axis=tf.range(tf.rank(dxpart) - tf.rank(x))
      )
      return dx1part, dx2part, dxpart, dkernel_args
  
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 3 for calling a global function with example usage: _backward_matmul_one_part(dcovx, kernel_fn, kernel_args, x1, x2, x, part_size, part_index, remainder_part_size) and returns: dx1part,  dx2part,  dxpart,  dkernel_args