def matmul_2d_index_select(array: Array, indices: Tuple[Array, Array]) -> Array:
  """Function subselects an array with 2D index using matrix multiplication.

  Args:
    array: [batch, seq, ...] array to be indexed over.
    indices: Tuple[[n_index], [n_index]] arrays of indices.

  Returns:
    Subselected array: array[indices[0], indices[1]].
  """
  batch_size = array.shape[0]
  seq_len = array.shape[1]
  flat_seq_len = seq_len * batch_size
  flat_array = jnp.reshape(array, (flat_seq_len, *array.shape[2:]))
  one_hot_indices = jax.nn.one_hot(
      indices[0] * seq_len + indices[1], flat_seq_len, dtype=array.dtype)
  value = jnp.einsum('b...,ib->i...', flat_array, one_hot_indices)
  return value
Function subselects an array with 2D index using matrix multiplication.

  Args:
    array: [batch, seq, ...] array to be indexed over.
    indices: Tuple[[n_index], [n_index]] arrays of indices.

  Returns:
    Subselected array: array[indices[0], indices[1]].
  
language/mentionmemory/utils/jax_utils.py