class LayoutLMv3TokenizerFast(PreTrainedTokenizerFast):
    vocab_files_names = VOCAB_FILES_NAMES
    model_input_names = ['input_ids', 'attention_mask']
    slow_tokenizer_class = LayoutLMv3Tokenizer

    def __init__(self, vocab_file=None, merges_file=None, tokenizer_file=None, errors='replace', bos_token='<s>', eos_token='</s>', sep_token='</s>', cls_token='<s>', unk_token='<unk>', pad_token='<pad>', mask_token='<mask>', add_prefix_space=True, trim_offsets=True, cls_token_box=[0, 0, 0, 0], sep_token_box=[0, 0, 0, 0], pad_token_box=[0, 0, 0, 0], pad_token_label=-100, only_label_first_subword=True, **kwargs):
        super().__init__(vocab_file, merges_file, tokenizer_file=tokenizer_file, errors=errors, bos_token=bos_token, eos_token=eos_token, sep_token=sep_token, cls_token=cls_token, unk_token=unk_token, pad_token=pad_token, mask_token=mask_token, add_prefix_space=add_prefix_space, trim_offsets=trim_offsets, cls_token_box=cls_token_box, sep_token_box=sep_token_box, pad_token_box=pad_token_box, pad_token_label=pad_token_label, only_label_first_subword=only_label_first_subword, **kwargs)
        tokenizer_component = 'post_processor'
        tokenizer_component_instance = getattr(self.backend_tokenizer, tokenizer_component, None)
        if tokenizer_component_instance:
            state = json.loads(tokenizer_component_instance.__getstate__())
            if 'sep' in state:
                state['sep'] = tuple(state['sep'])
            if 'cls' in state:
                state['cls'] = tuple(state['cls'])
            changes_to_apply = False
            if state.get('add_prefix_space', add_prefix_space) != add_prefix_space:
                state['add_prefix_space'] = add_prefix_space
                changes_to_apply = True
            if state.get('trim_offsets', trim_offsets) != trim_offsets:
                state['trim_offsets'] = trim_offsets
                changes_to_apply = True
            if changes_to_apply:
                component_class = getattr(processors, state.pop('type'))
                new_value = component_class(**state)
                setattr(self.backend_tokenizer, tokenizer_component, new_value)
        self.cls_token_box = cls_token_box
        self.sep_token_box = sep_token_box
        self.pad_token_box = pad_token_box
        self.pad_token_label = pad_token_label
        self.only_label_first_subword = only_label_first_subword

    @add_end_docstrings(LAYOUTLMV3_ENCODE_KWARGS_DOCSTRING, LAYOUTLMV3_ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING)
    def __call__(self, text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]], text_pair: Optional[Union[PreTokenizedInput, List[PreTokenizedInput]]]=None, boxes: Optional[Union[List[List[int]], List[List[List[int]]]]]=None, word_labels: Optional[Union[List[int], List[List[int]]]]=None, add_special_tokens: bool=True, padding: Union[bool, str, PaddingStrategy]=False, truncation: Union[bool, str, TruncationStrategy]=None, max_length: Optional[int]=None, stride: int=0, pad_to_multiple_of: Optional[int]=None, padding_side: Optional[str]=None, return_tensors: Optional[Union[str, TensorType]]=None, return_token_type_ids: Optional[bool]=None, return_attention_mask: Optional[bool]=None, return_overflowing_tokens: bool=False, return_special_tokens_mask: bool=False, return_offsets_mapping: bool=False, return_length: bool=False, verbose: bool=True, **kwargs) -> BatchEncoding:
        """
        Main method to tokenize and prepare for the model one or several sequence(s) or one or several pair(s) of
        sequences with word-level normalized bounding boxes and optional labels.

        Args:
            text (`str`, `List[str]`, `List[List[str]]`):
                The sequence or batch of sequences to be encoded. Each sequence can be a string, a list of strings
                (words of a single example or questions of a batch of examples) or a list of list of strings (batch of
                words).
            text_pair (`List[str]`, `List[List[str]]`):
                The sequence or batch of sequences to be encoded. Each sequence should be a list of strings
                (pretokenized string).
            boxes (`List[List[int]]`, `List[List[List[int]]]`):
                Word-level bounding boxes. Each bounding box should be normalized to be on a 0-1000 scale.
            word_labels (`List[int]`, `List[List[int]]`, *optional*):
                Word-level integer labels (for token classification tasks such as FUNSD, CORD).
        """

        def _is_valid_text_input(t):
            if isinstance(t, str):
                return True
            elif isinstance(t, (list, tuple)):
                if len(t) == 0:
                    return True
                elif isinstance(t[0], str):
                    return True
                elif isinstance(t[0], (list, tuple)):
                    return len(t[0]) == 0 or isinstance(t[0][0], str)
                else:
                    return False
            else:
                return False
        if text_pair is not None:
            if not _is_valid_text_input(text):
                raise ValueError('text input must of type `str` (single example) or `List[str]` (batch of examples). ')
            if not isinstance(text_pair, (list, tuple)):
                raise ValueError('Words must be of type `List[str]` (single pretokenized example), or `List[List[str]]` (batch of pretokenized examples).')
        elif not isinstance(text, (list, tuple)):
            raise ValueError('Words must be of type `List[str]` (single pretokenized example), or `List[List[str]]` (batch of pretokenized examples).')
        if text_pair is not None:
            is_batched = isinstance(text, (list, tuple))
        else:
            is_batched = isinstance(text, (list, tuple)) and text and isinstance(text[0], (list, tuple))
        words = text if text_pair is None else text_pair
        if boxes is None:
            raise ValueError('You must provide corresponding bounding boxes')
        if is_batched:
            if len(words) != len(boxes):
                raise ValueError('You must provide words and boxes for an equal amount of examples')
            for words_example, boxes_example in zip(words, boxes):
                if len(words_example) != len(boxes_example):
                    raise ValueError('You must provide as many words as there are bounding boxes')
        elif len(words) != len(boxes):
            raise ValueError('You must provide as many words as there are bounding boxes')
        if is_batched:
            if text_pair is not None and len(text) != len(text_pair):
                raise ValueError(f'batch length of `text`: {len(text)} does not match batch length of `text_pair`: {len(text_pair)}.')
            batch_text_or_text_pairs = list(zip(text, text_pair)) if text_pair is not None else text
            is_pair = bool(text_pair is not None)
            return self.batch_encode_plus(batch_text_or_text_pairs=batch_text_or_text_pairs, is_pair=is_pair, boxes=boxes, word_labels=word_labels, add_special_tokens=add_special_tokens, padding=padding, truncation=truncation, max_length=max_length, stride=stride, pad_to_multiple_of=pad_to_multiple_of, padding_side=padding_side, return_tensors=return_tensors, return_token_type_ids=return_token_type_ids, return_attention_mask=return_attention_mask, return_overflowing_tokens=return_overflowing_tokens, return_special_tokens_mask=return_special_tokens_mask, return_offsets_mapping=return_offsets_mapping, return_length=return_length, verbose=verbose, **kwargs)
        else:
            return self.encode_plus(text=text, text_pair=text_pair, boxes=boxes, word_labels=word_labels, add_special_tokens=add_special_tokens, padding=padding, truncation=truncation, max_length=max_length, stride=stride, pad_to_multiple_of=pad_to_multiple_of, padding_side=padding_side, return_tensors=return_tensors, return_token_type_ids=return_token_type_ids, return_attention_mask=return_attention_mask, return_overflowing_tokens=return_overflowing_tokens, return_special_tokens_mask=return_special_tokens_mask, return_offsets_mapping=return_offsets_mapping, return_length=return_length, verbose=verbose, **kwargs)

    @add_end_docstrings(LAYOUTLMV3_ENCODE_KWARGS_DOCSTRING, LAYOUTLMV3_ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING)
    def batch_encode_plus(self, batch_text_or_text_pairs: Union[List[TextInput], List[TextInputPair], List[PreTokenizedInput]], is_pair: Optional[bool]=None, boxes: Optional[List[List[List[int]]]]=None, word_labels: Optional[Union[List[int], List[List[int]]]]=None, add_special_tokens: bool=True, padding: Union[bool, str, PaddingStrategy]=False, truncation: Union[bool, str, TruncationStrategy]=None, max_length: Optional[int]=None, stride: int=0, pad_to_multiple_of: Optional[int]=None, padding_side: Optional[str]=None, return_tensors: Optional[Union[str, TensorType]]=None, return_token_type_ids: Optional[bool]=None, return_attention_mask: Optional[bool]=None, return_overflowing_tokens: bool=False, return_special_tokens_mask: bool=False, return_offsets_mapping: bool=False, return_length: bool=False, verbose: bool=True, **kwargs) -> BatchEncoding:
        padding_strategy, truncation_strategy, max_length, kwargs = self._get_padding_truncation_strategies(padding=padding, truncation=truncation, max_length=max_length, pad_to_multiple_of=pad_to_multiple_of, verbose=verbose, **kwargs)
        return self._batch_encode_plus(batch_text_or_text_pairs=batch_text_or_text_pairs, is_pair=is_pair, boxes=boxes, word_labels=word_labels, add_special_tokens=add_special_tokens, padding_strategy=padding_strategy, truncation_strategy=truncation_strategy, max_length=max_length, stride=stride, pad_to_multiple_of=pad_to_multiple_of, padding_side=padding_side, return_tensors=return_tensors, return_token_type_ids=return_token_type_ids, return_attention_mask=return_attention_mask, return_overflowing_tokens=return_overflowing_tokens, return_special_tokens_mask=return_special_tokens_mask, return_offsets_mapping=return_offsets_mapping, return_length=return_length, verbose=verbose, **kwargs)

    def tokenize(self, text: str, pair: Optional[str]=None, add_special_tokens: bool=False, **kwargs) -> List[str]:
        batched_input = [(text, pair)] if pair else [text]
        encodings = self._tokenizer.encode_batch(batched_input, add_special_tokens=add_special_tokens, is_pretokenized=False, **kwargs)
        return encodings[0].tokens

    @add_end_docstrings(LAYOUTLMV3_ENCODE_KWARGS_DOCSTRING, LAYOUTLMV3_ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING)
    def encode_plus(self, text: Union[TextInput, PreTokenizedInput], text_pair: Optional[PreTokenizedInput]=None, boxes: Optional[List[List[int]]]=None, word_labels: Optional[List[int]]=None, add_special_tokens: bool=True, padding: Union[bool, str, PaddingStrategy]=False, truncation: Union[bool, str, TruncationStrategy]=None, max_length: Optional[int]=None, stride: int=0, pad_to_multiple_of: Optional[int]=None, padding_side: Optional[str]=None, return_tensors: Optional[Union[str, TensorType]]=None, return_token_type_ids: Optional[bool]=None, return_attention_mask: Optional[bool]=None, return_overflowing_tokens: bool=False, return_special_tokens_mask: bool=False, return_offsets_mapping: bool=False, return_length: bool=False, verbose: bool=True, **kwargs) -> BatchEncoding:
        """
        Tokenize and prepare for the model a sequence or a pair of sequences. .. warning:: This method is deprecated,
        `__call__` should be used instead.

        Args:
            text (`str`, `List[str]`, `List[List[str]]`):
                The first sequence to be encoded. This can be a string, a list of strings or a list of list of strings.
            text_pair (`List[str]` or `List[int]`, *optional*):
                Optional second sequence to be encoded. This can be a list of strings (words of a single example) or a
                list of list of strings (words of a batch of examples).
        """
        padding_strategy, truncation_strategy, max_length, kwargs = self._get_padding_truncation_strategies(padding=padding, truncation=truncation, max_length=max_length, pad_to_multiple_of=pad_to_multiple_of, verbose=verbose, **kwargs)
        return self._encode_plus(text=text, boxes=boxes, text_pair=text_pair, word_labels=word_labels, add_special_tokens=add_special_tokens, padding_strategy=padding_strategy, truncation_strategy=truncation_strategy, max_length=max_length, stride=stride, pad_to_multiple_of=pad_to_multiple_of, padding_side=padding_side, return_tensors=return_tensors, return_token_type_ids=return_token_type_ids, return_attention_mask=return_attention_mask, return_overflowing_tokens=return_overflowing_tokens, return_special_tokens_mask=return_special_tokens_mask, return_offsets_mapping=return_offsets_mapping, return_length=return_length, verbose=verbose, **kwargs)

    def _batch_encode_plus(self, batch_text_or_text_pairs: Union[List[TextInput], List[TextInputPair], List[PreTokenizedInput]], is_pair: Optional[bool]=None, boxes: Optional[List[List[List[int]]]]=None, word_labels: Optional[List[List[int]]]=None, add_special_tokens: bool=True, padding_strategy: PaddingStrategy=PaddingStrategy.DO_NOT_PAD, truncation_strategy: TruncationStrategy=TruncationStrategy.DO_NOT_TRUNCATE, max_length: Optional[int]=None, stride: int=0, pad_to_multiple_of: Optional[int]=None, padding_side: Optional[str]=None, return_tensors: Optional[str]=None, return_token_type_ids: Optional[bool]=None, return_attention_mask: Optional[bool]=None, return_overflowing_tokens: bool=False, return_special_tokens_mask: bool=False, return_offsets_mapping: bool=False, return_length: bool=False, verbose: bool=True) -> BatchEncoding:
        if not isinstance(batch_text_or_text_pairs, list):
            raise TypeError(f'batch_text_or_text_pairs has to be a list (got {type(batch_text_or_text_pairs)})')
        self.set_truncation_and_padding(padding_strategy=padding_strategy, truncation_strategy=truncation_strategy, max_length=max_length, stride=stride, pad_to_multiple_of=pad_to_multiple_of, padding_side=padding_side)
        if is_pair:
            batch_text_or_text_pairs = [(text.split(), text_pair) for text, text_pair in batch_text_or_text_pairs]
        encodings = self._tokenizer.encode_batch(batch_text_or_text_pairs, add_special_tokens=add_special_tokens, is_pretokenized=True)
        tokens_and_encodings = [self._convert_encoding(encoding=encoding, return_token_type_ids=return_token_type_ids, return_attention_mask=return_attention_mask, return_overflowing_tokens=return_overflowing_tokens, return_special_tokens_mask=return_special_tokens_mask, return_offsets_mapping=True if word_labels is not None else return_offsets_mapping, return_length=return_length, verbose=verbose) for encoding in encodings]
        sanitized_tokens = {}
        for key in tokens_and_encodings[0][0].keys():
            stack = [e for item, _ in tokens_and_encodings for e in item[key]]
            sanitized_tokens[key] = stack
        sanitized_encodings = [e for _, item in tokens_and_encodings for e in item]
        if return_overflowing_tokens:
            overflow_to_sample_mapping = []
            for i, (toks, _) in enumerate(tokens_and_encodings):
                overflow_to_sample_mapping += [i] * len(toks['input_ids'])
            sanitized_tokens['overflow_to_sample_mapping'] = overflow_to_sample_mapping
        for input_ids in sanitized_tokens['input_ids']:
            self._eventual_warn_about_too_long_sequence(input_ids, max_length, verbose)
        token_boxes = []
        for batch_index in range(len(sanitized_tokens['input_ids'])):
            if return_overflowing_tokens:
                original_index = sanitized_tokens['overflow_to_sample_mapping'][batch_index]
            else:
                original_index = batch_index
            token_boxes_example = []
            for id, sequence_id, word_id in zip(sanitized_tokens['input_ids'][batch_index], sanitized_encodings[batch_index].sequence_ids, sanitized_encodings[batch_index].word_ids):
                if word_id is not None:
                    if is_pair and sequence_id == 0:
                        token_boxes_example.append(self.pad_token_box)
                    else:
                        token_boxes_example.append(boxes[original_index][word_id])
                elif id == self.cls_token_id:
                    token_boxes_example.append(self.cls_token_box)
                elif id == self.sep_token_id:
                    token_boxes_example.append(self.sep_token_box)
                elif id == self.pad_token_id:
                    token_boxes_example.append(self.pad_token_box)
                else:
                    raise ValueError('Id not recognized')
            token_boxes.append(token_boxes_example)
        sanitized_tokens['bbox'] = token_boxes
        if word_labels is not None:
            labels = []
            for batch_index in range(len(sanitized_tokens['input_ids'])):
                if return_overflowing_tokens:
                    original_index = sanitized_tokens['overflow_to_sample_mapping'][batch_index]
                else:
                    original_index = batch_index
                labels_example = []
                previous_token_empty = False
                for id, offset, word_id in zip(sanitized_tokens['input_ids'][batch_index], sanitized_tokens['offset_mapping'][batch_index], sanitized_encodings[batch_index].word_ids):
                    if word_id is not None:
                        if self.only_label_first_subword:
                            if offset[0] == 0 and (not previous_token_empty):
                                labels_example.append(word_labels[original_index][word_id])
                            else:
                                labels_example.append(self.pad_token_label)
                            if offset == (0, 0):
                                previous_token_empty = True
                            else:
                                previous_token_empty = False
                        else:
                            labels_example.append(word_labels[original_index][word_id])
                    else:
                        labels_example.append(self.pad_token_label)
                labels.append(labels_example)
            sanitized_tokens['labels'] = labels
            if not return_offsets_mapping:
                del sanitized_tokens['offset_mapping']
        return BatchEncoding(sanitized_tokens, sanitized_encodings, tensor_type=return_tensors)

    def _encode_plus(self, text: Union[TextInput, PreTokenizedInput], text_pair: Optional[PreTokenizedInput]=None, boxes: Optional[List[List[int]]]=None, word_labels: Optional[List[int]]=None, add_special_tokens: bool=True, padding_strategy: PaddingStrategy=PaddingStrategy.DO_NOT_PAD, truncation_strategy: TruncationStrategy=TruncationStrategy.DO_NOT_TRUNCATE, max_length: Optional[int]=None, stride: int=0, pad_to_multiple_of: Optional[int]=None, padding_side: Optional[str]=None, return_tensors: Optional[bool]=None, return_token_type_ids: Optional[bool]=None, return_attention_mask: Optional[bool]=None, return_overflowing_tokens: bool=False, return_special_tokens_mask: bool=False, return_offsets_mapping: bool=False, return_length: bool=False, verbose: bool=True, **kwargs) -> BatchEncoding:
        batched_input = [(text, text_pair)] if text_pair else [text]
        batched_boxes = [boxes]
        batched_word_labels = [word_labels] if word_labels is not None else None
        batched_output = self._batch_encode_plus(batched_input, is_pair=bool(text_pair is not None), boxes=batched_boxes, word_labels=batched_word_labels, add_special_tokens=add_special_tokens, padding_strategy=padding_strategy, truncation_strategy=truncation_strategy, max_length=max_length, stride=stride, pad_to_multiple_of=pad_to_multiple_of, padding_side=padding_side, return_tensors=return_tensors, return_token_type_ids=return_token_type_ids, return_attention_mask=return_attention_mask, return_overflowing_tokens=return_overflowing_tokens, return_special_tokens_mask=return_special_tokens_mask, return_offsets_mapping=return_offsets_mapping, return_length=return_length, verbose=verbose, **kwargs)
        if return_tensors is None and (not return_overflowing_tokens):
            batched_output = BatchEncoding({key: value[0] if len(value) > 0 and isinstance(value[0], list) else value for key, value in batched_output.items()}, batched_output.encodings)
        self._eventual_warn_about_too_long_sequence(batched_output['input_ids'], max_length, verbose)
        return batched_output

    def _pad(self, encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding], max_length: Optional[int]=None, padding_strategy: PaddingStrategy=PaddingStrategy.DO_NOT_PAD, pad_to_multiple_of: Optional[int]=None, padding_side: Optional[str]=None, return_attention_mask: Optional[bool]=None) -> dict:
        """
        Pad encoded inputs (on left/right and up to predefined length or max length in the batch)

        Args:
            encoded_inputs:
                Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`).
            max_length: maximum length of the returned list and optionally padding length (see below).
                Will truncate by taking into account the special tokens.
            padding_strategy: PaddingStrategy to use for padding.

                - PaddingStrategy.LONGEST Pad to the longest sequence in the batch
                - PaddingStrategy.MAX_LENGTH: Pad to the max length (default)
                - PaddingStrategy.DO_NOT_PAD: Do not pad
                The tokenizer padding sides are defined in self.padding_side:

                    - 'left': pads on the left of the sequences
                    - 'right': pads on the right of the sequences
            pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value.
                This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability
                `>= 7.5` (Volta).
            padding_side:
                The side on which the model should have padding applied. Should be selected between ['right', 'left'].
                Default value is picked from the class attribute of the same name.
            return_attention_mask:
                (optional) Set to False to avoid returning attention mask (default: set to model specifics)
        """
        if return_attention_mask is None:
            return_attention_mask = 'attention_mask' in self.model_input_names
        required_input = encoded_inputs[self.model_input_names[0]]
        if padding_strategy == PaddingStrategy.LONGEST:
            max_length = len(required_input)
        if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
            max_length = (max_length // pad_to_multiple_of + 1) * pad_to_multiple_of
        needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len(required_input) != max_length
        if return_attention_mask and 'attention_mask' not in encoded_inputs:
            encoded_inputs['attention_mask'] = [1] * len(required_input)
        if needs_to_be_padded:
            difference = max_length - len(required_input)
            padding_side = padding_side if padding_side is not None else self.padding_side
            if padding_side == 'right':
                if return_attention_mask:
                    encoded_inputs['attention_mask'] = encoded_inputs['attention_mask'] + [0] * difference
                if 'token_type_ids' in encoded_inputs:
                    encoded_inputs['token_type_ids'] = encoded_inputs['token_type_ids'] + [self.pad_token_type_id] * difference
                if 'bbox' in encoded_inputs:
                    encoded_inputs['bbox'] = encoded_inputs['bbox'] + [self.pad_token_box] * difference
                if 'labels' in encoded_inputs:
                    encoded_inputs['labels'] = encoded_inputs['labels'] + [self.pad_token_label] * difference
                if 'special_tokens_mask' in encoded_inputs:
                    encoded_inputs['special_tokens_mask'] = encoded_inputs['special_tokens_mask'] + [1] * difference
                encoded_inputs[self.model_input_names[0]] = required_input + [self.pad_token_id] * difference
            elif padding_side == 'left':
                if return_attention_mask:
                    encoded_inputs['attention_mask'] = [0] * difference + encoded_inputs['attention_mask']
                if 'token_type_ids' in encoded_inputs:
                    encoded_inputs['token_type_ids'] = [self.pad_token_type_id] * difference + encoded_inputs['token_type_ids']
                if 'bbox' in encoded_inputs:
                    encoded_inputs['bbox'] = [self.pad_token_box] * difference + encoded_inputs['bbox']
                if 'labels' in encoded_inputs:
                    encoded_inputs['labels'] = [self.pad_token_label] * difference + encoded_inputs['labels']
                if 'special_tokens_mask' in encoded_inputs:
                    encoded_inputs['special_tokens_mask'] = [1] * difference + encoded_inputs['special_tokens_mask']
                encoded_inputs[self.model_input_names[0]] = [self.pad_token_id] * difference + required_input
            else:
                raise ValueError('Invalid padding strategy:' + str(padding_side))
        return encoded_inputs

    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str]=None) -> Tuple[str]:
        files = self._tokenizer.model.save(save_directory, name=filename_prefix)
        return tuple(files)

    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
        output = [self.bos_token_id] + token_ids_0 + [self.eos_token_id]
        if token_ids_1 is None:
            return output
        return output + [self.eos_token_id] + token_ids_1 + [self.eos_token_id]

    def create_token_type_ids_from_sequences(self, token_ids_0: List[int], token_ids_1: Optional[List[int]]=None) -> List[int]:
        """
        Args:
        Create a mask from the two sequences passed to be used in a sequence-pair classification task. RoBERTa does not:
        make use of token type ids, therefore a list of zeros is returned.
            token_ids_0 (`List[int]`):
                List of IDs.
            token_ids_1 (`List[int]`, *optional*):
                Optional second list of IDs for sequence pairs.
        Returns:
            `List[int]`: List of zeros.
        """
        sep = [self.sep_token_id]
        cls = [self.cls_token_id]
        if token_ids_1 is None:
            return len(cls + token_ids_0 + sep) * [0]
        return len(cls + token_ids_0 + sep + sep + token_ids_1 + sep) * [0]

