src/transformers/models/layoutlmv3/modeling_tf_layoutlmv3.py/TFLayoutLMv3Encoder/_cal_pos_emb
class TFLayoutLMv3Encoder:    def _cal_pos_emb(
        self,
        dense_layer: tf.keras.layers.Dense,
        position_ids: tf.Tensor,
        num_buckets: int,
        max_distance: int,
    ):
        rel_pos_matrix = tf.expand_dims(position_ids, axis=-2) - tf.expand_dims(position_ids, axis=-1)
        rel_pos = self.relative_position_bucket(rel_pos_matrix, num_buckets, max_distance)
        rel_pos_one_hot = tf.one_hot(rel_pos, depth=num_buckets, dtype=self.compute_dtype)
        embedding = dense_layer(rel_pos_one_hot)
        # batch_size, seq_length, seq_length, num_heads --> batch_size, num_heads, seq_length, seq_length
        embedding = tf.transpose(embedding, [0, 3, 1, 2])
        embedding = tf.cast(embedding, dtype=self.compute_dtype)
        return embedding