code: def matmul(
      a: "_array.DistributedArray",
      b: "_array.DistributedArray",
      out: Optional["_array.DistributedArray"] = None,
      **kwargs,
  ) -> "_array.DistributedArray":
      """Matrix multiplication between distributed arrays.
  
      The arguments must have compatible :attr:`~DistributedArray.shape` and
      :attr:`~DistributedArray.index_map`.
  
      This operation converts its operands into the replica mode, and compute
      their product in the sum mode.
  
      Args:
          a, b: Input distributed arrays.
          out (optional): A location into which the result is stored. This option
              is currently not supported.
      Returns:
          The matrix product of the inputs.
  
      Example:
          >>> A = distributed_array(
          ...     cupy.arange(6).reshape(2, 3),
          ...     make_2d_index_map([0, 2], [0, 1, 3],
          ...                       [[{0}, {1, 2}]]))
          >>> B = distributed_array(
          ...     cupy.arange(12).reshape(3, 4),
          ...     make_2d_index_map([0, 1, 3], [0, 2, 4],
          ...                       [[{0}, {0}],
          ...                        [{1}, {2}]]))
          >>> C = A @ B
          >>> C.mode
          'sum'
          >>> C.all_chunks()
          {0: [array([[0, 0],
                      [0, 3]]),
               array([[0, 0],
                      [6, 9]])],
           1: [array([[20, 23],
                      [56, 65]])],
           2: [array([[26, 29],
                      [74, 83]])]}
          >>> C
          array([[20, 23, 26, 29],
                 [56, 68, 80, 92]])
  
      .. seealso:: :obj:`numpy.matmul`
      """
      if out is not None:
          raise RuntimeError("Argument `out` is not supported")
      for param in ("subok", "axes", "axis"):
          if param in kwargs:
              raise RuntimeError(f"Argument `{param}` is not supported")
      if not isinstance(a, _array.DistributedArray) or not isinstance(
          b, _array.DistributedArray
      ):
          raise RuntimeError(
              "Mixing a distributed array with a non-distributed array is not"
              " supported"
          )
  
      a = a._to_op_mode(_modes.REPLICA)
      b = b._to_op_mode(_modes.REPLICA)
  
      one_prepended = one_appended = False
      if a.ndim == 1:
          one_prepended = True
          a = _prepend_one_to_shape(a)
      if b.ndim == 1:
          one_appended = True
          b = _append_one_to_shape(b)
  
      n, m = a.shape[-2:]
      m2, p = b.shape[-2:]
      if m != m2 or a.shape[:-2] != b.shape[:-2]:
          raise ValueError("Shapes are incompatible")
  
      location_maps_a = _group_by_batch(a.shape, a.index_map)
      location_maps_b = _group_by_batch(b.shape, b.index_map)
      if location_maps_a.keys() != location_maps_b.keys():
          raise RuntimeError("Mismatched batch shapes")
  
      chunks_map: dict[int, list[_chunk._Chunk]] = {
          dev: [] for dev in a.devices
      }
      dtype = None
  
      for batch_idx in location_maps_a.keys():
          location_map_a = location_maps_a[batch_idx]
          location_map_b = location_maps_b[batch_idx]
  
          blocking = _find_blocking(location_map_a, location_map_b)
          plan = _make_execution_plan(
              blocking, location_map_a, location_map_b
          )
  
          index_prefix = _convert_to_slices(batch_idx)
          for block_a, block_b, dev in plan:
              loc_a = location_map_a[block_a]
              loc_b = location_map_b[block_b]
              chunk_a = a._chunks_map[dev][loc_a[dev]]
              chunk_b = b._chunks_map[dev][loc_b[dev]]
              chunk_a.flush(_modes.REPLICA)
              chunk_b.flush(_modes.REPLICA)
  
              index = index_prefix + (
                  slice(*block_a[0]),
                  slice(*block_b[1]),
              )
              with chunk_a.on_ready() as stream:
                  stream.wait_event(chunk_b.ready)
  
                  chunk_ab_array = cupy.linalg._product.matmul(
                      chunk_a.array, chunk_b.array, **kwargs
                  )
  
                  chunk_ab = _chunk._Chunk(
                      chunk_ab_array,
                      stream.record(),
                      index,
                      prevent_gc=(chunk_a, chunk_b),
                  )
                  chunks_map[dev].append(chunk_ab)
                  dtype = chunk_ab_array.dtype
  
      shape = a.shape[:-2] + (n, p)
      res = _array.DistributedArray(
          shape, dtype, chunks_map, _modes.SUM, a._comms
      )
  
      if one_prepended:
          res = _pop_front_from_shape(res)
      if one_appended:
          res = _pop_from_shape(res)
  
      return res
  
file_path: cupy/cupyx/distributed/array/_linalg.py
imports: import dataclasses
  import typing
  from typing import Callable, Optional
  import cupy
  from cupyx.distributed.array import _array
  from cupyx.distributed.array import _chunk
  from cupyx.distributed.array import _modes
label_desc: Use image node_id 11 for calling a global function with example usage: matmul(a, b, out) and returns: res