Construct a "fast" LayoutLMv3 tokenizer (backed by HuggingFace's *tokenizers* library). Based on BPE.

This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should
refer to this superclass for more information regarding those methods.

Args:
    vocab_file (`str`):
        Path to the vocabulary file.
    merges_file (`str`):
        Path to the merges file.
    errors (`str`, *optional*, defaults to `"replace"`):
        Paradigm to follow when decoding bytes to UTF-8. See
        [bytes.decode](https://docs.python.org/3/library/stdtypes.html#bytes.decode) for more information.
    bos_token (`str`, *optional*, defaults to `"<s>"`):
        The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.

        <Tip>

        When building a sequence using special tokens, this is not the token that is used for the beginning of
        sequence. The token used is the `cls_token`.

        </Tip>

    eos_token (`str`, *optional*, defaults to `"</s>"`):
        The end of sequence token.

        <Tip>

        When building a sequence using special tokens, this is not the token that is used for the end of sequence.
        The token used is the `sep_token`.

        </Tip>

    sep_token (`str`, *optional*, defaults to `"</s>"`):
        The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
        sequence classification or for a text and a question for question answering. It is also used as the last
        token of a sequence built with special tokens.
    cls_token (`str`, *optional*, defaults to `"<s>"`):
        The classifier token which is used when doing sequence classification (classification of the whole sequence
        instead of per-token classification). It is the first token of the sequence when built with special tokens.
    unk_token (`str`, *optional*, defaults to `"<unk>"`):
        The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
        token instead.
    pad_token (`str`, *optional*, defaults to `"<pad>"`):
        The token used for padding, for example when batching sequences of different lengths.
    mask_token (`str`, *optional*, defaults to `"<mask>"`):
        The token used for masking values. This is the token used when training this model with masked language
        modeling. This is the token which the model will try to predict.
    add_prefix_space (`bool`, *optional*, defaults to `False`):
        Whether or not to add an initial space to the input. This allows to treat the leading word just as any
        other word. (RoBERTa tokenizer detect beginning of words by the preceding space).
    trim_offsets (`bool`, *optional*, defaults to `True`):
        Whether the post processing step should trim offsets to avoid including whitespaces.
    cls_token_box (`List[int]`, *optional*, defaults to `[0, 0, 0, 0]`):
        The bounding box to use for the special [CLS] token.
    sep_token_box (`List[int]`, *optional*, defaults to `[0, 0, 0, 0]`):
        The bounding box to use for the special [SEP] token.
    pad_token_box (`List[int]`, *optional*, defaults to `[0, 0, 0, 0]`):
        The bounding box to use for the special [PAD] token.
    pad_token_label (`int`, *optional*, defaults to -100):
        The label to use for padding tokens. Defaults to -100, which is the `ignore_index` of PyTorch's
        CrossEntropyLoss.
    only_label_first_subword (`bool`, *optional*, defaults to `True`):
        Whether or not to only label the first subword, in case word labels are provided.


def from_text_vision_configs(cls, text_config: AlignTextConfig, vision_config: AlignVisionConfig, **kwargs):
    return cls(text_config=text_config.to_dict(), vision_config=vision_config.to_dict(), **kwargs)

Instantiate a [`AlignConfig`] (or a derived class) from align text model configuration and align vision model
configuration.

Returns:
    [`AlignConfig`]: An instance of a configuration object


def from_lengths(cls, lengths: Sequence[Union[Sequence[int], int]], num_row_partitions=None, dtype=dtypes.int64):
    if not isinstance(lengths, list):
        raise ValueError('lengths should be a list')
    for x in lengths:
        if not _is_int_or_tuple_of_ints(x):
            raise ValueError('element of lengths should be int or tuple of ints: instead %r' % (x,))
    if num_row_partitions is None:
        is_list = [not isinstance(x, int) for x in lengths]
        if any(is_list):
            num_row_partitions = len(is_list) - is_list[-1::-1].index(True) - 1
        else:
            num_row_partitions = 0
    if not isinstance(num_row_partitions, int):
        raise ValueError('num_row_partitions should be an int or None')
    if not lengths:
        if num_row_partitions > 0:
            raise ValueError('num_row_partitions==0 for a scalar shape')
        return DynamicRaggedShape([], [], dtype=dtype)
    if not num_row_partitions < len(lengths):
        raise ValueError('num_row_partitions should be less than `len(lengths)` if shape is not scalar.')
    if num_row_partitions > 0:
        row_partitions, nvals = _to_row_partitions_and_nvals_from_lengths(lengths[:num_row_partitions + 1])
        inner_shape = [nvals] + lengths[num_row_partitions + 1:]
        return DynamicRaggedShape(row_partitions, inner_shape, dtype=dtype)
    else:
        return DynamicRaggedShape([], lengths, dtype=dtype)

Creates a shape with the given lengths and num_row_partitions.

The lengths can either be a nonnegative int or a list of nonnegative ints.

If num_row_partitions is None, then the minimal num_row_partitions is used.

For example, [2, (3, 2)] is the shape of [[0, 0, 0], [0, 0]], and
[2, 2] is the shape of [[0, 0], [0, 0]]

This chooses the minimal num_row_partitions required (including zero).

The following table gives a few examples (where `RP(lengths)` is short
for `RowPartition.from_lengths(lengths)`):

For example:
from_lengths           | row_partitions            | inner_shape
---------------------- | --------------------------| -------------
[]                     | []                        | []
[2, (3, 2)]            | [RP([3, 2])]              | [5]
[2, 2]                 | []                        | [2, 2]
[2, (3, 2), 7]         | [RP([3, 2])]              | [5, 7]
[2, (2, 2), 3]         | [RP([2, 2])]              | [4, 3]
[2, 2, 3]              | []                        | [2, 2, 3]
[2, (2, 1), (2, 0, 3)] | [RP(2, 1), RP([2, 0, 3])] | [5]

If we want the row partitions to end with uniform row partitions, then
we can set num_row_partitions.

For example,
below URP(3, 12) is RowPartition.from_uniform_row_length(3, 12)

from_lengths   | num_row_partitions | row_partitions           | inner_shape
---------------| -------------------|--------------------------|------------
[2, (3, 2), 2] | 2                  | [RP([3, 2]), URP(2, 10)] | [10]
[2, 2]         | 1                  | [URP(2, 4)]              | [4]
[2, 2, 3]      | 0                  | []                       | [2, 2, 3]
[2, 2, 3]      | 1                  | [URP(2, 4)]              | [4, 3]
[2, 2, 3]      | 2                  | [URP(2, 4), URP(3, 12)]  | [12]

Representing the shapes from init():

from_lengths             | Tensor Example
------------------------ | ------------------------------
`[2, 3]`                 | `[[1, 2, 3], [4, 5, 6]]`
`[3, (2, 0, 3)]`         | `[[1, 2], [], [3, 4, 5]]`
`[2, (2, 1), 2]`         | `[[[1, 2], [3, 4]], [[5, 6]]]`
`[2, (2, 1), (2, 1, 2)]` | `[[[1, 2], [3]], [[4, 5]]]`

Args:
  lengths: the lengths of sublists along each axis.
  num_row_partitions: the num_row_partitions of the result or None
    indicating the minimum number of row_partitions.
  dtype: the dtype of the shape (tf.int32 or tf.int64).

Returns:
  a new DynamicRaggedShape


def update(lower: types.FloatTensor, upper: types.FloatTensor, estimate: types.FloatTensor, error: types.FloatTensor, tolerance: float, dtype: Optional[tf.DType]=None, name: Optional[str]=None) -> (types.FloatTensor, types.FloatTensor, types.FloatTensor):
    with tf.name_scope(name=name or 'adaptive_update'):
        lower = tf.convert_to_tensor(lower, dtype=dtype, name='lower')
        dtype = lower.dtype
        upper = tf.convert_to_tensor(upper, dtype=dtype, name='upper')
        relative_error = error / estimate
        condition = relative_error > tolerance
        num_bad_sub_intervals = tf.reduce_max(tf.math.count_nonzero(condition, axis=1, dtype=tf.int32), axis=0)
        indices = tf.math.top_k(relative_error, k=num_bad_sub_intervals, sorted=False).indices
        sum_all = tf.reduce_sum(estimate, axis=-1)
        sum_bad = tf.reduce_sum(tf.gather(estimate, indices, batch_dims=-1), axis=-1)
        sum_goods = sum_all - sum_bad
        filtered_lower = tf.gather(lower, indices, batch_dims=-1)
        filtered_upper = tf.gather(upper, indices, batch_dims=-1)
        mid_points = (filtered_lower + filtered_upper) / 2
        new_lower = tf.concat([filtered_lower, mid_points], axis=-1)
        new_upper = tf.concat([mid_points, filtered_upper], axis=-1)
        return (new_lower, new_upper, sum_goods)

Calculates new values for the limits for any adaptive quadrature.

Checks which intervals have estimated results that are within the provided
tolerance. The values for these intervals are added to the sum of good
estimations. The other intervals get divided in half.

#### Example
```python
  l = tf.constant([[[0.0], [1.0]]])
  u = tf.constant([[[1.0], [2.0]]])
  estimate = tf.constant([[[3.0], [4.0]]])
  err = tf.constant([[[0.01], [0.02]]])
  tol = 0.004
  update(l, u, estimate, err, tol)
  # tf.constant([[1.0, 1.5]]), tf.constant([[1.5, 2.0]]), tf.constant([3.0])
```

Args:
  lower: Represents the lower limits of integration. Must be a 2-dimensional
    tensor of shape `[batch_dim, n]` (where `n` is defined by the algorithm
    and represents the number of subintervals).
  upper: Same shape and dtype as `lower` representing the upper limits of
    intergation.
  estimate: Same shape and dtype as `lower` representing the integration
    results calculated with some quadrature method for the corresponding
    limits.
  error: Same shape and dtype as `lower` representing the estimated
    integration error for corresponding `estimate` values.
  tolerance: Represents the tolerance for the estimated error of the integral
    estimation, at which to stop further dividing the intervals.
  dtype: If supplied, the dtype for the `lower` and `upper`. Result will have
    the same dtype. Default value: None which maps to dtype of `lower`.
  name: The name to give to the ops created by this function. Default value:
    None which maps to 'adaptive_update'.

Returns:
  A tuple:
    * `Tensor` of shape `[batch_dim, new_n]`, containing values of the new
    lower limits,
    * `Tensor` of shape `[batch_dim, new_n]`, containing values of the new
    upper limits,
    * `Tensor` of shape `[batch_dim]`, containing sum values of the quadrature
    method results of the good intervals.


class Normalization(TFDataLayer):

    def __init__(self, axis=-1, mean=None, variance=None, invert=False, **kwargs):
        super().__init__(**kwargs)
        if axis is None:
            axis = ()
        elif isinstance(axis, int):
            axis = (axis,)
        else:
            axis = tuple(axis)
        self.axis = axis
        self.input_mean = mean
        self.input_variance = variance
        self.invert = invert
        self.supports_masking = True
        self._build_input_shape = None
        self.mean = None
        if (mean is not None) != (variance is not None):
            raise ValueError(f'When setting values directly, both `mean` and `variance` must be set. Received: mean={mean} and variance={variance}')

    def build(self, input_shape):
        if input_shape is None:
            return
        ndim = len(input_shape)
        self._build_input_shape = input_shape
        if any((a < -ndim or a >= ndim for a in self.axis)):
            raise ValueError(f'All `axis` values must be in the range [-ndim, ndim). Received inputs with ndim={ndim}, while axis={self.axis}')
        self._keep_axis = tuple(sorted([d if d >= 0 else d + ndim for d in self.axis]))
        for d in self._keep_axis:
            if input_shape[d] is None:
                raise ValueError(f'All `axis` values to be kept must have a known shape. Received axis={self.axis}, inputs.shape={input_shape}, with unknown axis at index {d}')
        self._reduce_axis = tuple((d for d in range(ndim) if d not in self._keep_axis))
        self._reduce_axis_mask = [0 if d in self._keep_axis else 1 for d in range(ndim)]
        self._broadcast_shape = [input_shape[d] if d in self._keep_axis else 1 for d in range(ndim)]
        mean_and_var_shape = tuple((input_shape[d] for d in self._keep_axis))
        self._mean_and_var_shape = mean_and_var_shape
        if self.input_mean is None:
            self.adapt_mean = self.add_weight(name='mean', shape=mean_and_var_shape, initializer='zeros', trainable=False)
            self.adapt_variance = self.add_weight(name='variance', shape=mean_and_var_shape, initializer='ones', trainable=False)
            self.count = self.add_weight(name='count', shape=(), dtype='int', initializer='zeros', trainable=False)
            self.built = True
            self.finalize_state()
        else:
            mean = ops.convert_to_tensor(self.input_mean)
            variance = ops.convert_to_tensor(self.input_variance)
            mean = ops.broadcast_to(mean, self._broadcast_shape)
            variance = ops.broadcast_to(variance, self._broadcast_shape)
            self.mean = ops.cast(mean, dtype=self.compute_dtype)
            self.variance = ops.cast(variance, dtype=self.compute_dtype)

    def adapt(self, data):
        """Computes the mean and variance of values in a dataset.

        Calling `adapt()` on a `Normalization` layer is an alternative to
        passing in `mean` and `variance` arguments during layer construction. A
        `Normalization` layer should always either be adapted over a dataset or
        passed `mean` and `variance`.

        During `adapt()`, the layer will compute a `mean` and `variance`
        separately for each position in each axis specified by the `axis`
        argument. To calculate a single `mean` and `variance` over the input
        data, simply pass `axis=None` to the layer.

        Arg:
            data: The data to train on. It can be passed either as a
                `tf.data.Dataset`, as a NumPy array, or as a backend-native
                eager tensor.
                If a dataset, *it must be batched*. Keras will assume that the
                data is batched, and if that assumption doesn't hold, the mean
                and variance may be incorrectly computed.
        """
        if isinstance(data, np.ndarray) or backend.is_tensor(data):
            input_shape = data.shape
        elif isinstance(data, tf.data.Dataset):
            input_shape = tuple(data.element_spec.shape)
            if len(input_shape) == 1:
                data = data.batch(128)
            input_shape = tuple(data.element_spec.shape)
        if not self.built:
            self.build(input_shape)
        else:
            for d in self._keep_axis:
                if input_shape[d] != self._build_input_shape[d]:
                    raise ValueError(f'The layer was built with input_shape={self._build_input_shape}, but adapt() is being called with data with an incompatible shape, data.shape={input_shape}')
        if isinstance(data, np.ndarray):
            total_mean = np.mean(data, axis=self._reduce_axis)
            total_var = np.var(data, axis=self._reduce_axis)
        elif backend.is_tensor(data):
            total_mean = ops.mean(data, axis=self._reduce_axis)
            total_var = ops.var(data, axis=self._reduce_axis)
        elif isinstance(data, tf.data.Dataset):
            total_mean = ops.zeros(self._mean_and_var_shape)
            total_var = ops.zeros(self._mean_and_var_shape)
            total_count = 0
            for batch in data:
                batch = backend.convert_to_tensor(batch, dtype=self.compute_dtype)
                batch_mean = ops.mean(batch, axis=self._reduce_axis)
                batch_var = ops.var(batch, axis=self._reduce_axis)
                if self._reduce_axis:
                    batch_reduce_shape = (batch.shape[d] for d in self._reduce_axis)
                    batch_count = math.prod(batch_reduce_shape)
                else:
                    batch_count = 1
                total_count += batch_count
                batch_weight = float(batch_count) / total_count
                existing_weight = 1.0 - batch_weight
                new_total_mean = total_mean * existing_weight + batch_mean * batch_weight
                total_var = (total_var + (total_mean - new_total_mean) ** 2) * existing_weight + (batch_var + (batch_mean - new_total_mean) ** 2) * batch_weight
                total_mean = new_total_mean
        else:
            raise NotImplementedError(f'Unsupported data type: {type(data)}')
        self.adapt_mean.assign(total_mean)
        self.adapt_variance.assign(total_var)
        self.finalize_state()

    def finalize_state(self):
        if self.input_mean is not None or not self.built:
            return
        self.mean = ops.reshape(self.adapt_mean, self._broadcast_shape)
        self.mean = ops.cast(self.mean, self.compute_dtype)
        self.variance = ops.reshape(self.adapt_variance, self._broadcast_shape)
        self.variance = ops.cast(self.variance, self.compute_dtype)

    def call(self, inputs):
        if self.mean is None:
            raise ValueError('You must call `.build(input_shape)` on the layer before using it.')
        inputs = self.backend.core.convert_to_tensor(inputs, dtype=self.compute_dtype)
        mean = self.convert_weight(self.mean)
        variance = self.convert_weight(self.variance)
        if self.invert:
            return self.backend.numpy.add(mean, self.backend.numpy.multiply(inputs, self.backend.numpy.maximum(self.backend.numpy.sqrt(variance), backend.epsilon())))
        else:
            return self.backend.numpy.divide(self.backend.numpy.subtract(inputs, mean), self.backend.numpy.maximum(self.backend.numpy.sqrt(variance), backend.epsilon()))

    def compute_output_shape(self, input_shape):
        return input_shape

    def get_config(self):
        config = super().get_config()
        config.update({'axis': self.axis, 'invert': self.invert, 'mean': np.array(self.input_mean).tolist(), 'variance': np.array(self.input_variance).tolist()})
        return config

    def load_own_variables(self, store):
        super().load_own_variables(store)
        self.finalize_state()

    def get_build_config(self):
        if self._build_input_shape:
            return {'input_shape': self._build_input_shape}

    def build_from_config(self, config):
        if config:
            self.build(config['input_shape'])

A preprocessing layer that normalizes continuous features.

This layer will shift and scale inputs into a distribution centered around
0 with standard deviation 1. It accomplishes this by precomputing the mean
and variance of the data, and calling `(input - mean) / sqrt(var)` at
runtime.

The mean and variance values for the layer must be either supplied on
construction or learned via `adapt()`. `adapt()` will compute the mean and
variance of the data and store them as the layer's weights. `adapt()` should
be called before `fit()`, `evaluate()`, or `predict()`.

Args:
    axis: Integer, tuple of integers, or None. The axis or axes that should
        have a separate mean and variance for each index in the shape.
        For example, if shape is `(None, 5)` and `axis=1`, the layer will
        track 5 separate mean and variance values for the last axis.
        If `axis` is set to `None`, the layer will normalize
        all elements in the input by a scalar mean and variance.
        When `-1`, the last axis of the input is assumed to be a
        feature dimension and is normalized per index.
        Note that in the specific case of batched scalar inputs where
        the only axis is the batch axis, the default will normalize
        each index in the batch separately.
        In this case, consider passing `axis=None`. Defaults to `-1`.
    mean: The mean value(s) to use during normalization. The passed value(s)
        will be broadcast to the shape of the kept axes above;
        if the value(s) cannot be broadcast, an error will be raised when
        this layer's `build()` method is called.
    variance: The variance value(s) to use during normalization. The passed
        value(s) will be broadcast to the shape of the kept axes above;
        if the value(s) cannot be broadcast, an error will be raised when
        this layer's `build()` method is called.
    invert: If `True`, this layer will apply the inverse transformation
        to its inputs: it would turn a normalized input back into its
        original form.

Examples:

Calculate a global mean and variance by analyzing the dataset in `adapt()`.

>>> adapt_data = np.array([1., 2., 3., 4., 5.], dtype='float32')
>>> input_data = np.array([1., 2., 3.], dtype='float32')
>>> layer = keras.layers.Normalization(axis=None)
>>> layer.adapt(adapt_data)
>>> layer(input_data)
array([-1.4142135, -0.70710677, 0.], dtype=float32)

Calculate a mean and variance for each index on the last axis.

>>> adapt_data = np.array([[0., 7., 4.],
...                        [2., 9., 6.],
...                        [0., 7., 4.],
...                        [2., 9., 6.]], dtype='float32')
>>> input_data = np.array([[0., 7., 4.]], dtype='float32')
>>> layer = keras.layers.Normalization(axis=-1)
>>> layer.adapt(adapt_data)
>>> layer(input_data)
array([-1., -1., -1.], dtype=float32)

Pass the mean and variance directly.

>>> input_data = np.array([[1.], [2.], [3.]], dtype='float32')
>>> layer = keras.layers.Normalization(mean=3., variance=2.)
>>> layer(input_data)
array([[-1.4142135 ],
       [-0.70710677],
       [ 0.        ]], dtype=float32)

Use the layer to de-normalize inputs (after adapting the layer).

>>> adapt_data = np.array([[0., 7., 4.],
...                        [2., 9., 6.],
...                        [0., 7., 4.],
...                        [2., 9., 6.]], dtype='float32')
>>> input_data = np.array([[1., 2., 3.]], dtype='float32')
>>> layer = keras.layers.Normalization(axis=-1, invert=True)
>>> layer.adapt(adapt_data)
>>> layer(input_data)
array([2., 10., 8.], dtype=float32)


def all(x, axis=None, keepdims=False):
    x = math_ops.cast(x, dtypes_module.bool)
    return math_ops.reduce_all(x, axis, keepdims)

Bitwise reduction (logical AND).

Args:
    x: Tensor or variable.
    axis: axis along which to perform the reduction.
    keepdims: whether the drop or broadcast the reduction axes.

Returns:
    A uint8 tensor (0s and 1s).


def read_graphs_event(self, offset):
    return debug_event_pb2.DebugEvent.FromString(self._get_reader(self._graphs_path).read(offset)[0])

Read a DebugEvent proto at a given offset from the .graphs file.

Args:
  offset: Offset to read the DebugEvent proto from.

Returns:
  A DebugEventProto.

Raises:
  `errors.DataLossError` if offset is at a wrong location.
  `IndexError` if offset is out of range of the file.


def get_profiles(self, cmd):
    if cmd not in self._views:
        raise ValueError('No autoprofiler for command: {}, was run'.format(cmd))
    return self._views[cmd]

Returns profiling results for each step at which `cmd` was run.

Args:
  cmd: string, profiling command used in an `add_auto_profiling` call.

Returns:
  dict[int: (MultiGraphNodeProto | GraphNodeProto)]. Keys are steps at which
  the profiling command was run. Values are the outputs of profiling.
  For "code" and "op" commands this will be a `MultiGraphNodeProto`, for
  "scope" and "graph" commands this will be a `GraphNodeProto.

Raises:
  ValueError: if `cmd` was never run (either because no session.run call was
  made or because there was no `add_auto_profiling` call with the specified
  `cmd`.


def _AssignTimestamps(pcoll, timestamp: Union[str, dict[str, str]], language: Optional[str]=None):
    timestamp_fn = _as_callable_for_pcoll(pcoll, timestamp, 'timestamp', language)
    T = TypeVar('T')
    return pcoll | beam.Map(lambda x: TimestampedValue(x, timestamp_fn(x))).with_input_types(T).with_output_types(T)

Assigns a new timestamp each element of its input.

This can be useful when reading records that have the timestamp embedded
in them, for example with various file types or other sources that by default
set all timestamps to the infinite past.

Note that the timestamp should only be set forward, as setting it backwards
may not cause it to hold back an already advanced watermark and the data
could become droppably late.

Args:
    timestamp: A field, callable, or expression giving the new timestamp.
    language: The language of the timestamp expression.
    error_handling: Whether and how to handle errors during timestamp
      evaluation.


class FlaubertSQuADHead(nn.Module):

    def __init__(self, config: FlaubertConfig):
        super().__init__()
        self.start_n_top = config.start_n_top
        self.end_n_top = config.end_n_top
        self.start_logits = FlaubertPoolerStartLogits(config)
        self.end_logits = FlaubertPoolerEndLogits(config)
        self.answer_class = FlaubertPoolerAnswerClass(config)

    @auto_docstring
    def forward(self, hidden_states: torch.FloatTensor, start_positions: Optional[torch.LongTensor]=None, end_positions: Optional[torch.LongTensor]=None, cls_index: Optional[torch.LongTensor]=None, is_impossible: Optional[torch.LongTensor]=None, p_mask: Optional[torch.FloatTensor]=None, return_dict: bool=False) -> Union[FlaubertSquadHeadOutput, Tuple[torch.FloatTensor]]:
        """
        hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`):
            Final hidden states of the model on the sequence tokens.
        start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            Positions of the first token for the labeled span.
        end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            Positions of the last token for the labeled span.
        cls_index (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            Position of the CLS token for each sentence in the batch. If `None`, takes the last token.
        is_impossible (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            Whether the question has a possible answer in the paragraph or not.
        p_mask (`torch.FloatTensor` of shape `(batch_size, seq_len)`, *optional*):
            Mask for tokens at invalid position, such as query and special symbols (PAD, SEP, CLS). 1.0 means token
            should be masked.
        """
        start_logits = self.start_logits(hidden_states, p_mask=p_mask)
        if start_positions is not None and end_positions is not None:
            for x in (start_positions, end_positions, cls_index, is_impossible):
                if x is not None and x.dim() > 1:
                    x.squeeze_(-1)
            end_logits = self.end_logits(hidden_states, start_positions=start_positions, p_mask=p_mask)
            loss_fct = CrossEntropyLoss()
            start_loss = loss_fct(start_logits, start_positions)
            end_loss = loss_fct(end_logits, end_positions)
            total_loss = (start_loss + end_loss) / 2
            if cls_index is not None and is_impossible is not None:
                cls_logits = self.answer_class(hidden_states, start_positions=start_positions, cls_index=cls_index)
                loss_fct_cls = nn.BCEWithLogitsLoss()
                cls_loss = loss_fct_cls(cls_logits, is_impossible)
                total_loss += cls_loss * 0.5
            return FlaubertSquadHeadOutput(loss=total_loss) if return_dict else (total_loss,)
        else:
            bsz, slen, hsz = hidden_states.size()
            start_log_probs = nn.functional.softmax(start_logits, dim=-1)
            start_top_log_probs, start_top_index = torch.topk(start_log_probs, self.start_n_top, dim=-1)
            start_top_index_exp = start_top_index.unsqueeze(-1).expand(-1, -1, hsz)
            start_states = torch.gather(hidden_states, -2, start_top_index_exp)
            start_states = start_states.unsqueeze(1).expand(-1, slen, -1, -1)
            hidden_states_expanded = hidden_states.unsqueeze(2).expand_as(start_states)
            p_mask = p_mask.unsqueeze(-1) if p_mask is not None else None
            end_logits = self.end_logits(hidden_states_expanded, start_states=start_states, p_mask=p_mask)
            end_log_probs = nn.functional.softmax(end_logits, dim=1)
            end_top_log_probs, end_top_index = torch.topk(end_log_probs, self.end_n_top, dim=1)
            end_top_log_probs = end_top_log_probs.view(-1, self.start_n_top * self.end_n_top)
            end_top_index = end_top_index.view(-1, self.start_n_top * self.end_n_top)
            start_states = torch.einsum('blh,bl->bh', hidden_states, start_log_probs)
            cls_logits = self.answer_class(hidden_states, start_states=start_states, cls_index=cls_index)
            if not return_dict:
                return (start_top_log_probs, start_top_index, end_top_log_probs, end_top_index, cls_logits)
            else:
                return FlaubertSquadHeadOutput(start_top_log_probs=start_top_log_probs, start_top_index=start_top_index, end_top_log_probs=end_top_log_probs, end_top_index=end_top_index, cls_logits=cls_logits)

A SQuAD head inspired by XLNet.

Args:
    config ([`FlaubertConfig`]):
        The config used by the model, will be used to grab the `hidden_size` of the model and the `layer_norm_eps`
        to use.


def compute_dtype(self):
    return self._compute_dtype

The compute dtype of this policy.

This is the dtype layers will do their computations in. Typically layers
output tensors with the compute dtype as well.

Note that even if the compute dtype is float16 or bfloat16, hardware devices
may not do individual adds, multiplies, and other fundamental operations in
float16 or bfloat16, but instead may do some of them in float32 for numeric
stability. The compute dtype is the dtype of the inputs and outputs of the
TensorFlow ops that the layer executes. Internally, many TensorFlow ops will
do certain internal calculations in float32 or some other device-internal
intermediate format with higher precision than float16/bfloat16, to increase
numeric stability.

For example, a `tf.keras.layers.Dense` layer, when run on a GPU with a
float16 compute dtype, will pass float16 inputs to `tf.linalg.matmul`. But,
`tf.linalg.matmul` will do use float32 intermediate math. The performance
benefit of float16 is still apparent, due to increased memory bandwidth and
the fact modern GPUs have specialized hardware for computing matmuls on
float16 inputs while still keeping intermediate computations in float32.

Returns:
  The compute dtype of this policy, as a string.


def incomplete_size(self, name=None):
    if name is None:
        name = '%s_BarrierIncompleteSize' % self._name
    return gen_data_flow_ops.barrier_incomplete_size(self._barrier_ref, name=name)

Compute the number of incomplete elements in the given barrier.

Args:
  name: A name for the operation (optional).

Returns:
  A single-element tensor containing the number of incomplete elements in
  the given barrier.


def trace(fun: Callable[[], Any], *, where: Optional[Callable[[base.HyperPrimitive], bool]]=None, require_hyper_name: bool=False, per_thread: bool=True) -> DynamicEvaluationContext:
    context = DynamicEvaluationContext(where=where, require_hyper_name=require_hyper_name, per_thread=per_thread)
    with context.collect():
        fun()
    return context

Trace the hyper primitives called within a function by executing it.

See examples in :class:`pyglove.hyper.DynamicEvaluationContext`.

Args:
  fun: Function in which the search space is defined.
  where: A callable object that decide whether a hyper primitive should be
    included when being instantiated under `collect`.
    If None, all hyper primitives under `collect` will be included.
  require_hyper_name: If True, all hyper primitives defined in this scope
    will need to carry their names, which is usually a good idea when the
    function that instantiates the hyper primtives need to be called multiple
    times.
  per_thread: If True, the context manager will be applied to current thread
    only. Otherwise, it will be applied on current process.

Returns:
    An DynamicEvaluationContext that can be passed to `pg.sample`.


def _scale_boxes(boxes, target_sizes):
    if isinstance(target_sizes, (list, tuple)):
        image_height = torch.tensor([i[0] for i in target_sizes])
        image_width = torch.tensor([i[1] for i in target_sizes])
    elif isinstance(target_sizes, torch.Tensor):
        image_height, image_width = target_sizes.unbind(1)
    else:
        raise ValueError('`target_sizes` must be a list, tuple or torch.Tensor')
    scale_factor = torch.stack([image_width, image_height, image_width, image_height], dim=1)
    scale_factor = scale_factor.unsqueeze(1).to(boxes.device)
    boxes = boxes * scale_factor
    return boxes

Scale batch of bounding boxes to the target sizes.

Args:
    boxes (`torch.Tensor` of shape `(batch_size, num_boxes, 4)`):
        Bounding boxes to scale. Each box is expected to be in (x1, y1, x2, y2) format.
    target_sizes (`List[Tuple[int, int]]` or `torch.Tensor` of shape `(batch_size, 2)`):
        Target sizes to scale the boxes to. Each target size is expected to be in (height, width) format.

Returns:
    `torch.Tensor` of shape `(batch_size, num_boxes, 4)`: Scaled bounding boxes.


def set_callback(self, property_name, callback):
    if property_name not in self._config:
        raise KeyError('%s is not a valid property name.' % property_name)
    if not callable(callback):
        raise TypeError('The callback object provided is not callable.')
    self._set_callbacks[property_name] = callback

Set a set-callback for given property.

Args:
  property_name: Name of the property.
  callback: The callback as a `callable` of signature:
      def cbk(config):
    where config is the config after it is set to the new value.
    The callback is invoked each time the set() method is called with the
    matching property_name.

Raises:
  KeyError: If property_name does not exist.
  TypeError: If `callback` is not callable.


def configure_collective_ops(self, collective_leader='', scoped_allocator_enabled_ops=('CollectiveReduce',), use_nccl_communication=False, device_filters=None):
    if self._collective_leader is not None:
        if self._collective_leader != collective_leader or self._collective_scoped_allocator_enabled_ops != scoped_allocator_enabled_ops or self._collective_use_nccl_communication != use_nccl_communication or (self._collective_device_filters != device_filters):
            raise ValueError('Collective ops are already configured.')
        else:
            return
    if self._context_handle is not None:
        raise RuntimeError('Collective ops must be configured at program startup')
    self._collective_leader = collective_leader
    self._collective_scoped_allocator_enabled_ops = scoped_allocator_enabled_ops
    self._collective_use_nccl_communication = use_nccl_communication
    self._collective_device_filters = device_filters

Configure collective ops.

  Collective group leader is necessary for collective ops to run, other
  configurations are mainly for the purpose of performance.

Args:
  collective_leader: a device string for collective leader, e.g.
    "/job:worker/replica:0/task:0"; empty string means local execution of
    collective ops.
  scoped_allocator_enabled_ops: a tuple or a list of op names for scoped
    allocator to run with.
  use_nccl_communication: whether to use nccl communication for collective
    ops.
  device_filters: a tuple or a list of device strings. If set, corresponding
    task can only see the devices filtered by these device filters.

Raises:
  RuntimeError: if this method is not called at program startup.


def dataclass(cls=None, *, kw_only=False, replace=True, repr=True, auto_cast=True, contextvars=True, allow_unfrozen=False):
    if cls is None:
        return functools.partial(dataclass, kw_only=kw_only, replace=replace, repr=repr, auto_cast=auto_cast, allow_unfrozen=allow_unfrozen)
    if kw_only:
        cls = _make_kw_only(cls)
    if repr:
        cls = add_repr(cls)
    if replace:
        cls = _add_replace(cls)
    if allow_unfrozen:
        cls = frozen_utils.add_unfrozen(cls)
    descriptor_fns = []
    if auto_cast:
        descriptor_fns.append(helpers.DescriptorInfo(annotation=cast_utils.AutoCast, descriptor_fn=cast_utils.make_auto_cast_descriptor))
    if contextvars:
        descriptor_fns.append(helpers.DescriptorInfo(annotation=context.ContextVar, descriptor_fn=context.make_contextvar_descriptor))
    cls = helpers.wrap_new(cls, descriptor_fns)
    return cls

Augment a dataclass with additional features.

`auto_cast`: Auto-convert init assignements to the annotated class.

```python
@edc.dataclass
class A:
  path: edc.AutoCast[epath.Path]
  some_enum: edc.AutoCast[MyEnum]
  x: edc.AutoCast[str]

a = A(
    path='/some/path',
    some_enum='A',
    x=123
)
# Fields annotated with `AutoCast` are automatically casted to their type
assert a.path == epath.Path('/some/path')
assert a.some_enum is MyEnum.A
assert a.x == '123'
```

`allow_unfrozen`: allow nested dataclass to be updated. This add two methods:

 * `.unfrozen()`: Create a lazy deep-copy of the current dataclass. Updates
   to nested attributes will be propagated to the top-level dataclass.
 * `.frozen()`: Returns the frozen dataclass, after it was mutated.

Example:

```python
old_x = X(y=Y(z=123))

x = old_x.unfrozen()
x.y.z = 456
x = x.frozen()

assert x == X(y=Y(z=123))  # Only new x is mutated
assert old_x == X(y=Y(z=456))  # Old x is not mutated
```

Note:

* Only the last `.frozen()` call resolve the dataclass by calling `.replace`
  recursivelly.
* Dataclass returned by `.unfrozen()` and nested attributes are not the
  original dataclass but proxy objects which track the mutations. As such,
  those object are not compatible with `isinstance()`, `jax.tree.map`,...
* Only the top-level dataclass need to be `allow_unfrozen=True`
* Avoid using `unfrozen` if 2 attributes of the dataclass point to the
  same nested dataclass. Updates on one attribute might not be reflected on
  the other.

  ```python
  y = Y(y=123)
  x = X(x0=y, x1=y)  # Same instance assigned twice in `x0` and `x1`
  x = x.unfrozen()
  x.x0.y = 456  # Changes in `x0` not reflected in `x1`
  x = x.frozen()

  assert x == X(x0=Y(y=456), x1=Y(y=123))
  ```

  This is because only attributes which are accessed are tracked, so `etils`
  do not know the object exist somewhere else in the attribute tree.

* After `.frozen()` has been called, any of the temporary sub-attribute
  become invalid:

  ```python
  a = a.unfrozen()
  y = a.y
  a = a.frozen()

  y.x  # Raise error (created between the unfrozen/frozen call)
  a.y.x  # Work
  ```

`contextvars`: Fields annotated as `edc.ContextVar` are wrapped in
a `contextvars.ContextVar`. Afterward each thread / asyncio coroutine will
have its own version of the fields (similarly to `threading.local`).

The contextvars are lazily initialized at first usage.

Example:

```python
@edc.dataclass
@dataclasses.dataclass
class Context:
  thread_id: edc.ContextVar[int] = dataclasses.field(
      default_factory=threading.get_native_id
  )
  stack: edc.ContextVar[list[str]] = dataclasses.field(default_factory=list)

# Global context object
context = Context(thread_id=0)

def worker():
  # Inside each thread, the worker use its own context
  assert context.thread_id != 0
  context.stack.append(1)

with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
  for _ in range(10):
    executor.submit(worker)
```

Args:
  cls: The dataclass to decorate
  kw_only: If True, make the dataclass `__init__` keyword-only.
  replace: If `True`, add a `.replace(` alias of `dataclasses.replace`.
  repr: If `True`, the class `__repr__` will return a pretty-printed `str`
    (one attribute per line)
  auto_cast: If `True`, fields annotated as `x: edc.AutoCast[Cls]` will be
    converted to `x: Cls = edc.field(validator=Cls)`.
  contextvars: It `True`, fields annotated as `x: edc.AutoCast[T]` are
    converted to `contextvars`. This allow to have a `threading.local`-like
    API for contextvars.
  allow_unfrozen: If `True`, add `.frozen`, `.unfrozen` methods.

Returns:
  Decorated class


def show_backref(target, max_depth=3):
    if objgraph is None:
        raise NotImplementedError('objgraph is not installed.')
    string_io = io.StringIO()
    objgraph.show_backrefs(target, max_depth=max_depth, output=string_io)
    graph = string_io.getvalue()
    string_io.close()
    return graph

Returns a dot graph of all the objects that are referencing the target.

A object referencing graph is useful to debug memory leak like circular
reference. objgraph provides a good visualization of the memory graph than
most python built-in utilities like gc.get_referrers(), which are not
human-readable sometimes.

The dot graph will be written to a string IO object, and can be rendered with
graphviz in operating system.
E.g. dot -Tpng {$dot_graph} -o output.png
Args:
  target: The target object for the memory graph.
  max_depth: The maximum depth of the graph. By default 3 layers of references
    are used. Increases this a lot may result in the graph growing too big.

Returns:
  A string that contains the object reference graph.
Raises:
  NotImplementedError: if objgraph is not installed.


def load_model_from_hdf5(filepath, custom_objects=None, compile=True):
    if h5py is None:
        raise ImportError('`load_model` requires h5py.')
    if not custom_objects:
        custom_objects = {}
    opened_new_file = not isinstance(filepath, h5py.File)
    if opened_new_file:
        f = h5py.File(filepath, mode='r')
    else:
        f = filepath
    model = None
    try:
        model_config = f.attrs.get('model_config')
        if model_config is None:
            raise ValueError('No model found in config file.')
        if hasattr(model_config, 'decode'):
            model_config = model_config.decode('utf-8')
        model_config = json_utils.decode(model_config)
        model = model_config_lib.model_from_config(model_config, custom_objects=custom_objects)
        load_weights_from_hdf5_group(f['model_weights'], model.layers)
        if compile:
            training_config = f.attrs.get('training_config')
            if hasattr(training_config, 'decode'):
                training_config = training_config.decode('utf-8')
            if training_config is None:
                logging.warning('No training configuration found in the save file, so the model was *not* compiled. Compile it manually.')
                return model
            training_config = json_utils.decode(training_config)
            model.compile(**saving_utils.compile_args_from_training_config(training_config, custom_objects), from_serialized=True)
            saving_utils.try_build_compiled_arguments(model)
            if 'optimizer_weights' in f:
                try:
                    model.optimizer._create_all_weights(model.trainable_variables)
                except (NotImplementedError, AttributeError):
                    logging.warning('Error when creating the weights of optimizer {}, making it impossible to restore the saved optimizer state. As a result, your model is starting with a freshly initialized optimizer.')
                optimizer_weight_values = load_optimizer_weights_from_hdf5_group(f)
                try:
                    model.optimizer.set_weights(optimizer_weight_values)
                except ValueError:
                    logging.warning('Error in loading the saved optimizer state. As a result, your model is starting with a freshly initialized optimizer.')
    finally:
        if opened_new_file:
            f.close()
    return model

Loads a model saved via `save_model_to_hdf5`.

Args:
    filepath: One of the following:
        - String, path to the saved model
        - `h5py.File` object from which to load the model
    custom_objects: Optional dictionary mapping names
        (strings) to custom classes or functions to be
        considered during deserialization.
    compile: Boolean, whether to compile the model
        after loading.

Returns:
    A Keras model instance. If an optimizer was found
    as part of the saved model, the model is already
    compiled. Otherwise, the model is uncompiled and
    a warning will be displayed. When `compile` is set
    to False, the compilation is omitted without any
    warning.

Raises:
    ImportError: if h5py is not available.
    ValueError: In case of an invalid savefile.


def get_special_tokens_mask(self, token_ids_0: List[int], token_ids_1: Optional[List[int]]=None, already_has_special_tokens: bool=False) -> List[int]:
    if already_has_special_tokens:
        return super().get_special_tokens_mask(token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True)
    if token_ids_1 is not None:
        return [1] + [0] * len(token_ids_0) + [1] + [0] * len(token_ids_1) + [1]
    return [1] + [0] * len(token_ids_0) + [1]

Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
special tokens using the tokenizer `prepare_for_model` method.

Args:
    token_ids_0 (`List[int]`):
        List of IDs.
    token_ids_1 (`List[int]`, *optional*):
        Optional second list of IDs for sequence pairs.
    already_has_special_tokens (`bool`, *optional*, defaults to `False`):
        Whether or not the token list is already formatted with special tokens for the model.

Returns:
    `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.


def take_along_axis(x, indices, axis=None):
    if any_symbolic_tensors((x, indices)):
        return TakeAlongAxis(axis=axis).symbolic_call(x, indices)
    return backend.numpy.take_along_axis(x, indices, axis=axis)

Select values from `x` at the 1-D `indices` along the given axis.

Args:
    x: Source tensor.
    indices: The indices of the values to extract.
    axis: The axis over which to select values. By default, the flattened
        input tensor is used.

Returns:
    The corresponding tensor of values.


def compute_gradient(x, x_shape, y, y_shape, x_init_value=None, delta=0.001, init_targets=None, extra_feed_dict=None):
    if extra_feed_dict is None:
        extra_feed_dict = {}
    if isinstance(x, list):
        return _compute_gradient_list(x, x_shape, y, y_shape, x_init_value, delta, init_targets, extra_feed_dict=extra_feed_dict)
    else:
        if init_targets is not None:
            assert isinstance(init_targets, (list, tuple))
            for init in init_targets:
                init.run()
        dx, dy = _compute_dx_and_dy(x, y, y_shape)
        ret = _compute_gradient(x, x_shape, dx, y, y_shape, dy, x_init_value, delta, extra_feed_dict=extra_feed_dict)
        return ret

Computes and returns the theoretical and numerical Jacobian.

If `x` or `y` is complex, the Jacobian will still be real but the
corresponding Jacobian dimension(s) will be twice as large.  This is required
even if both input and output is complex since TensorFlow graphs are not
necessarily holomorphic, and may have gradients not expressible as complex
numbers.  For example, if `x` is complex with shape `[m]` and `y` is complex
with shape `[n]`, each Jacobian `J` will have shape `[m * 2, n * 2]` with

    J[:m, :n] = d(Re y)/d(Re x)
    J[:m, n:] = d(Im y)/d(Re x)
    J[m:, :n] = d(Re y)/d(Im x)
    J[m:, n:] = d(Im y)/d(Im x)

Args:
  x: a tensor or list of tensors
  x_shape: the dimensions of x as a tuple or an array of ints. If x is a list,
  then this is the list of shapes.
  y: a tensor
  y_shape: the dimensions of y as a tuple or an array of ints.
  x_init_value: (optional) a numpy array of the same shape as "x"
    representing the initial value of x. If x is a list, this should be a list
    of numpy arrays.  If this is none, the function will pick a random tensor
    as the initial value.
  delta: (optional) the amount of perturbation.
  init_targets: list of targets to run to initialize model params.
  extra_feed_dict: dict that allows fixing specified tensor values
    during the Jacobian calculation.

Returns:
  Two 2-d numpy arrays representing the theoretical and numerical
  Jacobian for dy/dx. Each has "x_size" rows and "y_size" columns
  where "x_size" is the number of elements in x and "y_size" is the
  number of elements in y. If x is a list, returns a list of two numpy arrays.


def post_process_panoptic_sample(out_logits: np.ndarray, masks: np.ndarray, boxes: np.ndarray, processed_size: Tuple[int, int], target_size: Tuple[int, int], is_thing_map: Dict, threshold=0.85) -> Dict:
    scores, labels = score_labels_from_class_probabilities(out_logits)
    keep = (labels != out_logits.shape[-1] - 1) & (scores > threshold)
    cur_scores = scores[keep]
    cur_classes = labels[keep]
    cur_boxes = center_to_corners_format(boxes[keep])
    if len(cur_boxes) != len(cur_classes):
        raise ValueError('Not as many boxes as there are classes')
    cur_masks = masks[keep]
    cur_masks = resize(cur_masks[:, None], processed_size, resample=PILImageResampling.BILINEAR)
    cur_masks = safe_squeeze(cur_masks, 1)
    b, h, w = cur_masks.shape
    cur_masks = cur_masks.reshape(b, -1)
    stuff_equiv_classes = defaultdict(list)
    for k, label in enumerate(cur_classes):
        if not is_thing_map[label]:
            stuff_equiv_classes[label].append(k)
    seg_img = get_segmentation_image(cur_masks, processed_size, target_size, stuff_equiv_classes, deduplicate=True)
    area = get_mask_area(cur_masks, processed_size, n_classes=len(cur_scores))
    if cur_classes.size() > 0:
        filtered_small = np.array([a <= 4 for a in area], dtype=bool)
        while filtered_small.any():
            cur_masks = cur_masks[~filtered_small]
            cur_scores = cur_scores[~filtered_small]
            cur_classes = cur_classes[~filtered_small]
            seg_img = get_segmentation_image(cur_masks, (h, w), target_size, stuff_equiv_classes, deduplicate=True)
            area = get_mask_area(seg_img, target_size, n_classes=len(cur_scores))
            filtered_small = np.array([a <= 4 for a in area], dtype=bool)
    else:
        cur_classes = np.ones((1, 1), dtype=np.int64)
    segments_info = [{'id': i, 'isthing': is_thing_map[cat], 'category_id': int(cat), 'area': a} for i, (cat, a) in enumerate(zip(cur_classes, area))]
    del cur_classes
    with io.BytesIO() as out:
        PIL.Image.fromarray(seg_img).save(out, format='PNG')
        predictions = {'png_string': out.getvalue(), 'segments_info': segments_info}
    return predictions

Converts the output of [`DetrForSegmentation`] into panoptic segmentation predictions for a single sample.

Args:
    out_logits (`torch.Tensor`):
        The logits for this sample.
    masks (`torch.Tensor`):
        The predicted segmentation masks for this sample.
    boxes (`torch.Tensor`):
        The prediced bounding boxes for this sample. The boxes are in the normalized format `(center_x, center_y,
        width, height)` and values between `[0, 1]`, relative to the size the image (disregarding padding).
    processed_size (`Tuple[int, int]`):
        The processed size of the image `(height, width)`, as returned by the preprocessing step i.e. the size
        after data augmentation but before batching.
    target_size (`Tuple[int, int]`):
        The target size of the image, `(height, width)` corresponding to the requested final size of the
        prediction.
    is_thing_map (`Dict`):
        A dictionary mapping class indices to a boolean value indicating whether the class is a thing or not.
    threshold (`float`, *optional*, defaults to 0.85):
        The threshold used to binarize the segmentation masks.


def forward(self, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.FloatTensor]=None, head_mask: Optional[torch.FloatTensor]=None, inputs_embeds: Optional[torch.FloatTensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[Tuple[torch.FloatTensor], BaseModelOutput]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    encoder_outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, head_mask=head_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    return encoder_outputs

input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
    Indices of input sequence tokens in the vocabulary. UMT5 is a model with relative position embeddings so you
    should be able to pad the inputs on both the right and the left.

    Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
    [`PreTrainedTokenizer.__call__`] for detail.

    To know more on how to prepare `input_ids` for pretraining take a look a [UMT5 Training](./umt5#training).

Example:

```python
>>> from transformers import AutoTokenizer, UMT5EncoderModel

>>> tokenizer = AutoTokenizer.from_pretrained("google/umt5-small")
>>> model = UMT5EncoderModel.from_pretrained("google/umt5-small")
>>> input_ids = tokenizer(
...     "Studies have been shown that owning a dog is good for you", return_tensors="pt"
... ).input_ids  # Batch size 1
>>> outputs = model(input_ids=input_ids)
>>> last_hidden_states = outputs.last_hidden_state
```


def build_inputs_with_special_tokens(self, token_ids_0: List[int], token_ids_1: Optional[List[int]]=None) -> List[int]:
    if token_ids_1 is None:
        return [self.sep_token_id] + token_ids_0
    sep = [self.sep_token_id]
    return sep + token_ids_0 + sep + sep + token_ids_1

Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
adding special tokens. An XLM-RoBERTa sequence has the following format:

- single sequence: `<s> X </s>`
- pair of sequences: `<s> A </s></s> B </s>`

Args:
    token_ids_0 (`List[int]`):
        List of IDs to which the special tokens will be added.
    token_ids_1 (`List[int]`, *optional*):
        Optional second list of IDs for sequence pairs.

Returns:
    `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.


def streaming_restore(status, session=None):
    if context.executing_eagerly():
        return
    if session is None:
        session = get_session()
    if isinstance(status, NameBasedSaverStatus):
        raise NotImplementedError('Streaming restore not supported from name-based checkpoints when graph building. File a feature request if this limitation bothers you. As a workaround, consider either using tf.train.Checkpoint to load name-based checkpoints or enabling eager execution.')
    status.run_restore_ops(session=session)
    status._checkpoint.new_restore_ops_callback = lambda ops: session.run(ops, feed_dict=status._feed_dict)

When graph building, runs restore ops as soon as they come in.

Args:
  status: A _LoadStatus objects from an object-based saver's restore().
    Streaming restore from name-based checkpoints is not currently supported.
  session: A session to run new restore ops in.


def arccos(self: EventSetOrNode) -> EventSetOrNode:
    from temporian.core.operators.unary import arccos
    return arccos(self)

Calculates the inverse cosine of an [`EventSet`][temporian.EventSet]'s features.

Can only be used on floating point features.

Example:
    ```python
    >>> a = tp.event_set(
    ...     timestamps=[1, 2, 3],
    ...     features={"M": [1.0, 0, -1.0]},
    ... )
    >>> a.arccos()
    indexes: ...
            timestamps: [1. 2. 3.]
            'M': [0.     1.5708 3.1416]
    ...

    ```

Returns:
    EventSetOrNode with inverse cosine of input features.


class PegasusXEncoder(PegasusXPreTrainedModel):

    def __init__(self, config: PegasusXConfig, embed_tokens: Optional[nn.Embedding]=None):
        super().__init__(config)
        self.dropout = config.dropout
        self.layerdrop = config.encoder_layerdrop
        embed_dim = config.d_model
        padding_idx = config.pad_token_id
        self.max_source_positions = config.max_position_embeddings
        embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
        if embed_tokens is not None:
            self.embed_tokens = embed_tokens
        else:
            self.embed_tokens = PegasusXScaledWordEmbedding(config.vocab_size, embed_dim, padding_idx, embed_scale=embed_scale)
        self.embed_global = nn.Embedding(config.num_global_tokens, embed_dim)
        self.embed_positions = PegasusXSinusoidalPositionalEmbedding(embed_dim)
        self.layers = nn.ModuleList([PegasusXEncoderLayer(stagger_blocks_this_layer=i % 2 == 1 and config.stagger_local_blocks, config=config) for i in range(config.encoder_layers)])
        self.layer_norm = nn.LayerNorm(config.d_model)
        self.gradient_checkpointing = False
        self.post_init()

    def resize_position_embeddings(self, new_num_position_embeddings: int):
        """
        Resizes position embeddings matrix of the model if `new_num_position_embeddings !=
        config.max_position_embeddings`.

        Arguments:
            new_num_position_embeddings (`int`):
                The number of new position embeddings. If position embeddings are learned, increasing the size will add
                newly initialized vectors at the end, whereas reducing the size will remove vectors from the end. If
                position embeddings are not learned (*e.g.* sinusoidal position embeddings), increasing the size will
                add correct vectors at the end following the position encoding algorithm, whereas reducing the size
                will remove vectors from the end.
        """
        logger.info(f'Setting `config.max_position_embeddings={new_num_position_embeddings}`...')
        self.config.max_position_embeddings = new_num_position_embeddings
        self.embed_positions = PegasusXSinusoidalPositionalEmbedding(self.config.d_model)
        self.embed_positions.to(self.device)

    def get_position_embeddings(self) -> nn.Embedding:
        """
        Returns the position embeddings matrix
        """
        return self.embed_positions

    def forward(self, input_ids=None, attention_mask=None, inputs_embeds=None, output_attentions=None, output_hidden_states=None, return_dict=None):
        """
        Args:
            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
                Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you
                provide it.

                Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
                [`PreTrainedTokenizer.__call__`] for details.

                [What are input IDs?](../glossary#input-ids)
            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:

                - 1 for tokens that are **not masked**,
                - 0 for tokens that are **masked**.

                [What are attention masks?](../glossary#attention-mask)

            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
                Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
                This is useful if you want more control over how to convert `input_ids` indices into associated vectors
                than the model's internal embedding lookup matrix.
            output_attentions (`bool`, *optional*):
                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
                returned tensors for more detail.
            output_hidden_states (`bool`, *optional*):
                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
                for more detail.
            return_dict (`bool`, *optional*):
                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
        """
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time')
        elif input_ids is not None:
            self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
            input_shape = input_ids.size()
            input_ids = input_ids.view(-1, input_shape[-1])
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
        else:
            raise ValueError('You have to specify either input_ids or inputs_embeds')
        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)
        embed_pos = self.embed_positions(inputs_embeds)
        hidden_states = inputs_embeds + embed_pos
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        batch_size, seq_len, _ = hidden_states.shape
        if attention_mask is None:
            attention_mask = torch.ones(*input_shape, dtype=inputs_embeds.dtype, device=inputs_embeds.device)
        attention_mask = attention_mask.to(dtype=hidden_states.dtype)
        mask_min_value = torch.finfo(hidden_states.dtype).min
        inverted_mask = 1.0 - attention_mask
        attention_mask = inverted_mask.masked_fill(inverted_mask.to(torch.bool), mask_min_value)
        if seq_len % self.config.block_size != 0:
            pad_len = self.config.block_size - seq_len % self.config.block_size
            hidden_states = nn.functional.pad(hidden_states, pad=(0, 0, 0, pad_len), value=0)
            attention_mask = nn.functional.pad(attention_mask, pad=(0, pad_len), value=mask_min_value)
        global_hidden_states = self.embed_global(torch.arange(self.config.num_global_tokens, device=hidden_states.device)[None].expand(batch_size, -1))
        encoder_states = () if output_hidden_states else None
        all_attentions = () if output_attentions else None
        for idx, encoder_layer in enumerate(self.layers):
            if output_hidden_states:
                encoder_states = encoder_states + (hidden_states,)
            to_drop = False
            if self.training:
                dropout_probability = torch.rand([])
                if dropout_probability < self.layerdrop:
                    to_drop = True
            if to_drop:
                layer_outputs = (None, None)
            else:
                if self.gradient_checkpointing and self.training:
                    layer_outputs = self._gradient_checkpointing_func(encoder_layer.__call__, hidden_states, global_hidden_states, attention_mask, output_attentions)
                else:
                    layer_outputs = encoder_layer(hidden_states, global_hidden_states, attention_mask, output_attentions=output_attentions)
                hidden_states = layer_outputs[0]
                global_hidden_states = layer_outputs[1]
            if output_attentions:
                all_attentions = all_attentions + (layer_outputs[2],)
        hidden_states = hidden_states[:, :seq_len]
        hidden_states = self.layer_norm(hidden_states)
        if output_hidden_states:
            encoder_states = encoder_states + ((hidden_states, global_hidden_states),)
        if not return_dict:
            return tuple((v for v in [hidden_states, encoder_states, all_attentions] if v is not None))
        return BaseModelOutput(last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions)

Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
[`PegasusXEncoderLayer`].

Args:
    config: PegasusXConfig
    embed_tokens (nn.Embedding): output embedding


def masks_to_boxes(masks: np.ndarray) -> np.ndarray:
    if masks.size == 0:
        return np.zeros((0, 4))
    h, w = masks.shape[-2:]
    y = np.arange(0, h, dtype=np.float32)
    x = np.arange(0, w, dtype=np.float32)
    y, x = np.meshgrid(y, x, indexing='ij')
    x_mask = masks * np.expand_dims(x, axis=0)
    x_max = x_mask.reshape(x_mask.shape[0], -1).max(-1)
    x = np.ma.array(x_mask, mask=~np.array(masks, dtype=bool))
    x_min = x.filled(fill_value=100000000.0)
    x_min = x_min.reshape(x_min.shape[0], -1).min(-1)
    y_mask = masks * np.expand_dims(y, axis=0)
    y_max = y_mask.reshape(x_mask.shape[0], -1).max(-1)
    y = np.ma.array(y_mask, mask=~np.array(masks, dtype=bool))
    y_min = y.filled(fill_value=100000000.0)
    y_min = y_min.reshape(y_min.shape[0], -1).min(-1)
    return np.stack([x_min, y_min, x_max, y_max], 1)

Compute the bounding boxes around the provided panoptic segmentation masks.

Args:
    masks: masks in format `[number_masks, height, width]` where N is the number of masks

Returns:
    boxes: bounding boxes in format `[number_masks, 4]` in xyxy format


def batch_sizes_for_worker(global_batch_size, num_workers, num_replicas_per_worker, worker_index):
    num_subbatches = num_workers * num_replicas_per_worker
    offset = worker_index * num_replicas_per_worker
    const_value = tensor_util.constant_value(global_batch_size)
    if const_value is not None:
        global_batch_size = const_value
    floor = global_batch_size // num_subbatches
    num_ceil = global_batch_size - num_subbatches * floor
    if const_value is not None:
        worker_0 = [floor + 1] * num_ceil + [floor] * (num_subbatches - num_ceil)
        return ops.convert_to_tensor(worker_0[offset:] + worker_0[:offset], dtype=dtypes.int64, name='batch_sizes')
    worker_0 = array_ops.ones(num_subbatches, dtype=dtypes.int64)
    worker_0 = floor * worker_0 + array_ops.concat([array_ops.ones(num_ceil, dtype=dtypes.int64), array_ops.zeros(num_subbatches - num_ceil, dtype=dtypes.int64)], axis=0)
    return array_ops.concat([worker_0[offset:], worker_0[:offset]], axis=0)

Determines how to rebatch a dataset for the given worker.

Given the global batch size, number of workers, number of replicas per worker,
and worker index, returns the correct batch sizes for rebatching a dataset
on worker `worker_index` of `num_workers`, such that each global step (across
all workers and replicas) will consume global_batch_size elements. The
returned value should be passed as the `batch_sizes` input parameter to
`tf.data.experimental.rebatch()`. The returned batch sizes meet the following
constraints:

Let G = global_batch_size, W = num_workers, R = num_replicas_per_worker
(A) for any worker, len(batch_sizes) = W * R
(B) for any worker, sum(batch_sizes) == G
(C) for any global step (i.e. R iterations on each worker), the sum of batches
    consumed by replicas across all workers is G.
(D) any two batch sizes of any two replicas differs by at most one.

For example, suppose we have G = 7, W = 2, R = 2, and suppose we have two
files which each contain 7 elements:

```python
# WORKER 0
batch_sizes_0 = batch_sizes_for_worker(global_batch_size=global_batch_size,
                                       num_workers=2,
                                       num_replicas_per_worker=2,
                                       worker_index=0)
print(batch_sizes_0)
>> [2, 2, 2, 1]

dataset_0 = tf.data.Dataset.from_tensor_slices(["file_a", "file_b"])
dataset_0 = dataset_0.shard(num_shards, index=0)
dataset_0 = dataset_0.batch(7)
dataset_0 = dataset_0.apply(tf.data.experimental.rebatch(batch_sizes_0))
for elem in dataset_0:
  print(elem)
>> [[A0, A1], [A2, A3], [A4, A5], [A6]]

# WORKER 1
batch_sizes_1 = batch_sizes_for_worker(global_batch_size=global_batch_size,
                                       num_workers=2,
                                       num_replicas_per_worker=2,
                                       worker_index=1)
print(batch_sizes_1)
>> [2, 1, 2, 2]

dataset_1 = tf.data.Dataset.from_tensor_slices(["file_a", "file_b"])
dataset_1 = dataset_1.shard(num_shards, index=1)
dataset_1 = dataset_1.batch(7)
dataset_1 = dataset_1.apply(tf.data.experimental.rebatch(batch_sizes_1))
for elem in dataset_1:
  print(elem)
>> [[B0, B1], [B2], [B3, B4], [B5, B6]]
```

The above example will produce the following elements:

Step 1:
  Worker 0 Replica 0: [A0, A1]
  Worker 0 Replica 1: [A2, A3]
  Worker 1 Replica 0: [B0, B1]
  Worker 1 Replica 1: [B2]
Total batch size = 7

Step 2:
  Worker 0 Replica 0: [A4, A5]
  Worker 0 Replica 1: [A6]
  Worker 1 Replica 0: [B3, B4]
  Worker 1 Replica 1: [B5, B6]
Total batch size = 7

Args:
  global_batch_size: A `tf.int64` scalar, representing the global batch size.
  num_workers: An integer representing the number of workers the dataset will
    be distributed across.
  num_replicas_per_worker: An integer representing the number of replicas per
    worker. All workers are assumed to have the same number of replicas.
  worker_index: An integer index of the worker to be rebatched.

Returns:
  A `tf.int64` vector, representing the batch sizes to rebatch the dataset
  into.


def sigmoid_focal_loss(inputs: torch.Tensor, targets: torch.Tensor, num_boxes: int, alpha: float=0.25, gamma: float=2):
    prob = inputs.sigmoid()
    ce_loss = nn.functional.binary_cross_entropy_with_logits(inputs, targets, reduction='none')
    p_t = prob * targets + (1 - prob) * (1 - targets)
    loss = ce_loss * (1 - p_t) ** gamma
    if alpha >= 0:
        alpha_t = alpha * targets + (1 - alpha) * (1 - targets)
        loss = alpha_t * loss
    return loss.sum() / num_boxes

Loss used in RetinaNet for dense detection: https://huggingface.co/papers/1708.02002.

Args:
    inputs (`torch.FloatTensor` of arbitrary shape):
        The predictions for each example.
    targets (`torch.FloatTensor` with the same shape as `inputs`)
        A tensor storing the binary classification label for each element in the `inputs` (0 for the negative class
        and 1 for the positive class).
    num_boxes (`int`):
        The total number of boxes in the batch.
    alpha (`float`, *optional*, defaults to 0.25):
        Optional weighting factor in the range (0,1) to balance positive vs. negative examples.
    gamma (`int`, *optional*, defaults to 2):
        Exponent of the modulating factor (1 - p_t) to balance easy vs hard examples.

Returns:
    Loss tensor


def pad(boxes, top, left, height=None, width=None, bounding_box_format='xyxy'):
    if bounding_box_format != 'xyxy':
        raise NotImplementedError
    box_utils = BoundingBox()
    if backend_utils.in_tf_graph():
        box_utils.backend.set_backend('tensorflow')
    outputs = box_utils.pad(boxes, top, left)
    box_utils.backend.reset()
    return outputs

Pads bounding boxes by adding top and left offsets.

This function adds padding to the bounding boxes by increasing the 'top'
and 'left' coordinates by the specified amounts. The method assume the
input bounding_box_format is `xyxy`.

Args:
    boxes: Bounding boxes to pad. Shape `(N, 4)` or `(batch, N, 4)`.
    top: Vertical padding to add.
    left: Horizontal padding to add.
    height: Image height. Defaults to None.
    width: Image width. Defaults to None.
    bounding_box_format: The format of the input bounding boxes. Defaults to
        `"xyxy"`.

Returns:
    Padded bounding boxes in the original format.