def matmul_2d_index_add(array: Array, indices: Tuple[Array, Array],
                        values: Array) -> Array:
  """Convenience function for 2D index add.

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

  Returns:
    Added array: array[indices[0], indices[1]] + values.
  """
  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)
  encoding_addition = jnp.einsum('ib,i...->b...', one_hot_indices, values)
  added_array = flat_array + encoding_addition
  added_array = jnp.reshape(added_array, array.shape)
  return added_array
Convenience function for 2D index add.

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

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