code: def _chunked_matmul(
      kernel_fn,
      kernel_args,
      x1,
      x2,
      x,
      num_matmul_parts,
      operator_shape,
  ):
      """Implements while-loop matmul with recomputed gradients."""
      # We can't use kwargs with tf.custom_gradient in graph mode, so we close over.
      kernel_args_structure = tf.nest.pack_sequence_as(
          kernel_args, [0] * len(tf.nest.flatten(kernel_args))
      )
  
      @tf.custom_gradient
      def _chunked_matmul_cgrad(x1, x2, x, *kernel_args):
          """Chunk-at-a-time matrix multiplication and backprop."""
          kernel_args = tf.nest.pack_sequence_as(
              kernel_args_structure, kernel_args
          )
          kernel = kernel_fn(*kernel_args)
          fwd_ax_size = tf.shape(x2)[-kernel.feature_ndims - 1]
          fwd_part_size = fwd_ax_size // num_matmul_parts
  
          dist_ctx = tf.distribute.get_replica_context()
          replica_id = dist_ctx.replica_id_in_sync_group
          num_replicas = dist_ctx.num_replicas_in_sync
          replica_num_parts = (
              num_matmul_parts // num_replicas
              + tf.cast(
                  num_matmul_parts % num_replicas > replica_id,
                  num_matmul_parts.dtype,
              )
          )
          replica_begin = (
              num_matmul_parts // num_replicas
          ) * replica_id + tf.minimum(
              replica_id, num_matmul_parts % num_replicas
          )
  
          def cond(i, _):
              return i < replica_begin + replica_num_parts
  
          def body(i, covx):
              return i + 1, covx + _forward_matmul_one_part(
                  kernel, x1, x2, x, fwd_part_size, i
              )
  
          result_batch_shape = tf.broadcast_dynamic_shape(
              operator_shape[:-2], tf.shape(x)[:-2]
          )
          result_shape = tf.concat(
              [
                  result_batch_shape,
                  [operator_shape[-2], tf.shape(x)[-1]],
              ],
              axis=0,
          )
          _, covx = tf.while_loop(
              cond,
              body,
              (replica_begin, tf.zeros(result_shape, dtype=x.dtype)),
              back_prop=False,
              parallel_iterations=1,
          )
  
          remainder = _forward_matmul_one_part(
              kernel,
              x1,
              x2,
              x,
              fwd_part_size,
              num_matmul_parts,
              remainder_part_size=fwd_ax_size
              - (num_matmul_parts * fwd_part_size),
          )
  
          covx = (
              dist_ctx.all_reduce(tf.distribute.ReduceOp.SUM, covx)
              + remainder
          )
  
          del result_batch_shape, result_shape
  
          def grad_fn(dcovx):
              """Chunk-at-a-time backprop."""
              # Backward, we partition along the `x1`-defined axis.
              bwd_ax_size = tf.shape(x1)[-kernel.feature_ndims - 1]
              bwd_part_size = bwd_ax_size // num_matmul_parts
  
              dist_ctx = tf.distribute.get_replica_context()
              replica_id = dist_ctx.replica_id_in_sync_group
              num_replicas = dist_ctx.num_replicas_in_sync
              replica_num_parts = (
                  num_matmul_parts // num_replicas
                  + tf.cast(
                      num_matmul_parts % num_replicas > replica_id,
                      num_matmul_parts.dtype,
                  )
              )
              replica_begin = (
                  num_matmul_parts // num_replicas
              ) * replica_id + tf.minimum(
                  replica_id, num_matmul_parts % num_replicas
              )
  
              def bw_cond(i, *_):
                  return i < replica_begin + replica_num_parts
  
              def bw_body(i, dx1, dx2, dx, dkernel_args):
                  """tf.while_loop body for backprop."""
                  (
                      dx1part,
                      dx2part,
                      dxpart,
                      dkernel_argspart,
                  ) = _backward_matmul_one_part(
                      dcovx,
                      kernel_fn,
                      kernel_args,
                      x1,
                      x2,
                      x,
                      bwd_part_size,
                      i,
                  )
                  dx1, dx2, dx, dkernel_args = tf.nest.pack_sequence_as(
                      (dx1, dx2, dx, dkernel_args),
                      [
                          a + b
                          for a, b in zip(  # pylint: disable=g-complex-comprehension
                              tf.nest.flatten(
                                  (dx1, dx2, dx, dkernel_args)
                              ),
                              tf.nest.flatten(
                                  (
                                      dx1part,
                                      dx2part,
                                      dxpart,
                                      dkernel_argspart,
                                  )
                              ),
                          )
                      ],
                  )
                  return i + 1, dx1, dx2, dx, dkernel_args
  
              _, dx1, dx2, dx, dkernel_args = tf.while_loop(
                  bw_cond,
                  bw_body,
                  (replica_begin,)
                  + tf.nest.map_structure(
                      tf.zeros_like, (x1, x2, x, kernel_args)
                  ),
                  back_prop=False,
                  parallel_iterations=1,
              )
              (
                  dx1rem,
                  dx2rem,
                  dxrem,
                  dkernel_argsrem,
              ) = _backward_matmul_one_part(
                  dcovx,
                  kernel_fn,
                  kernel_args,
                  x1,
                  x2,
                  x,
                  bwd_part_size,
                  num_matmul_parts,
                  remainder_part_size=bwd_ax_size
                  - (num_matmul_parts * bwd_part_size),
              )
              flat_xdevice = tf.nest.flatten(
                  (dx1, dx2, dx, dkernel_args)
              )
              flat_remainder = tf.nest.flatten(
                  (dx1rem, dx2rem, dxrem, dkernel_argsrem)
              )
              return tuple(
                  dist_ctx.all_reduce(tf.distribute.ReduceOp.SUM, a) + b
                  for a, b in zip(flat_xdevice, flat_remainder)
              )
  
          return covx, grad_fn
  
      return _chunked_matmul_cgrad(
          x1, x2, x, *tf.nest.flatten(kernel_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 4 for calling a global function with example usage: _chunked_matmul(kernel_fn, kernel_args, x1, x2, x, num_matmul_parts, operator_shape) and returns: _chunked_matmul_cgrad, covx,  grad_fn, unknown, unknown,  unknown, tuple, unknown, unknown,  dx1,  dx2,  dx,  dkernel_args