src/transformers/models/layoutlmv3/modeling_tf_layoutlmv3.py/TFLayoutLMv3MainLayer/get_extended_attention_mask
class TFLayoutLMv3MainLayer:    def get_extended_attention_mask(self, attention_mask: tf.Tensor) -> tf.Tensor:
        # Adapted from transformers.modelling_utils.ModuleUtilsMixin.get_extended_attention_mask

        n_dims = len(attention_mask.shape)

        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        if n_dims == 3:
            extended_attention_mask = tf.expand_dims(attention_mask, axis=1)
        elif n_dims == 2:
            # Provided a padding mask of dimensions [batch_size, seq_length].
            # Make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length].
            extended_attention_mask = tf.expand_dims(attention_mask, axis=1)  # (batch_size, 1, seq_length)
            extended_attention_mask = tf.expand_dims(extended_attention_mask, axis=1)  # (batch_size, 1, 1, seq_length)
        else:
            raise ValueError(f"Wrong shape for attention_mask (shape {attention_mask.shape}).")

        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
        # masked positions, this operation will create a tensor which is 0.0 for
        # positions we want to attend and -10000.0 for masked positions.
        # Since we are adding it to the raw scores before the softmax, this is
        # effectively the same as removing these entirely.
        extended_attention_mask = tf.cast(extended_attention_mask, self.compute_dtype)
        extended_attention_mask = (1.0 - extended_attention_mask) * LARGE_NEGATIVE

        return extended_attention_mask