def matmul_slice(array: Array, indices: Array) -> Array:
  """Convenience function for indexing that differs along first dimension.

  The function supports batched and not-batched input:

  (1) Input is NOT batched, meaning that `indices` has shape [n_index] and
  `array` has shape [array_len, ...]. The output is precisely array[indices].
  (2) Input is batched, so `indices` has shape [batch, n_index] and
  `array` has shape [batch, array_len, ...].

  Values of the `indices` should be from 0 to array_len - 1.

  Args:
    array: array to be indexed over with shape either [batch, array_len, ...] or
      [array_len, ...].
    indices: array of indices with shape either [batch, n_index] or [n_index].

  Returns:
    Sliced values.
  """
  if indices.ndim == 1:
    # Case (1): The input is not batched, so neither `array` not `indices` have
    # batch dimension
    array_len = array.shape[0]
    one_hot_indices = jax.nn.one_hot(indices, array_len, dtype=array.dtype)
    value = jnp.einsum('p...,ip->i...', array, one_hot_indices)
  else:
    # Case (2): The input is batched, so both `array` and `indices` have
    # extra batch dimension (the first one).
    array_len = array.shape[1]
    one_hot_indices = jax.nn.one_hot(indices, array_len, dtype=array.dtype)
    value = jnp.einsum('bp...,bip->bi...', array, one_hot_indices)
  return value
Convenience function for indexing that differs along first dimension.

  The function supports batched and not-batched input:

  (1) Input is NOT batched, meaning that `indices` has shape [n_index] and
  `array` has shape [array_len, ...]. The output is precisely array[indices].
  (2) Input is batched, so `indices` has shape [batch, n_index] and
  `array` has shape [batch, array_len, ...].

  Values of the `indices` should be from 0 to array_len - 1.

  Args:
    array: array to be indexed over with shape either [batch, array_len, ...] or
      [array_len, ...].
    indices: array of indices with shape either [batch, n_index] or [n_index].

  Returns:
    Sliced values.
  
language/mentionmemory/utils/jax_utils.py