src/transformers/models/deberta_v2/modeling_tf_deberta_v2.py/take_along_axis
def take_along_axis(x, indices):
    # Only a valid port of np.take_along_axis when the gather axis is -1

    # TPU + gathers and reshapes don't go along well -- see https://github.com/huggingface/transformers/issues/18239
    if isinstance(tf.distribute.get_strategy(), tf.distribute.TPUStrategy):
        # [B, S, P] -> [B, S, P, D]
        one_hot_indices = tf.one_hot(indices, depth=x.shape[-1], dtype=x.dtype)

        # if we ignore the first two dims, this is equivalent to multiplying a matrix (one hot) by a vector (x)
        # grossly abusing notation: [B, S, P, D] . [B, S, D] = [B, S, P]
        gathered = tf.einsum("ijkl,ijl->ijk", one_hot_indices, x)

    # GPUs, on the other hand, prefer gathers instead of large one-hot+matmuls
    else:
        gathered = tf.gather(x, indices, batch_dims=2)

    return gathered