def matmul_index_add(array: Array, indices: Array, values: Array) -> Array:
  """Convenience function for index add that differs along first dimension.

  Args:
    array: [batch, seq, dim] array to be added to.
    indices: [batch, n_index] array of indices to add to.
    values: [batch, n_index, dim] array of values to add.

  Returns:
    Added array.
  """
  seq_len = array.shape[1]
  one_hot_indices = jax.nn.one_hot(indices, seq_len, dtype=array.dtype)
  encoding_addition = jnp.einsum('bqp,bqd->bpd', one_hot_indices, values)
  added_array = array + encoding_addition
  return added_array
Convenience function for index add that differs along first dimension.

  Args:
    array: [batch, seq, dim] array to be added to.
    indices: [batch, n_index] array of indices to add to.
    values: [batch, n_index, dim] array of values to add.

  Returns:
    Added array.
  
language/mentionmemory/utils/jax_utils.py