class Wav2Vec2ProcessorWithLM(ProcessorMixin):
    feature_extractor_class = 'AutoFeatureExtractor'
    tokenizer_class = 'Wav2Vec2CTCTokenizer'

    def __init__(self, feature_extractor: 'FeatureExtractionMixin', tokenizer: 'PreTrainedTokenizerBase', decoder: 'BeamSearchDecoderCTC'):
        from pyctcdecode import BeamSearchDecoderCTC
        super().__init__(feature_extractor, tokenizer)
        if not isinstance(decoder, BeamSearchDecoderCTC):
            raise TypeError(f'`decoder` has to be of type {BeamSearchDecoderCTC.__class__}, but is {type(decoder)}')
        if feature_extractor.__class__.__name__ not in ['Wav2Vec2FeatureExtractor', 'SeamlessM4TFeatureExtractor']:
            raise ValueError(f'`feature_extractor` has to be of type `Wav2Vec2FeatureExtractor` or `SeamlessM4TFeatureExtractor`, but is {type(feature_extractor)}')
        missing_decoder_tokens = self.get_missing_alphabet_tokens(decoder, tokenizer)
        if len(missing_decoder_tokens) > 0:
            raise ValueError(f"The tokens {missing_decoder_tokens} are defined in the tokenizer's vocabulary, but not in the decoder's alphabet. Make sure to include {missing_decoder_tokens} in the decoder's alphabet.")
        self.decoder = decoder
        self.current_processor = self.feature_extractor
        self._in_target_context_manager = False

    def save_pretrained(self, save_directory):
        super().save_pretrained(save_directory)
        self.decoder.save_to_dir(save_directory)

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
        """
        Instantiate a [`Wav2Vec2ProcessorWithLM`] from a pretrained Wav2Vec2 processor.

        <Tip>

        This class method is simply calling the feature extractor's
        [`~feature_extraction_utils.FeatureExtractionMixin.from_pretrained`], Wav2Vec2CTCTokenizer's
        [`~tokenization_utils_base.PreTrainedTokenizerBase.from_pretrained`], and
        [`pyctcdecode.BeamSearchDecoderCTC.load_from_hf_hub`].

        Please refer to the docstrings of the methods above for more information.

        </Tip>

        Args:
            pretrained_model_name_or_path (`str` or `os.PathLike`):
                This can be either:

                - a string, the *model id* of a pretrained feature_extractor hosted inside a model repo on
                  huggingface.co.
                - a path to a *directory* containing a feature extractor file saved using the
                  [`~SequenceFeatureExtractor.save_pretrained`] method, e.g., `./my_model_directory/`.
                - a path or url to a saved feature extractor JSON *file*, e.g.,
                  `./my_model_directory/preprocessor_config.json`.
            **kwargs
                Additional keyword arguments passed along to both [`SequenceFeatureExtractor`] and
                [`PreTrainedTokenizer`]
        """
        requires_backends(cls, 'pyctcdecode')
        from pyctcdecode import BeamSearchDecoderCTC
        feature_extractor, tokenizer = super()._get_arguments_from_pretrained(pretrained_model_name_or_path, **kwargs)
        if os.path.isdir(pretrained_model_name_or_path) or os.path.isfile(pretrained_model_name_or_path):
            unigram_encoding = kwargs.get('unigram_encoding', 'utf-8')
            decoder = BeamSearchDecoderCTC.load_from_dir(pretrained_model_name_or_path, unigram_encoding)
        else:
            kwargs.pop('_from_auto', None)
            kwargs.pop('trust_remote_code', None)
            language_model_filenames = os.path.join(BeamSearchDecoderCTC._LANGUAGE_MODEL_SERIALIZED_DIRECTORY, '*')
            alphabet_filename = BeamSearchDecoderCTC._ALPHABET_SERIALIZED_FILENAME
            allow_patterns = [language_model_filenames, alphabet_filename]
            decoder = BeamSearchDecoderCTC.load_from_hf_hub(pretrained_model_name_or_path, allow_patterns=allow_patterns, **kwargs)
        for attribute in ['alpha', 'beta', 'unk_score_offset', 'score_boundary']:
            value = kwargs.pop(attribute, None)
            if value is not None:
                cls._set_language_model_attribute(decoder, attribute, value)
        missing_decoder_tokens = cls.get_missing_alphabet_tokens(decoder, tokenizer)
        if len(missing_decoder_tokens) > 0:
            raise ValueError(f"The tokens {missing_decoder_tokens} are defined in the tokenizer's vocabulary, but not in the decoder's alphabet. Make sure to include {missing_decoder_tokens} in the decoder's alphabet.")
        return cls(feature_extractor=feature_extractor, tokenizer=tokenizer, decoder=decoder)

    @staticmethod
    def _set_language_model_attribute(decoder: 'BeamSearchDecoderCTC', attribute: str, value: float):
        setattr(decoder.model_container[decoder._model_key], attribute, value)

    @property
    def language_model(self):
        return self.decoder.model_container[self.decoder._model_key]

    @staticmethod
    def get_missing_alphabet_tokens(decoder, tokenizer):
        from pyctcdecode.alphabet import BLANK_TOKEN_PTN, UNK_TOKEN, UNK_TOKEN_PTN
        tokenizer_vocab_list = list(tokenizer.get_vocab().keys())
        for i, token in enumerate(tokenizer_vocab_list):
            if BLANK_TOKEN_PTN.match(token):
                tokenizer_vocab_list[i] = ''
            if token == tokenizer.word_delimiter_token:
                tokenizer_vocab_list[i] = ' '
            if UNK_TOKEN_PTN.match(token):
                tokenizer_vocab_list[i] = UNK_TOKEN
        missing_tokens = set(tokenizer_vocab_list) - set(decoder._alphabet.labels)
        return missing_tokens

    def __call__(self, *args, **kwargs):
        """
        When used in normal mode, this method forwards all its arguments to the feature extractor's
        [`~FeatureExtractionMixin.__call__`] and returns its output. If used in the context
        [`~Wav2Vec2ProcessorWithLM.as_target_processor`] this method forwards all its arguments to
        Wav2Vec2CTCTokenizer's [`~Wav2Vec2CTCTokenizer.__call__`]. Please refer to the docstring of the above two
        methods for more information.
        """
        if self._in_target_context_manager:
            return self.current_processor(*args, **kwargs)
        if 'raw_speech' in kwargs:
            warnings.warn('Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead.')
            audio = kwargs.pop('raw_speech')
        else:
            audio = kwargs.pop('audio', None)
        sampling_rate = kwargs.pop('sampling_rate', None)
        text = kwargs.pop('text', None)
        if len(args) > 0:
            audio = args[0]
            args = args[1:]
        if audio is None and text is None:
            raise ValueError('You need to specify either an `audio` or `text` input to process.')
        if audio is not None:
            inputs = self.feature_extractor(audio, *args, sampling_rate=sampling_rate, **kwargs)
        if text is not None:
            encodings = self.tokenizer(text, **kwargs)
        if text is None:
            return inputs
        elif audio is None:
            return encodings
        else:
            inputs['labels'] = encodings['input_ids']
            return inputs

    def pad(self, *args, **kwargs):
        """
        When used in normal mode, this method forwards all its arguments to the feature extractor's
        [`~FeatureExtractionMixin.pad`] and returns its output. If used in the context
        [`~Wav2Vec2ProcessorWithLM.as_target_processor`] this method forwards all its arguments to
        Wav2Vec2CTCTokenizer's [`~Wav2Vec2CTCTokenizer.pad`]. Please refer to the docstring of the above two methods
        for more information.
        """
        if self._in_target_context_manager:
            return self.current_processor.pad(*args, **kwargs)
        input_features = kwargs.pop('input_features', None)
        labels = kwargs.pop('labels', None)
        if len(args) > 0:
            input_features = args[0]
            args = args[1:]
        if input_features is not None:
            input_features = self.feature_extractor.pad(input_features, *args, **kwargs)
        if labels is not None:
            labels = self.tokenizer.pad(labels, **kwargs)
        if labels is None:
            return input_features
        elif input_features is None:
            return labels
        else:
            input_features['labels'] = labels['input_ids']
            return input_features

    def batch_decode(self, logits: np.ndarray, pool: Optional[Pool]=None, num_processes: Optional[int]=None, beam_width: Optional[int]=None, beam_prune_logp: Optional[float]=None, token_min_logp: Optional[float]=None, hotwords: Optional[Iterable[str]]=None, hotword_weight: Optional[float]=None, alpha: Optional[float]=None, beta: Optional[float]=None, unk_score_offset: Optional[float]=None, lm_score_boundary: Optional[bool]=None, output_word_offsets: bool=False, n_best: int=1):
        """
        Batch decode output logits to audio transcription with language model support.

        <Tip>

        This function makes use of Python's multiprocessing. Currently, multiprocessing is available only on Unix
        systems (see this [issue](https://github.com/kensho-technologies/pyctcdecode/issues/65)).

        If you are decoding multiple batches, consider creating a `Pool` and passing it to `batch_decode`. Otherwise,
        `batch_decode` will be very slow since it will create a fresh `Pool` for each call. See usage example below.

        </Tip>

        Args:
            logits (`np.ndarray`):
                The logits output vector of the model representing the log probabilities for each token.
            pool (`multiprocessing.Pool`, *optional*):
                An optional user-managed pool. If not set, one will be automatically created and closed. The pool
                should be instantiated *after* `Wav2Vec2ProcessorWithLM`. Otherwise, the LM won't be available to the
                pool's sub-processes.

                <Tip>

                Currently, only pools created with a 'fork' context can be used. If a 'spawn' pool is passed, it will
                be ignored and sequential decoding will be used instead.

                </Tip>

            num_processes (`int`, *optional*):
                If `pool` is not set, number of processes on which the function should be parallelized over. Defaults
                to the number of available CPUs.
            beam_width (`int`, *optional*):
                Maximum number of beams at each step in decoding. Defaults to pyctcdecode's DEFAULT_BEAM_WIDTH.
            beam_prune_logp (`int`, *optional*):
                Beams that are much worse than best beam will be pruned Defaults to pyctcdecode's DEFAULT_PRUNE_LOGP.
            token_min_logp (`int`, *optional*):
                Tokens below this logp are skipped unless they are argmax of frame Defaults to pyctcdecode's
                DEFAULT_MIN_TOKEN_LOGP.
            hotwords (`List[str]`, *optional*):
                List of words with extra importance, can be OOV for LM
            hotword_weight (`int`, *optional*):
                Weight factor for hotword importance Defaults to pyctcdecode's DEFAULT_HOTWORD_WEIGHT.
            alpha (`float`, *optional*):
                Weight for language model during shallow fusion
            beta (`float`, *optional*):
                Weight for length score adjustment of during scoring
            unk_score_offset (`float`, *optional*):
                Amount of log score offset for unknown tokens
            lm_score_boundary (`bool`, *optional*):
                Whether to have kenlm respect boundaries when scoring
            output_word_offsets (`bool`, *optional*, defaults to `False`):
                Whether or not to output word offsets. Word offsets can be used in combination with the sampling rate
                and model downsampling rate to compute the time-stamps of transcribed words.
            n_best (`int`, *optional*, defaults to `1`):
                Number of best hypotheses to return. If `n_best` is greater than 1, the returned `text` will be a list
                of lists of strings, `logit_score` will be a list of lists of floats, and `lm_score` will be a list of
                lists of floats, where the length of the outer list will correspond to the batch size and the length of
                the inner list will correspond to the number of returned hypotheses . The value should be >= 1.

                <Tip>

                Please take a look at the Example of [`~Wav2Vec2ProcessorWithLM.decode`] to better understand how to
                make use of `output_word_offsets`. [`~Wav2Vec2ProcessorWithLM.batch_decode`] works the same way with
                batched output.

                </Tip>

        Returns:
            [`~models.wav2vec2.Wav2Vec2DecoderWithLMOutput`].

        Example:
            See [Decoding multiple audios](#decoding-multiple-audios).
        """
        from pyctcdecode.constants import DEFAULT_BEAM_WIDTH, DEFAULT_HOTWORD_WEIGHT, DEFAULT_MIN_TOKEN_LOGP, DEFAULT_PRUNE_LOGP
        beam_width = beam_width if beam_width is not None else DEFAULT_BEAM_WIDTH
        beam_prune_logp = beam_prune_logp if beam_prune_logp is not None else DEFAULT_PRUNE_LOGP
        token_min_logp = token_min_logp if token_min_logp is not None else DEFAULT_MIN_TOKEN_LOGP
        hotword_weight = hotword_weight if hotword_weight is not None else DEFAULT_HOTWORD_WEIGHT
        self.decoder.reset_params(alpha=alpha, beta=beta, unk_score_offset=unk_score_offset, lm_score_boundary=lm_score_boundary)
        logits_list = [array[(array != -100.0).all(axis=-1)] for array in logits]
        if pool is None:
            default_context = get_start_method()
            if default_context == 'fork':
                cm = pool = get_context().Pool(num_processes)
            else:
                logger.warning('Parallel batch decoding is not currently supported in this platform. Falling back to sequential decoding.')
                cm = nullcontext()
        else:
            cm = nullcontext()
            if num_processes is not None:
                logger.warning('Parameter `num_process` was passed, but it will be ignored since `pool` was also specified.')
        with cm:
            decoded_beams = self.decoder.decode_beams_batch(pool=pool, logits_list=logits_list, beam_width=beam_width, beam_prune_logp=beam_prune_logp, token_min_logp=token_min_logp, hotwords=hotwords, hotword_weight=hotword_weight)
        batch_texts, logit_scores, lm_scores, word_offsets = ([], [], [], [])
        for d in decoded_beams:
            batch_texts.append([beam[0] for beam in d])
            logit_scores.append([beam[-2] for beam in d])
            lm_scores.append([beam[-1] for beam in d])
            word_offsets.append([[{'word': word, 'start_offset': start_offset, 'end_offset': end_offset} for word, (start_offset, end_offset) in beam[1]] for beam in d])
        word_offsets = word_offsets if output_word_offsets else None
        if n_best == 1:
            return Wav2Vec2DecoderWithLMOutput(text=[hyps[0] for hyps in batch_texts], logit_score=[hyps[0] for hyps in logit_scores], lm_score=[hyps[0] for hyps in lm_scores], word_offsets=[hyps[0] for hyps in word_offsets] if word_offsets is not None else None)
        else:
            return Wav2Vec2DecoderWithLMOutput(text=[hyps[:n_best] for hyps in batch_texts], logit_score=[hyps[:n_best] for hyps in logit_scores], lm_score=[hyps[:n_best] for hyps in lm_scores], word_offsets=[hyps[:n_best] for hyps in word_offsets] if word_offsets is not None else None)

    def decode(self, logits: np.ndarray, beam_width: Optional[int]=None, beam_prune_logp: Optional[float]=None, token_min_logp: Optional[float]=None, hotwords: Optional[Iterable[str]]=None, hotword_weight: Optional[float]=None, alpha: Optional[float]=None, beta: Optional[float]=None, unk_score_offset: Optional[float]=None, lm_score_boundary: Optional[bool]=None, output_word_offsets: bool=False, n_best: int=1):
        """
        Decode output logits to audio transcription with language model support.

        Args:
            logits (`np.ndarray`):
                The logits output vector of the model representing the log probabilities for each token.
            beam_width (`int`, *optional*):
                Maximum number of beams at each step in decoding. Defaults to pyctcdecode's DEFAULT_BEAM_WIDTH.
            beam_prune_logp (`int`, *optional*):
                A threshold to prune beams with log-probs less than best_beam_logp + beam_prune_logp. The value should
                be <= 0. Defaults to pyctcdecode's DEFAULT_PRUNE_LOGP.
            token_min_logp (`int`, *optional*):
                Tokens with log-probs below token_min_logp are skipped unless they are have the maximum log-prob for an
                utterance. Defaults to pyctcdecode's DEFAULT_MIN_TOKEN_LOGP.
            hotwords (`List[str]`, *optional*):
                List of words with extra importance which can be missing from the LM's vocabulary, e.g. ["huggingface"]
            hotword_weight (`int`, *optional*):
                Weight multiplier that boosts hotword scores. Defaults to pyctcdecode's DEFAULT_HOTWORD_WEIGHT.
            alpha (`float`, *optional*):
                Weight for language model during shallow fusion
            beta (`float`, *optional*):
                Weight for length score adjustment of during scoring
            unk_score_offset (`float`, *optional*):
                Amount of log score offset for unknown tokens
            lm_score_boundary (`bool`, *optional*):
                Whether to have kenlm respect boundaries when scoring
            output_word_offsets (`bool`, *optional*, defaults to `False`):
                Whether or not to output word offsets. Word offsets can be used in combination with the sampling rate
                and model downsampling rate to compute the time-stamps of transcribed words.
            n_best (`int`, *optional*, defaults to `1`):
                Number of best hypotheses to return. If `n_best` is greater than 1, the returned `text` will be a list
                of strings, `logit_score` will be a list of floats, and `lm_score` will be a list of floats, where the
                length of these lists will correspond to the number of returned hypotheses. The value should be >= 1.

                <Tip>

                Please take a look at the example below to better understand how to make use of `output_word_offsets`.

                </Tip>

        Returns:
            [`~models.wav2vec2.Wav2Vec2DecoderWithLMOutput`].

        Example:

        ```python
        >>> # Let's see how to retrieve time steps for a model
        >>> from transformers import AutoTokenizer, AutoProcessor, AutoModelForCTC
        >>> from datasets import load_dataset
        >>> import datasets
        >>> import torch

        >>> # import model, feature extractor, tokenizer
        >>> model = AutoModelForCTC.from_pretrained("patrickvonplaten/wav2vec2-base-100h-with-lm")
        >>> processor = AutoProcessor.from_pretrained("patrickvonplaten/wav2vec2-base-100h-with-lm")

        >>> # load first sample of English common_voice
        >>> dataset = load_dataset("mozilla-foundation/common_voice_11_0", "en", split="train", streaming=True, trust_remote_code=True)
        >>> dataset = dataset.cast_column("audio", datasets.Audio(sampling_rate=16_000))
        >>> dataset_iter = iter(dataset)
        >>> sample = next(dataset_iter)

        >>> # forward sample through model to get greedily predicted transcription ids
        >>> input_values = processor(sample["audio"]["array"], return_tensors="pt").input_values
        >>> with torch.no_grad():
        ...     logits = model(input_values).logits[0].cpu().numpy()

        >>> # retrieve word stamps (analogous commands for `output_char_offsets`)
        >>> outputs = processor.decode(logits, output_word_offsets=True)
        >>> # compute `time_offset` in seconds as product of downsampling ratio and sampling_rate
        >>> time_offset = model.config.inputs_to_logits_ratio / processor.feature_extractor.sampling_rate

        >>> word_offsets = [
        ...     {
        ...         "word": d["word"],
        ...         "start_time": round(d["start_offset"] * time_offset, 2),
        ...         "end_time": round(d["end_offset"] * time_offset, 2),
        ...     }
        ...     for d in outputs.word_offsets
        ... ]
        >>> # compare word offsets with audio `en_train_0/common_voice_en_19121553.mp3` online on the dataset viewer:
        >>> # https://huggingface.co/datasets/mozilla-foundation/common_voice_11_0/viewer/en
        >>> word_offsets[:4]
        [{'word': 'THE', 'start_time': 0.68, 'end_time': 0.78}, {'word': 'TRACK', 'start_time': 0.88, 'end_time': 1.1}, {'word': 'APPEARS', 'start_time': 1.18, 'end_time': 1.66}, {'word': 'ON', 'start_time': 1.86, 'end_time': 1.92}]
        ```"""
        from pyctcdecode.constants import DEFAULT_BEAM_WIDTH, DEFAULT_HOTWORD_WEIGHT, DEFAULT_MIN_TOKEN_LOGP, DEFAULT_PRUNE_LOGP
        beam_width = beam_width if beam_width is not None else DEFAULT_BEAM_WIDTH
        beam_prune_logp = beam_prune_logp if beam_prune_logp is not None else DEFAULT_PRUNE_LOGP
        token_min_logp = token_min_logp if token_min_logp is not None else DEFAULT_MIN_TOKEN_LOGP
        hotword_weight = hotword_weight if hotword_weight is not None else DEFAULT_HOTWORD_WEIGHT
        self.decoder.reset_params(alpha=alpha, beta=beta, unk_score_offset=unk_score_offset, lm_score_boundary=lm_score_boundary)
        decoded_beams = self.decoder.decode_beams(logits, beam_width=beam_width, beam_prune_logp=beam_prune_logp, token_min_logp=token_min_logp, hotwords=hotwords, hotword_weight=hotword_weight)
        word_offsets = None
        if output_word_offsets:
            word_offsets = [[{'word': word, 'start_offset': start_offset, 'end_offset': end_offset} for word, (start_offset, end_offset) in beam[2]] for beam in decoded_beams]
        logit_scores = [beam[-2] for beam in decoded_beams]
        lm_scores = [beam[-1] for beam in decoded_beams]
        hypotheses = [beam[0] for beam in decoded_beams]
        if n_best > len(decoded_beams):
            logger.info('N-best size is larger than the number of generated hypotheses, all hypotheses will be returned.')
        if n_best == 1:
            return Wav2Vec2DecoderWithLMOutput(text=hypotheses[0], logit_score=logit_scores[0], lm_score=lm_scores[0], word_offsets=word_offsets[0] if word_offsets is not None else None)
        else:
            return Wav2Vec2DecoderWithLMOutput(text=hypotheses[:n_best], logit_score=logit_scores[:n_best], lm_score=lm_scores[:n_best], word_offsets=word_offsets[:n_best] if word_offsets is not None else None)

    @contextmanager
    def as_target_processor(self):
        """
        Temporarily sets the processor for processing the target. Useful for encoding the labels when fine-tuning
        Wav2Vec2.
        """
        warnings.warn('`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your labels by using the argument `text` of the regular `__call__` method (either in the same call as your audio inputs, or in a separate call.')
        self._in_target_context_manager = True
        self.current_processor = self.tokenizer
        yield
        self.current_processor = self.feature_extractor
        self._in_target_context_manager = False

Constructs a Wav2Vec2 processor which wraps a Wav2Vec2 feature extractor, a Wav2Vec2 CTC tokenizer and a decoder
with language model support into a single processor for language model boosted speech recognition decoding.

Args:
    feature_extractor ([`Wav2Vec2FeatureExtractor`] or [`SeamlessM4TFeatureExtractor`]):
        An instance of [`Wav2Vec2FeatureExtractor`] or [`SeamlessM4TFeatureExtractor`]. The feature extractor is a required input.
    tokenizer ([`Wav2Vec2CTCTokenizer`]):
        An instance of [`Wav2Vec2CTCTokenizer`]. The tokenizer is a required input.
    decoder (`pyctcdecode.BeamSearchDecoderCTC`):
        An instance of [`pyctcdecode.BeamSearchDecoderCTC`]. The decoder is a required input.


def _resource_apply_sparse_duplicate_indices(self, grad, handle, indices):
    summed_grad, unique_indices = _deduplicate_indexed_slices(values=grad, indices=indices)
    return self._resource_apply_sparse(summed_grad, handle, unique_indices)

Add ops to apply sparse gradients to `handle`, with repeated indices.

Optimizers which override this method must deal with repeated indices. See
the docstring of `_apply_sparse_duplicate_indices` for details. By default
the correct behavior, to sum non-unique indices and their associated
gradients, is enforced by first pre-processing `grad` and `indices` and
passing them on to `_resource_apply_sparse`. Optimizers which deal correctly
with duplicate indices may instead override this method to avoid the
overhead of summing.

Args:
  grad: a `Tensor` representing the gradient for the affected indices.
  handle: a `Tensor` of dtype `resource` which points to the variable
   to be updated.
  indices: a `Tensor` of integral type representing the indices for
   which the gradient is nonzero. Indices may be repeated.

Returns:
  An `Operation` which updates the value of the variable.


def normalize_moments(counts, mean_ss, variance_ss, shift, name=None):
    with ops.name_scope(name, 'normalize', [counts, mean_ss, variance_ss, shift]):
        divisor = math_ops.reciprocal(counts, name='divisor')
        if shift is not None:
            shifted_mean = math_ops.multiply(mean_ss, divisor, name='shifted_mean')
            mean = math_ops.add(shifted_mean, shift, name='mean')
        else:
            shifted_mean = math_ops.multiply(mean_ss, divisor, name='mean')
            mean = shifted_mean
        variance = math_ops.subtract(math_ops.multiply(variance_ss, divisor), math_ops.square(shifted_mean), name='variance')
    return (mean, variance)

Calculate the mean and variance of based on the sufficient statistics.

Args:
  counts: A `Tensor` containing the total count of the data (one value).
  mean_ss: A `Tensor` containing the mean sufficient statistics: the (possibly
    shifted) sum of the elements to average over.
  variance_ss: A `Tensor` containing the variance sufficient statistics: the
    (possibly shifted) squared sum of the data to compute the variance over.
  shift: A `Tensor` containing the value by which the data is shifted for
    numerical stability, or `None` if no shift was performed.
  name: Name used to scope the operations that compute the moments.

Returns:
  Two `Tensor` objects: `mean` and `variance`.


def ParseCode(unformatted_source, filename='<unknown>'):
    if not unformatted_source.endswith(os.linesep):
        unformatted_source += os.linesep
    try:
        ast_tree = ast.parse(unformatted_source, filename)
        ast.fix_missing_locations(ast_tree)
        readline = StringIO(unformatted_source).readline
        tokens = tokenize.generate_tokens(readline)
    except Exception:
        raise
    logical_lines = _CreateLogicalLines(tokens)
    split_penalty_visitor.SplitPenalty(logical_lines).visit(ast_tree)
    return logical_lines

Parse a string of Python code into logical lines.

This provides an alternative entry point to YAPF.

Arguments:
  unformatted_source: (unicode) The code to format.
  filename: (unicode) The name of the file being reformatted.

Returns:
  A list of LogicalLines.

Raises:
  An exception is raised if there's an error during AST parsing.


def __init__(self, model_name: str, *, max_seq_length: Optional[int]=None, **kwargs):
    if not SentenceTransformer:
        raise ImportError('sentence-transformers is required to use HuggingfaceTextEmbeddings.Please install it with using `pip install sentence-transformers`.')
    super().__init__(type_adapter=create_rag_adapter(), **kwargs)
    self.model_name = model_name
    self.max_seq_length = max_seq_length
    self.model_class = SentenceTransformer

Utilizes huggingface SentenceTransformer embeddings for RAG pipeline.

Args:
    model_name: Name of the sentence-transformers model to use
    max_seq_length: Maximum sequence length for the model
    **kwargs: Additional arguments passed to
    :class:`~apache_beam.ml.transforms.base.EmbeddingsManager`
    constructor including ModelHandler arguments


def replace_with_aqlm_linear(model, quantization_config=None, linear_weights_not_to_quantize=None, current_key_name=None, has_been_replaced=False):
    if not is_aqlm_available():
        raise ValueError('AQLM is not available. Please install it with `pip install aqlm[cpu,gpu]`')
    if not is_accelerate_available():
        raise ValueError(f"AQLM requires Accelerate to be installed: `pip install 'accelerate>={ACCELERATE_MIN_VERSION}'`")
    if linear_weights_not_to_quantize is None:
        linear_weights_not_to_quantize = []
    from accelerate import init_empty_weights
    from aqlm import QuantizedLinear
    for name, module in model.named_children():
        if current_key_name is None:
            current_key_name = []
        current_key_name.append(name)
        if isinstance(module, nn.Linear):
            if '.'.join(current_key_name) + '.weight' not in linear_weights_not_to_quantize:
                with init_empty_weights():
                    in_features = module.in_features
                    out_features = module.out_features
                    model._modules[name] = QuantizedLinear(in_features, out_features, bias=module.bias is not None, in_group_size=quantization_config.in_group_size, out_group_size=quantization_config.out_group_size, num_codebooks=quantization_config.num_codebooks, nbits_per_codebook=quantization_config.nbits_per_codebook)
                    has_been_replaced = True
                    model._modules[name].source_cls = type(module)
                    model._modules[name].requires_grad_(False)
        if len(list(module.children())) > 0:
            _, has_been_replaced = replace_with_aqlm_linear(module, quantization_config=quantization_config, linear_weights_not_to_quantize=linear_weights_not_to_quantize, current_key_name=current_key_name, has_been_replaced=has_been_replaced)
        current_key_name.pop(-1)
    return (model, has_been_replaced)

Public method that recursively replaces the Linear layers of the given model with AQLM quantized layers.
`accelerate` is needed to use this method. Returns the converted model and a boolean that indicates if the
conversion has been successful or not.

Args:
    model (`torch.nn.Module`):
        The model to convert, can be any `torch.nn.Module` instance.
    quantization_config (`AqlmConfig`):
        The quantization config object that contains the quantization parameters.
    linear_weights_not_to_quantize (`list[str]`, *optional*):
        A list of nn.Linear weights to not convert. If a parameter path is in the list (e.g. `lm_head.weight`), the corresponding module will not be
        converted.
    current_key_name (`list`, *optional*):
        A list that contains the current key name. This is used for recursion and should not be passed by the user.
    has_been_replaced (`bool`, *optional*):
        A boolean that indicates if the conversion has been successful or not. This is used for recursion and
        should not be passed by the user.


Returns a Python dict of the object config.

A constraint config is a Python dictionary (JSON-serializable) that can
be used to reinstantiate the same object.

Returns:
    Python dict containing the configuration of the constraint object.


def convert_nested_bidirectional(weights):
    num_weights_per_layer = len(weights) // 2
    forward_weights = preprocess_weights_for_loading(layer.forward_layer, weights[:num_weights_per_layer], original_keras_version, original_backend)
    backward_weights = preprocess_weights_for_loading(layer.backward_layer, weights[num_weights_per_layer:], original_keras_version, original_backend)
    return forward_weights + backward_weights

Converts layers nested in `Bidirectional` wrapper.

This function uses `preprocess_weights_for_loading()` for converting
layers.

Args:
    weights: List of weights values (Numpy arrays).

Returns:
    A list of weights values (Numpy arrays).


class GotOcr2VisionConfig(PretrainedConfig):
    base_config_key = 'vision_config'

    def __init__(self, hidden_size=768, output_channels=256, num_hidden_layers=12, num_attention_heads=12, num_channels=3, image_size=1024, patch_size=16, hidden_act='gelu', layer_norm_eps=1e-06, attention_dropout=0.0, initializer_range=1e-10, qkv_bias=True, use_abs_pos=True, use_rel_pos=True, window_size=14, global_attn_indexes=[2, 5, 8, 11], mlp_dim=3072, **kwargs):
        super().__init__(**kwargs)
        self.hidden_size = hidden_size
        self.output_channels = output_channels
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.num_channels = num_channels
        self.image_size = image_size
        self.patch_size = patch_size
        self.hidden_act = hidden_act
        self.layer_norm_eps = layer_norm_eps
        self.attention_dropout = attention_dropout
        self.initializer_range = initializer_range
        self.qkv_bias = qkv_bias
        self.use_abs_pos = use_abs_pos
        self.use_rel_pos = use_rel_pos
        self.window_size = window_size
        self.global_attn_indexes = global_attn_indexes
        self.mlp_dim = mlp_dim

This is the configuration class to store the configuration of a [`GotOcr2VisionModel`]. It is used to instantiate a GOT_OCR2
vision encoder according to the specified arguments, defining the model architecture. Instantiating a configuration
defaults will yield a similar configuration to that of the SAM ViT-h
[facebook/sam-vit-huge](https://huggingface.co/facebook/sam-vit-huge) architecture.

Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
documentation from [`PretrainedConfig`] for more information.

Args:
    hidden_size (`int`, *optional*, defaults to 768):
        Dimensionality of the encoder layers and the pooler layer.
    output_channels (`int`, *optional*, defaults to 256):
        Dimensionality of the output channels in the Patch Encoder.
    num_hidden_layers (`int`, *optional*, defaults to 12):
        Number of hidden layers in the Transformer encoder.
    num_attention_heads (`int`, *optional*, defaults to 12):
        Number of attention heads for each attention layer in the Transformer encoder.
    num_channels (`int`, *optional*, defaults to 3):
        Number of channels in the input image.
    image_size (`int`, *optional*, defaults to 1024):
        Expected resolution. Target size of the resized input image.
    patch_size (`int`, *optional*, defaults to 16):
        Size of the patches to be extracted from the input image.
    hidden_act (`str`, *optional*, defaults to `"gelu"`):
        The non-linear activation function (function or string)
    layer_norm_eps (`float`, *optional*, defaults to 1e-06):
        The epsilon used by the layer normalization layers.
    attention_dropout (`float`, *optional*, defaults to 0.0):
        The dropout ratio for the attention probabilities.
    initializer_range (`float`, *optional*, defaults to 1e-10):
        The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
    qkv_bias (`bool`, *optional*, defaults to `True`):
        Whether to add a bias to query, key, value projections.
    use_abs_pos (`bool`, *optional*, defaults to `True`):
        Whether to use absolute position embedding.
    use_rel_pos (`bool`, *optional*, defaults to `True`):
        Whether to use relative position embedding.
    window_size (`int`, *optional*, defaults to 14):
        Window size for relative position.
    global_attn_indexes (`List[int]`, *optional*, defaults to `[2, 5, 8, 11]`):
        The indexes of the global attention layers.
    mlp_dim (`int`, *optional*, defaults to 3072):
        The dimensionality of the MLP layer in the Transformer encoder.


def _build_attention_equation(rank, attn_axes):
    target_notation = ''
    for i in range(rank):
        target_notation += _index_to_einsum_variable(i)
    batch_dims = tuple(np.delete(range(rank), attn_axes + (rank - 1,)))
    letter_offset = rank
    source_notation = ''
    for i in range(rank):
        if i in batch_dims or i == rank - 1:
            source_notation += target_notation[i]
        else:
            source_notation += _index_to_einsum_variable(letter_offset)
            letter_offset += 1
    product_notation = ''.join([target_notation[i] for i in batch_dims] + [target_notation[i] for i in attn_axes] + [source_notation[i] for i in attn_axes])
    dot_product_equation = '%s,%s->%s' % (source_notation, target_notation, product_notation)
    attn_scores_rank = len(product_notation)
    combine_equation = '%s,%s->%s' % (product_notation, source_notation, target_notation)
    return (dot_product_equation, combine_equation, attn_scores_rank)

Builds einsum equations for the attention computation.

Query, key, value inputs after projection are expected to have the shape as:
`(bs, <non-attention dims>, <attention dims>, num_heads, channels)`.
`bs` and `<non-attention dims>` are treated as `<batch dims>`.

The attention operations can be generalized:
1. Query-key dot product:
    (<batch dims>, <query attention dims>, num_heads, channels),
    (<batch dims>, <key attention dims>, num_heads, channels) ->
    (<batch dims>, num_heads, <query attention dims>, <key attention dims>)
2. Combination:
    (<batch dims>, num_heads, <query attention dims>, <key attention dims>),
    (<batch dims>, <value attention dims>, num_heads, channels) -> (<batch
    dims>, <query attention dims>, num_heads, channels)

Args:
    rank: Rank of query, key, value tensors.
    attn_axes: List/tuple of axes, `[-1, rank)`,
        that attention will be applied to.

Returns:
    Einsum equations.


def __init__(self, target='', graph=None, config=None):
    if not config:
        gpu_options = config_pb2.GPUOptions(allow_growth=True)
        config = config_pb2.ConfigProto(gpu_options=gpu_options)
    config.graph_options.place_pruned_graph = True
    super(InteractiveSession, self).__init__(target, graph, config)
    with InteractiveSession._count_lock:
        if InteractiveSession._active_session_count > 0:
            logging.error('An interactive session is already active. This can cause out-of-memory errors or some other unexpected errors (due to the unpredictable timing of garbage collection) in some cases. You must explicitly call `InteractiveSession.close()` to release resources held by the other session(s). Please use `tf.Session()` if you intend to productionize.')
        InteractiveSession._active_session_count += 1
    self._explicitly_closed = False
    self._default_session = self.as_default()
    self._default_session.enforce_nesting = False
    self._default_session.__enter__()
    self._explicit_graph = graph
    if self._explicit_graph is not None:
        self._default_graph = graph.as_default()
        self._default_graph.enforce_nesting = False
        self._default_graph.__enter__()

Creates a new interactive TensorFlow session.

If no `graph` argument is specified when constructing the session,
the default graph will be launched in the session. If you are
using more than one graph (created with `tf.Graph()`) in the same
process, you will have to use different sessions for each graph,
but each graph can be used in multiple sessions. In this case, it
is often clearer to pass the graph to be launched explicitly to
the session constructor.

Args:
  target: (Optional.) The execution engine to connect to. Defaults to using
    an in-process engine.
  graph: (Optional.) The `Graph` to be launched (described above).
  config: (Optional) `ConfigProto` proto used to configure the session.


def __call__(self, *args, **kwargs):
    self._assert_built_as_v1()
    if not hasattr(self, '_thread_local'):
        raise RuntimeError('You must call `super().__init__()` in the layer constructor.')
    if args:
        inputs = args[0]
        args = args[1:]
    elif self._call_fn_args[0] in kwargs:
        inputs = kwargs.pop(self._call_fn_args[0])
    else:
        raise ValueError('The first argument to `Layer.call` must always be passed.')
    call_context = base_layer_utils.call_context()
    input_list = nest.flatten(inputs)
    build_graph = tf_utils.are_all_symbolic_tensors(input_list)
    if any((isinstance(x, (np.ndarray, float, int)) for x in input_list)):

        def _convert_non_tensor(x):
            if isinstance(x, (np.ndarray, float, int)):
                return tensor_conversion.convert_to_tensor_v2_with_dispatch(x)
            return x
        inputs = nest.map_structure(_convert_non_tensor, inputs)
        input_list = nest.flatten(inputs)
    mask_arg_passed_by_framework = False
    input_masks = self._collect_input_masks(inputs, args, kwargs)
    if self._expects_mask_arg and input_masks is not None and (not self._call_arg_was_passed('mask', args, kwargs)):
        mask_arg_passed_by_framework = True
        kwargs['mask'] = input_masks
    training_value = None
    training_arg_passed_by_framework = False
    if self._call_arg_was_passed('training', args, kwargs):
        training_value = self._get_call_arg_value('training', args, kwargs)
        if not self._expects_training_arg:
            kwargs.pop('training')
    if training_value is None:
        if call_context.training is not None:
            training_value = call_context.training
        elif backend.global_learning_phase_is_set():
            training_value = backend.learning_phase()
        elif build_graph:
            with backend.get_graph().as_default():
                if base_layer_utils.is_in_keras_graph():
                    training_value = backend.learning_phase()
        if self._expects_training_arg and training_value is not None:
            if tensor_util.is_tf_type(training_value):
                training_value = math_ops.cast(training_value, dtypes.bool)
            else:
                training_value = bool(training_value)
            args, kwargs = self._set_call_arg_value('training', training_value, args, kwargs)
            training_arg_passed_by_framework = True
    if build_graph and base_layer_utils.needs_keras_history(inputs):
        base_layer_utils.create_keras_history(inputs)
    with call_context.enter(self, inputs, build_graph, training_value):
        if build_graph:
            input_spec.assert_input_compatibility(self.input_spec, inputs, self.name)
            graph = backend.get_graph()
            with graph.as_default(), backend.name_scope(self._name_scope()):
                self._maybe_build(inputs)
                cast_inputs = self._maybe_cast_inputs(inputs)
                if base_layer_utils.is_subclassed(self) and (not base_layer_utils.from_saved_model(self)):
                    call_fn = autograph.tf_convert(self.call, ag_ctx.control_status_ctx())
                else:
                    call_fn = self.call
                if not self.dynamic:
                    try:
                        with autocast_variable.enable_auto_cast_variables(self._compute_dtype_object):
                            outputs = call_fn(cast_inputs, *args, **kwargs)
                    except errors.OperatorNotAllowedInGraphError as e:
                        raise TypeError('You are attempting to use Python control flow in a layer that was not declared to be dynamic. Pass `dynamic=True` to the class constructor.\nEncountered error:\n"""\n' + str(e) + '\n"""')
                else:
                    outputs = self._symbolic_call(inputs)
                if outputs is None:
                    raise ValueError("A layer's `call` method should return a Tensor or a list of Tensors, not None (layer: " + self.name + ').')
                if base_layer_utils.have_all_keras_metadata(inputs):
                    if training_arg_passed_by_framework:
                        args, kwargs = self._set_call_arg_value('training', None, args, kwargs, pop_kwarg_if_none=True)
                    if mask_arg_passed_by_framework:
                        kwargs.pop('mask')
                    outputs = self._set_connectivity_metadata((inputs,) + args, kwargs, outputs)
                self._handle_activity_regularization(inputs, outputs)
                self._set_mask_metadata(inputs, outputs, input_masks)
                if hasattr(self, '_set_inputs') and (not self.inputs):
                    self._set_inputs(inputs, outputs)
        else:
            with backend.name_scope(self._name_scope()):
                self._maybe_build(inputs)
                cast_inputs = self._maybe_cast_inputs(inputs)
                with autocast_variable.enable_auto_cast_variables(self._compute_dtype_object):
                    outputs = self.call(cast_inputs, *args, **kwargs)
                self._handle_activity_regularization(inputs, outputs)
                self._set_mask_metadata(inputs, outputs, input_masks)
    return outputs

Wraps `call`, applying pre- and post-processing steps.

Args:
  *args: Positional arguments to be passed to `self.call`.
  **kwargs: Keyword arguments to be passed to `self.call`.

Returns:
  Output tensor(s).

Note:
  - The following optional keyword arguments are reserved for specific uses:
    * `training`: Boolean scalar tensor of Python boolean indicating
      whether the `call` is meant for training or inference.
    * `mask`: Boolean input mask.
  - If the layer's `call` method takes a `mask` argument (as some Keras
    layers do), its default value will be set to the mask generated
    for `inputs` by the previous layer (if `input` did come from
    a layer that generated a corresponding mask, i.e. if it came from
    a Keras layer with masking support.

Raises:
  ValueError: if the layer's `call` method returns None (an invalid value).
  RuntimeError: if `super().__init__()` was not called in the constructor.


def _shape_invariant_to_type_spec(self, shape):
    raise NotImplementedError(f'{type(self).__name__}._shape_invariant_to_type_spec')

Returns a TypeSpec given a shape invariant (used by `tf.while_loop`).

Args:
  shape: A `tf.TensorShape` object.  The shape invariant for this
    `CompositeTensor`, or `None` if a default shape invariant should be used
    (based on the value of this `CompositeTensor`).

Returns:
  A nested structure whose values are `tf.TensorShape` objects, specifying
  the shape invariants for the tensors that comprise this `CompositeTensor`.


def _process_single_batch(model, inputs, targets, output_loss_metrics=None, sample_weights=None, training=False):
    with backend.eager_learning_phase_scope(1 if training else 0), training_utils.RespectCompiledTrainableState(model):
        with GradientTape() as tape:
            outs, total_loss, output_losses, masks = _model_loss(model, inputs, targets, output_loss_metrics=output_loss_metrics, sample_weights=sample_weights, training=training)
            if isinstance(model.optimizer, loss_scale_optimizer.LossScaleOptimizer):
                scaled_total_loss = model.optimizer.get_scaled_loss(total_loss)
            else:
                scaled_total_loss = total_loss
        if training:
            trainable_weights = model.trainable_weights
            if trainable_weights:
                if hasattr(model, '_backwards'):
                    model._backwards(tape, scaled_total_loss)
                else:
                    grads = tape.gradient(scaled_total_loss, trainable_weights)
                    if isinstance(model.optimizer, loss_scale_optimizer.LossScaleOptimizer):
                        grads = model.optimizer.get_unscaled_gradients(grads)
                    model.optimizer.apply_gradients(zip(grads, trainable_weights))
            else:
                logging.warning('The list of trainable weights is empty. Make sure that you are not setting model.trainable to False before compiling the model.')
        return (outs, total_loss, output_losses, masks)

Calculate the loss and gradient for one input batch.

   The model weights are updated if training is set to True.

Args:
    model: Model whose loss has to be calculated.
    inputs: List of input arrays.
    targets: List of target arrays.
    output_loss_metrics: List of metrics that are used to aggregated output
      loss values.
    sample_weights: Optional list of sample weight arrays.
    training: The boolean represents if the weights of the model are updated.
            'fit' methods will set this to True while 'evaluate' methods will
            set this to False.

Returns:
    output of the model, total loss, the loss and the mask
    associated with each output.

Raises:
    ValueError: If the model has no loss to optimize.


def asset_or_nothing_price(*, volatilities: types.RealTensor, strikes: types.RealTensor, expiries: types.RealTensor, spots: types.RealTensor=None, forwards: types.RealTensor=None, discount_rates: types.RealTensor=None, dividend_rates: types.RealTensor=None, discount_factors: types.RealTensor=None, is_call_options: types.BoolTensor=None, is_normal_volatility: bool=False, dtype: tf.DType=None, name: str=None) -> types.RealTensor:
    if (spots is None) == (forwards is None):
        raise ValueError('Either spots or forwards must be supplied but not both.')
    if discount_rates is not None and discount_factors is not None:
        raise ValueError('At most one of discount_rates and discount_factors may be supplied')
    with tf.name_scope(name or 'asset_or_nothing_price'):
        strikes = tf.convert_to_tensor(strikes, dtype=dtype, name='strikes')
        dtype = strikes.dtype
        volatilities = tf.convert_to_tensor(volatilities, dtype=dtype, name='volatilities')
        expiries = tf.convert_to_tensor(expiries, dtype=dtype, name='expiries')
        if discount_rates is not None:
            discount_rates = tf.convert_to_tensor(discount_rates, dtype=dtype, name='discount_rates')
            discount_factors = tf.exp(-discount_rates * expiries)
        elif discount_factors is not None:
            discount_factors = tf.convert_to_tensor(discount_factors, dtype=dtype, name='discount_factors')
            discount_rates = -tf.math.log(discount_factors) / expiries
        else:
            discount_rates = tf.convert_to_tensor(0.0, dtype=dtype, name='discount_rates')
            discount_factors = tf.convert_to_tensor(1.0, dtype=dtype, name='discount_factors')
        if dividend_rates is None:
            dividend_rates = tf.convert_to_tensor(0.0, dtype=dtype, name='dividend_rates')
        if forwards is not None:
            forwards = tf.convert_to_tensor(forwards, dtype=dtype, name='forwards')
        else:
            spots = tf.convert_to_tensor(spots, dtype=dtype, name='spots')
            forwards = spots * tf.exp((discount_rates - dividend_rates) * expiries)
        sqrt_var = volatilities * tf.math.sqrt(expiries)
        if not is_normal_volatility:
            d1 = tf.math.divide_no_nan(tf.math.log(forwards / strikes), sqrt_var) + sqrt_var / 2
            undiscounted_calls = tf.where(sqrt_var > 0, forwards * _ncdf(d1), tf.where(forwards > strikes, forwards, 0.0))
        else:
            d1 = tf.math.divide_no_nan(forwards - strikes, sqrt_var)
            undiscounted_calls = tf.where(sqrt_var > 0.0, forwards * _ncdf(d1) + sqrt_var * tf.math.exp(-0.5 * d1 ** 2) / np.sqrt(2 * np.pi), tf.where(forwards > strikes, forwards, 0.0))
        if is_call_options is None:
            return discount_factors * undiscounted_calls
        undiscounted_puts = forwards - undiscounted_calls
        predicate = tf.broadcast_to(is_call_options, tf.shape(undiscounted_calls))
        return discount_factors * tf.where(predicate, undiscounted_calls, undiscounted_puts)

Computes the Black Scholes price for a batch of asset-or-nothing options.

The asset-or-nothing call (resp. put) pays out one unit of the underlying
asset if the spot is above (resp. below) the strike at maturity.

#### Example

```python
  # Price a batch of 5 asset_or_nothing call and put options.
  volatilities = np.array([0.0001, 102.0, 2.0, 0.1, 0.4])
  forwards = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
  # Strikes will automatically be broadcasted to shape [5].
  strikes = np.array([3.0])
  # Expiries will be broadcast to shape [5], i.e. each option has strike=3
  # and expiry = 1.
  expiries = 1.0
  computed_prices = tff.black_scholes.asset_or_nothing_price(
      volatilities=volatilities,
      strikes=strikes,
      expiries=expiries,
      forwards=forwards)
# Expected print output of prices:
# [0., 2., 2.52403424, 3.99315108, 4.65085383]
```

#### References:

[1] Hull, John C., Options, Futures and Other Derivatives. Pearson, 2018.
[2] https://en.wikipedia.org/wiki/Binary_option#Asset-or-nothing_call

Args:
  volatilities: Real `Tensor` of any shape and dtype. The volatilities to
    expiry of the options to price.
  strikes: A real `Tensor` of the same dtype and compatible shape as
    `volatilities`. The strikes of the options to be priced.
  expiries: A real `Tensor` of same dtype and compatible shape as
    `volatilities`. The expiry of each option.
  spots: A real `Tensor` of any shape that broadcasts to the shape of the
    `volatilities`. The current spot price of the underlying. Either this
    argument or the `forwards` (but not both) must be supplied.
  forwards: A real `Tensor` of any shape that broadcasts to the shape of
    `volatilities`. The forwards to maturity. Either this argument or the
    `spots` must be supplied but both must not be supplied.
  discount_rates: An optional real `Tensor` of same dtype as the
    `volatilities` and of the shape that broadcasts with `volatilities`. If
    not `None`, discount factors are calculated as e^(-rT), where r are the
    discount rates, or risk free rates. At most one of discount_rates and
    discount_factors can be supplied.
    Default value: `None`, equivalent to r = 0 and discount factors = 1 when
      discount_factors also not given.
  dividend_rates: An optional real `Tensor` of same dtype as the
    `volatilities` and of the shape that broadcasts with `volatilities`.
    Default value: `None`, equivalent to q = 0.
  discount_factors: An optional real `Tensor` of same dtype as the
    `volatilities`. If not `None`, these are the discount factors to expiry
    (i.e. e^(-rT)). Mutually exclusive with discount_rates. If neither is
    given, no discounting is applied (i.e. the undiscounted option price is
    returned). If `spots` is supplied and `discount_factors` is not `None`
    then this is also used to compute the forwards to expiry. At most one of
    `discount_rates` and `discount_factors` can be supplied.
    Default value: `None`, which maps to e^(-rT) calculated from
      discount_rates.
  is_call_options: A boolean `Tensor` of a shape compatible with
    `volatilities`. Indicates whether the option is a call (if True) or a put
    (if False). If not supplied, call options are assumed.
  is_normal_volatility: An optional Python boolean specifying whether the
    `volatilities` correspond to lognormal Black volatility (if False) or
    normal Black volatility (if True).
    Default value: False, which corresponds to lognormal volatility.
  dtype: Optional `tf.DType`. If supplied, the dtype to be used for conversion
    of any supplied non-`Tensor` arguments to `Tensor`.
    Default value: `None` which maps to the default dtype inferred by
      TensorFlow.
  name: str. The name for the ops created by this function.
    Default value: `None`, which is mapped to the default name
      `asset_or_nothing_price`.

Returns:
  option_prices: A `Tensor` of the same shape as `forwards`. The Black
  Scholes price of the options.

Raises:
  ValueError: If both `forwards` and `spots` are supplied or if neither is
    supplied.
  ValueError: If both `discount_rates` and `discount_factors` is supplied.


def he_uniform(seed=None):
    return VarianceScaling(scale=2.0, mode='fan_in', distribution='uniform', seed=seed)

He uniform variance scaling initializer.

It draws samples from a uniform distribution within [-limit, limit]
where `limit` is `sqrt(6 / fan_in)`
where `fan_in` is the number of input units in the weight tensor.

Args:
    seed: A Python integer. Used to seed the random generator.

Returns:
    An initializer.

References:
    [He et al., 2015]
    (https://www.cv-foundation.org/openaccess/content_iccv_2015/html/He_Delving_Deep_into_ICCV_2015_paper.html)
    # pylint: disable=line-too-long
    ([pdf](https://www.cv-foundation.org/openaccess/content_iccv_2015/papers/He_Delving_Deep_into_ICCV_2015_paper.pdf))


class XLAOptions(collections.namedtuple('XLAOptions', ['use_spmd_for_xla_partitioning', 'enable_xla_dynamic_padder'])):

    def __new__(cls, use_spmd_for_xla_partitioning=True, enable_xla_dynamic_padder=True):
        return super(XLAOptions, cls).__new__(cls, use_spmd_for_xla_partitioning, enable_xla_dynamic_padder)

XLA compilation options.

Attributes:
  use_spmd_for_xla_partitioning: Boolean. Whether to use XLA's SPMD
    partitioner instead of MPMD partitioner when compiler partitioning is
    requested.
  enable_xla_dynamic_padder: Boolean. Whether to enable XLA dynamic padder
    infrastructure to handle dynamic shapes inputs inside XLA. True by
    default. Disabling this may cause correctness issues with dynamic shapes
    inputs, as XLA will just assume the inputs are with padded shapes. However
    users can optionally set it to False to improve device time if masking is
    already handled in the user side.


def size(self, name=None):
    with ops.name_scope(name, '%s_Size' % self.name, [self.resource_handle]):
        return gen_lookup_ops.lookup_table_size_v2(self.resource_handle)

Compute the number of elements in this table.

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

Returns:
  A scalar tensor containing the number of elements in this table.


def layer_test(layer_cls, kwargs=None, input_shape=None, input_dtype=None, input_data=None, expected_output=None, expected_output_dtype=None, expected_output_shape=None, validate_training=True, adapt_data=None, custom_objects=None, test_harness=None, supports_masking=None):
    if input_data is None:
        if input_shape is None:
            raise ValueError('input_shape is None')
        if not input_dtype:
            input_dtype = 'float32'
        input_data_shape = list(input_shape)
        for i, e in enumerate(input_data_shape):
            if e is None:
                input_data_shape[i] = np.random.randint(1, 4)
        input_data = 10 * np.random.random(input_data_shape)
        if input_dtype[:5] == 'float':
            input_data -= 0.5
        input_data = input_data.astype(input_dtype)
    elif input_shape is None:
        input_shape = input_data.shape
    if input_dtype is None:
        input_dtype = input_data.dtype
    if expected_output_dtype is None:
        expected_output_dtype = input_dtype
    if dtypes.as_dtype(expected_output_dtype) == dtypes.string:
        if test_harness:
            assert_equal = test_harness.assertAllEqual
        else:
            assert_equal = string_test
    elif test_harness:
        assert_equal = test_harness.assertAllClose
    else:
        assert_equal = numeric_test
    kwargs = kwargs or {}
    layer = layer_cls(**kwargs)
    if supports_masking is not None and layer.supports_masking != supports_masking:
        raise AssertionError('When testing layer %s, the `supports_masking` property is %rbut expected to be %r.\nFull kwargs: %s' % (layer_cls.__name__, layer.supports_masking, supports_masking, kwargs))
    if adapt_data is not None:
        layer.adapt(adapt_data)
    weights = layer.get_weights()
    layer.set_weights(weights)
    if 'weights' in tf_inspect.getargspec(layer_cls.__init__):
        kwargs['weights'] = weights
        layer = layer_cls(**kwargs)
    x = layers.Input(shape=input_shape[1:], dtype=input_dtype)
    y = layer(x)
    if backend.dtype(y) != expected_output_dtype:
        raise AssertionError('When testing layer %s, for input %s, found output dtype=%s but expected to find %s.\nFull kwargs: %s' % (layer_cls.__name__, x, backend.dtype(y), expected_output_dtype, kwargs))

    def assert_shapes_equal(expected, actual):
        """Asserts that the output shape from the layer matches the actual shape."""
        if len(expected) != len(actual):
            raise AssertionError('When testing layer %s, for input %s, found output_shape=%s but expected to find %s.\nFull kwargs: %s' % (layer_cls.__name__, x, actual, expected, kwargs))
        for expected_dim, actual_dim in zip(expected, actual):
            if isinstance(expected_dim, tensor_shape.Dimension):
                expected_dim = expected_dim.value
            if isinstance(actual_dim, tensor_shape.Dimension):
                actual_dim = actual_dim.value
            if expected_dim is not None and expected_dim != actual_dim:
                raise AssertionError('When testing layer %s, for input %s, found output_shape=%s but expected to find %s.\nFull kwargs: %s' % (layer_cls.__name__, x, actual, expected, kwargs))
    if expected_output_shape is not None:
        assert_shapes_equal(tensor_shape.TensorShape(expected_output_shape), y.shape)
    model = models.Model(x, y)
    computed_output_shape = tuple(layer.compute_output_shape(tensor_shape.TensorShape(input_shape)).as_list())
    computed_output_signature = layer.compute_output_signature(tensor_spec.TensorSpec(shape=input_shape, dtype=input_dtype))
    actual_output = model.predict(input_data)
    actual_output_shape = actual_output.shape
    assert_shapes_equal(computed_output_shape, actual_output_shape)
    assert_shapes_equal(computed_output_signature.shape, actual_output_shape)
    if computed_output_signature.dtype != actual_output.dtype:
        raise AssertionError('When testing layer %s, for input %s, found output_dtype=%s but expected to find %s.\nFull kwargs: %s' % (layer_cls.__name__, x, actual_output.dtype, computed_output_signature.dtype, kwargs))
    if expected_output is not None:
        assert_equal(actual_output, expected_output)
    model_config = model.get_config()
    recovered_model = models.Model.from_config(model_config, custom_objects)
    if model.weights:
        weights = model.get_weights()
        recovered_model.set_weights(weights)
        output = recovered_model.predict(input_data)
        assert_equal(output, actual_output)
    layer_weights = layer.get_weights()
    if validate_training:
        model = models.Model(x, layer(x))
        if _thread_local_data.run_eagerly is not None:
            model.compile('rmsprop', 'mse', weighted_metrics=['acc'], run_eagerly=should_run_eagerly())
        else:
            model.compile('rmsprop', 'mse', weighted_metrics=['acc'])
        model.train_on_batch(input_data, actual_output)
    layer_config = layer.get_config()
    layer_config['batch_input_shape'] = input_shape
    layer = layer.__class__.from_config(layer_config)
    if adapt_data is not None:
        layer.adapt(adapt_data)
    model = models.Sequential()
    model.add(layers.Input(shape=input_shape[1:], dtype=input_dtype))
    model.add(layer)
    layer.set_weights(layer_weights)
    actual_output = model.predict(input_data)
    actual_output_shape = actual_output.shape
    for expected_dim, actual_dim in zip(computed_output_shape, actual_output_shape):
        if expected_dim is not None:
            if expected_dim != actual_dim:
                raise AssertionError('When testing layer %s **after deserialization**, for input %s, found output_shape=%s but expected to find inferred shape %s.\nFull kwargs: %s' % (layer_cls.__name__, x, actual_output_shape, computed_output_shape, kwargs))
    if expected_output is not None:
        assert_equal(actual_output, expected_output)
    model_config = model.get_config()
    recovered_model = models.Sequential.from_config(model_config, custom_objects)
    if model.weights:
        weights = model.get_weights()
        recovered_model.set_weights(weights)
        output = recovered_model.predict(input_data)
        assert_equal(output, actual_output)
    return actual_output

Test routine for a layer with a single input and single output.

Args:
  layer_cls: Layer class object.
  kwargs: Optional dictionary of keyword arguments for instantiating the
    layer.
  input_shape: Input shape tuple.
  input_dtype: Data type of the input data.
  input_data: Numpy array of input data.
  expected_output: Numpy array of the expected output.
  expected_output_dtype: Data type expected for the output.
  expected_output_shape: Shape tuple for the expected shape of the output.
  validate_training: Whether to attempt to validate training on this layer.
    This might be set to False for non-differentiable layers that output
    string or integer values.
  adapt_data: Optional data for an 'adapt' call. If None, adapt() will not
    be tested for this layer. This is only relevant for PreprocessingLayers.
  custom_objects: Optional dictionary mapping name strings to custom objects
    in the layer class. This is helpful for testing custom layers.
  test_harness: The Tensorflow test, if any, that this function is being
    called in.
  supports_masking: Optional boolean to check the `supports_masking` property
    of the layer. If None, the check will not be performed.

Returns:
  The output data (Numpy array) returned by the layer, for additional
  checks to be done by the calling code.

Raises:
  ValueError: if `input_shape is None`.


class BufferedQuantileTracker(WindowedTracker, QuantileTracker):

    def __init__(self, window_mode, q, **kwargs):
        super().__init__(window_mode, **kwargs)
        QuantileTracker.__init__(self, q)
        self._sorted_items = SortedList()

    def push(self, x):
        """Pushes a new value, maintains the sorted list, and manages the window.

    Args:
      x: The new value to be pushed.
    """
        if not math.isnan(x):
            self._sorted_items.add(x)
        if self._window_mode == WindowMode.SLIDING:
            if len(self._queue) >= self._window_size and (not math.isnan((old_x := self.pop()))):
                self._sorted_items.discard(old_x)
            super().push(x)

    @staticmethod
    def _get_helper(sorted_items, q):
        n = len(sorted_items)
        if n < 1:
            return float('nan')
        pos = q * (n - 1)
        lo = math.floor(pos)
        lo_value = typing.cast(float, sorted_items[lo])
        hi = min(lo + 1, n - 1)
        hi_value: float = typing.cast(float, sorted_items[hi])
        return lo_value + (hi_value - lo_value) * (pos - lo)

    def get(self):
        """Returns the current quantile value using the sorted list.

    Calculates the quantile using linear interpolation on the sorted values.

    Returns:
      float: The calculated quantile value. Returns NaN if the window is empty.
    """
        return self._get_helper(self._sorted_items, self._q)

Abstract base class for buffered quantile trackers.

Warning:
  Buffered quantile trackers are NOT truly incremental in the sense that they
  don't update the quantile in constant time per new data point. They maintain
  a sorted list of all values in the window.

Args:
  window_mode: A `WindowMode` enum specifying whether the window is `LANDMARK`
    or `SLIDING`.
  q: The quantile to calculate, a float between 0 and 1 (inclusive).
  **kwargs: Keyword arguments passed to the parent class constructor.


def _maybe_colocate_with(self, value):
    if not self._colocate_with_first_write_call:
        yield
    else:
        if not self._colocate_with:
            self._colocate_with.append(value)
        with ops.colocate_with(self._colocate_with[0]):
            yield

Colocate operations with an internal colocation group or `value`.

Args:
  value: `Tensor`, the tensor to try to colocate with.

Yields:
  Does not yield anything, but the new context is a colocation context.

If no internal colocation group is set, colocate with `value` and set
the internal colocation group to be value.


class MaxNorm(Constraint):

    def __init__(self, max_value=2, axis=0):
        self.max_value = max_value
        self.axis = axis

    @doc_controls.do_not_generate_docs
    def __call__(self, w):
        norms = backend.sqrt(math_ops.reduce_sum(math_ops.square(w), axis=self.axis, keepdims=True))
        desired = backend.clip(norms, 0, self.max_value)
        return w * (desired / (backend.epsilon() + norms))

    @doc_controls.do_not_generate_docs
    def get_config(self):
        return {'max_value': self.max_value, 'axis': self.axis}

MaxNorm weight constraint.

Constrains the weights incident to each hidden unit
to have a norm less than or equal to a desired value.

Also available via the shortcut function `tf.keras.constraints.max_norm`.

Args:
  max_value: the maximum norm value for the incoming weights.
  axis: integer, axis along which to calculate weight norms.
    For instance, in a `Dense` layer the weight matrix
    has shape `(input_dim, output_dim)`,
    set `axis` to `0` to constrain each weight vector
    of length `(input_dim,)`.
    In a `Conv2D` layer with `data_format="channels_last"`,
    the weight tensor has shape
    `(rows, cols, input_depth, output_depth)`,
    set `axis` to `[0, 1, 2]`
    to constrain the weights of each filter tensor of size
    `(rows, cols, input_depth)`.


def multiplex(cond, a, b, name=None):
    return examples_multiplex_dense(cond=cond, a=a, b=b, name=name)

Return elements chosen from `a` or `b` depending on `cond`.

This is similar to `np.where` and `tf.where`, but simplified to only handle
the case of dense tensors, no optional parameters, no broadcasting, etc..

>>> multiplex([True, False, False, True], [1,2,3,4], [100,200,300,400])
<tf.Tensor: shape=(4,), dtype=int32, numpy=array([  1, 200, 300,   4], ...)>

Args:
  cond: tf.Tensor of type bool. Where True, yield `a`, otherwise yield `b`.
  a: tf.Tensor with the same type and shape as `b`.
  b: tf.Tensor with the same type and shape as `a`.
  name: An optional name for the op.

Returns:
  A tf.Tensor with elements from `a` where `cond` is True, and elements
  from `b` elsewhere.


def get(self):
    if self._n < 1:
        return float('nan')
    return self._mean

Returns the current incremental mean.

Returns:
  float: The current incremental mean value.
         Returns NaN if no valid (non-NaN) values have been pushed.


def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub: bool=False, **kwargs):
    use_auth_token = kwargs.pop('use_auth_token', None)
    if use_auth_token is not None:
        warnings.warn('The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.', FutureWarning)
        if kwargs.get('token', None) is not None:
            raise ValueError('`token` and `use_auth_token` are both specified. Please set only the argument `token`.')
        kwargs['token'] = use_auth_token
    if os.path.isfile(save_directory):
        raise AssertionError(f'Provided path ({save_directory}) should be a directory, not a file')
    os.makedirs(save_directory, exist_ok=True)
    if push_to_hub:
        commit_message = kwargs.pop('commit_message', None)
        repo_id = kwargs.pop('repo_id', save_directory.split(os.path.sep)[-1])
        repo_id = self._create_repo(repo_id, **kwargs)
        files_timestamps = self._get_files_timestamps(save_directory)
    if self._auto_class is not None:
        custom_object_save(self, save_directory, config=self)
    output_video_processor_file = os.path.join(save_directory, VIDEO_PROCESSOR_NAME)
    self.to_json_file(output_video_processor_file)
    logger.info(f'Video processor saved in {output_video_processor_file}')
    if push_to_hub:
        self._upload_modified_files(save_directory, repo_id, files_timestamps, commit_message=commit_message, token=kwargs.get('token'))
    return [output_video_processor_file]

Save an video processor object to the directory `save_directory`, so that it can be re-loaded using the
[`~video_processing_utils.VideoProcessorBase.from_pretrained`] class method.

Args:
    save_directory (`str` or `os.PathLike`):
        Directory where the video processor JSON file will be saved (will be created if it does not exist).
    push_to_hub (`bool`, *optional*, defaults to `False`):
        Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the
        repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
        namespace).
    kwargs (`Dict[str, Any]`, *optional*):
        Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.


def __init__(self, project, sub_name, expected_msg=None, expected_msg_len=None, timeout=DEFAULT_TIMEOUT, with_attributes=False, strip_attributes=None, sleep_time=DEFAULT_SLEEP_TIME, max_messages_in_one_pull=DEFAULT_MAX_MESSAGES_IN_ONE_PULL, pull_timeout=DEFAULT_PULL_TIMEOUT):
    if pubsub is None:
        raise ImportError('PubSub dependencies are not installed.')
    if not project:
        raise ValueError('Invalid project %s.' % project)
    if not sub_name:
        raise ValueError('Invalid subscription %s.' % sub_name)
    if not expected_msg_len and (not expected_msg):
        raise ValueError('Required expected_msg: {} or expected_msg_len: {}.'.format(expected_msg, expected_msg_len))
    if expected_msg and (not isinstance(expected_msg, list)):
        raise ValueError('Invalid expected messages %s.' % expected_msg)
    if expected_msg_len and (not isinstance(expected_msg_len, int)):
        raise ValueError('Invalid expected messages %s.' % expected_msg_len)
    self.project = project
    self.sub_name = sub_name
    self.expected_msg = expected_msg
    self.expected_msg_len = expected_msg_len or len(self.expected_msg)
    self.timeout = timeout
    self.messages = None
    self.messages_all_details = None
    self.with_attributes = with_attributes
    self.strip_attributes = strip_attributes
    self.sleep_time = sleep_time
    self.max_messages_in_one_pull = max_messages_in_one_pull
    self.pull_timeout = pull_timeout

Initialize PubSubMessageMatcher object.

Args:
  project: A name string of project.
  sub_name: A name string of subscription which is attached to output.
  expected_msg: A string list that contains expected message data pulled
    from the subscription. See also: with_attributes.
  expected_msg_len: Number of expected messages pulled from the
    subscription.
  timeout: Timeout in seconds to wait for all expected messages appears.
  with_attributes: If True, will match against both message data and
    attributes. If True, expected_msg should be a list of ``PubsubMessage``
    objects. Otherwise, it should be a list of ``bytes``.
  strip_attributes: List of strings. If with_attributes==True, strip the
    attributes keyed by these values from incoming messages.
    If a key is missing, will add an attribute with an error message as
    value to prevent a successful match.
  sleep_time: Time in seconds between which the pulls from pubsub are done.
  max_messages_in_one_pull: Maximum number of messages pulled from pubsub
    at once.
  pull_timeout: Time in seconds after which the pull from pubsub is repeated


class Speech2Text2Decoder(Speech2Text2PreTrainedModel):

    def __init__(self, config: Speech2Text2Config):
        super().__init__(config)
        self.dropout = config.dropout
        self.layerdrop = config.decoder_layerdrop
        self.padding_idx = config.pad_token_id
        self.max_target_positions = config.max_target_positions
        self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
        self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model, self.padding_idx)
        self.embed_positions = Speech2Text2SinusoidalPositionalEmbedding(self.max_target_positions, config.d_model, self.padding_idx)
        self.layers = nn.ModuleList([Speech2Text2DecoderLayer(config) for _ in range(config.decoder_layers)])
        self.gradient_checkpointing = False
        self.post_init()

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, value):
        self.embed_tokens = value

    def forward(self, input_ids=None, attention_mask=None, encoder_hidden_states=None, encoder_attention_mask=None, head_mask=None, cross_attn_head_mask=None, past_key_values=None, inputs_embeds=None, use_cache=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 [`Speech2Text2Tokenizer`]. 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)
            encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):
                Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention
                of the decoder.
            encoder_attention_mask (`torch.LongTensor` of shape `(batch_size, encoder_sequence_length)`, *optional*):
                Mask to avoid performing cross-attention on padding tokens indices of encoder input_ids. 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)
            head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
                Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:

                - 1 indicates the head is **not masked**,
                - 0 indicates the head is **masked**.

            cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
                Mask to nullify selected heads of the attention modules in encoder to avoid performing cross-attention
                on hidden heads. Mask values selected in `[0, 1]`:

                - 1 indicates the head is **not masked**,
                - 0 indicates the head is **masked**.

            past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
                Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
                shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of
                shape `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.

                Contains pre-computed hidden-states (key and values in the self-attention blocks and in the
                cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.

                If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those
                that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of
                all `decoder_input_ids` of shape `(batch_size, sequence_length)`.
            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
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        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 decoder_input_ids and decoder_inputs_embeds at the same time')
        elif input_ids is not None:
            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 decoder_input_ids or decoder_inputs_embeds')
        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
        attention_mask = _prepare_4d_causal_attention_mask(attention_mask, input_shape, inputs_embeds, past_key_values_length)
        if encoder_hidden_states is not None and encoder_attention_mask is not None:
            encoder_attention_mask = _prepare_4d_attention_mask(encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1])
        positions = self.embed_positions(input_ids, past_key_values_length=past_key_values_length)
        hidden_states = inputs_embeds + positions
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        if self.gradient_checkpointing and self.training:
            if use_cache:
                logger.warning_once('`use_cache = True` is incompatible with gradient checkpointing. Setting `use_cache = False`...')
                use_cache = False
        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None
        all_cross_attentions = () if output_attentions and encoder_hidden_states is not None else None
        next_decoder_cache = () if use_cache else None
        for attn_mask, mask_name in zip([head_mask, cross_attn_head_mask], ['head_mask', 'cross_attn_head_mask']):
            if attn_mask is not None:
                if attn_mask.size()[0] != len(self.layers):
                    raise ValueError(f'The `{mask_name}` should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}.')
        for idx, decoder_layer in enumerate(self.layers):
            if output_hidden_states:
                all_hidden_states += (hidden_states,)
            if self.training:
                dropout_probability = torch.rand([])
                if dropout_probability < self.layerdrop:
                    continue
            past_key_value = past_key_values[idx] if past_key_values is not None else None
            if self.gradient_checkpointing and self.training:
                layer_outputs = self._gradient_checkpointing_func(decoder_layer.__call__, hidden_states, attention_mask, encoder_hidden_states, encoder_attention_mask, head_mask[idx] if head_mask is not None else None, cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None, None)
            else:
                layer_outputs = decoder_layer(hidden_states, attention_mask=attention_mask, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, layer_head_mask=head_mask[idx] if head_mask is not None else None, cross_attn_layer_head_mask=cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None, past_key_value=past_key_value, output_attentions=output_attentions, use_cache=use_cache)
            hidden_states = layer_outputs[0]
            if use_cache:
                next_decoder_cache += (layer_outputs[3 if output_attentions else 1],)
            if output_attentions:
                all_self_attns += (layer_outputs[1],)
                if encoder_hidden_states is not None:
                    all_cross_attentions += (layer_outputs[2],)
        if output_hidden_states:
            all_hidden_states += (hidden_states,)
        next_cache = next_decoder_cache if use_cache else None
        if not return_dict:
            return tuple((v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_cross_attentions] if v is not None))
        return BaseModelOutputWithPastAndCrossAttentions(last_hidden_state=hidden_states, past_key_values=next_cache, hidden_states=all_hidden_states, attentions=all_self_attns, cross_attentions=all_cross_attentions)

Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`Speech2Text2DecoderLayer`]

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


def call(self, input_ids: TFModelInputType | None=None, attention_mask: np.ndarray | tf.Tensor | None=None, token_type_ids: np.ndarray | tf.Tensor | None=None, position_ids: np.ndarray | tf.Tensor | None=None, head_mask: np.ndarray | tf.Tensor | None=None, inputs_embeds: np.ndarray | tf.Tensor | None=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, labels: np.ndarray | tf.Tensor | None=None, next_sentence_label: np.ndarray | tf.Tensor | None=None, training: Optional[bool]=False) -> Union[TFBertForPreTrainingOutput, Tuple[tf.Tensor]]:
    outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, training=training)
    sequence_output, pooled_output = outputs[:2]
    prediction_scores = self.mlm(sequence_output=sequence_output, training=training)
    seq_relationship_score = self.nsp(pooled_output=pooled_output)
    total_loss = None
    if labels is not None and next_sentence_label is not None:
        d_labels = {'labels': labels}
        d_labels['next_sentence_label'] = next_sentence_label
        total_loss = self.hf_compute_loss(labels=d_labels, logits=(prediction_scores, seq_relationship_score))
    if not return_dict:
        output = (prediction_scores, seq_relationship_score) + outputs[2:]
        return (total_loss,) + output if total_loss is not None else output
    return TFBertForPreTrainingOutput(loss=total_loss, prediction_logits=prediction_scores, seq_relationship_logits=seq_relationship_score, hidden_states=outputs.hidden_states, attentions=outputs.attentions)

labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
    Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
    config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
    loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
next_sentence_label (`tf.Tensor` of shape `(batch_size,)`, *optional*):
    Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair
    (see `input_ids` docstring) Indices should be in `[0, 1]`:

    - 0 indicates sequence B is a continuation of sequence A,
    - 1 indicates sequence B is a random sequence.
kwargs (`Dict[str, any]`, *optional*, defaults to `{}`):
    Used to hide legacy arguments that have been deprecated.

Return:

Examples:

```python
>>> import tensorflow as tf
>>> from transformers import AutoTokenizer, TFBertForPreTraining

>>> tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
>>> model = TFBertForPreTraining.from_pretrained("google-bert/bert-base-uncased")
>>> input_ids = tokenizer("Hello, my dog is cute", add_special_tokens=True, return_tensors="tf")
>>> # Batch size 1

>>> outputs = model(input_ids)
>>> prediction_logits, seq_relationship_logits = outputs[:2]
```


def _matmul_2d(a, b, **kwargs):
    ragged_err = 'The matrices in `a` and `b` may not be ragged in their innermost dimension.'
    checks = []
    if isinstance(a, ragged_tensor.RaggedTensor):
        original_size = array_ops.size(a.flat_values)
        a = a.to_tensor()
        checks.append(check_ops.assert_equal(original_size, array_ops.size(a), message=ragged_err))
    if isinstance(b, ragged_tensor.RaggedTensor):
        original_size = array_ops.size(b.flat_values)
        b = b.to_tensor()
        checks.append(check_ops.assert_equal(original_size, array_ops.size(b), message=ragged_err))
    with ops.control_dependencies(checks):
        return math_ops.matmul(a, b, **kwargs)

Multiplies potentially ragged 2D tensors.

Args:
  a: A 2D Tensor or RaggedTensor with `shape=[I, J]`
  b: A 2D Tensor or RaggedTensor with `shape=[J, K]`
  **kwargs: Additional arguments for `tf.matmul` (e.g. transpose_a).

Returns:
  A 2D Tensor with `shape=[I, K]`.


A function that is executed upon a test being skipped.

Implementation is optional.

Args:
  record: records.TestResultRecord, a copy of the test record for
    this test, containing all information of the test execution
    including exception objects.


def forward(self, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.FloatTensor]=None, token_type_ids: Optional[torch.LongTensor]=None, position_ids: Optional[torch.LongTensor]=None, head_mask: Optional[torch.FloatTensor]=None, inputs_embeds: Optional[torch.FloatTensor]=None, encoder_hidden_states: Optional[torch.FloatTensor]=None, encoder_attention_mask: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, **kwargs) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    if labels is not None:
        use_cache = False
    outputs = self.roberta(input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, past_key_values=past_key_values, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    sequence_output = outputs[0]
    prediction_scores = self.lm_head(sequence_output)
    lm_loss = None
    if labels is not None:
        lm_loss = self.loss_function(prediction_scores, labels, vocab_size=self.config.vocab_size, **kwargs)
    if not return_dict:
        output = (prediction_scores,) + outputs[2:]
        return (lm_loss,) + output if lm_loss is not None else output
    return CausalLMOutputWithCrossAttentions(loss=lm_loss, logits=prediction_scores, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, cross_attentions=outputs.cross_attentions)

labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
    Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
    `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are
    ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`

Example:

```python
>>> from transformers import AutoTokenizer, RobertaForCausalLM, RobertaConfig
>>> import torch

>>> tokenizer = AutoTokenizer.from_pretrained("FacebookAI/roberta-base")
>>> config = RobertaConfig.from_pretrained("FacebookAI/roberta-base")
>>> config.is_decoder = True
>>> model = RobertaForCausalLM.from_pretrained("FacebookAI/roberta-base", config=config)
>>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
>>> outputs = model(**inputs)
>>> prediction_logits = outputs.logits
```


def pad(self, images: List[np.ndarray], annotations: Optional[Union[AnnotationType, List[AnnotationType]]]=None, constant_values: Union[float, Iterable[float]]=0, return_pixel_mask: bool=True, return_tensors: Optional[Union[str, TensorType]]=None, data_format: Optional[ChannelDimension]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None, update_bboxes: bool=True, pad_size: Optional[Dict[str, int]]=None) -> BatchFeature:
    pad_size = pad_size if pad_size is not None else self.pad_size
    if pad_size is not None:
        padded_size = (pad_size['height'], pad_size['width'])
    else:
        padded_size = get_max_height_width(images, input_data_format=input_data_format)
    annotation_list = annotations if annotations is not None else [None] * len(images)
    padded_images = []
    padded_annotations = []
    for image, annotation in zip(images, annotation_list):
        padded_image, padded_annotation = self._pad_image(image, padded_size, annotation, constant_values=constant_values, data_format=data_format, input_data_format=input_data_format, update_bboxes=update_bboxes)
        padded_images.append(padded_image)
        padded_annotations.append(padded_annotation)
    data = {'pixel_values': padded_images}
    if return_pixel_mask:
        masks = [make_pixel_mask(image=image, output_size=padded_size, input_data_format=input_data_format) for image in images]
        data['pixel_mask'] = masks
    encoded_inputs = BatchFeature(data=data, tensor_type=return_tensors)
    if annotations is not None:
        encoded_inputs['labels'] = [BatchFeature(annotation, tensor_type=return_tensors) for annotation in padded_annotations]
    return encoded_inputs

Pads a batch of images to the bottom and right of the image with zeros to the size of largest height and width
in the batch and optionally returns their corresponding pixel mask.

Args:
    images (List[`np.ndarray`]):
        Images to pad.
    annotations (`AnnotationType` or `List[AnnotationType]`, *optional*):
        Annotations to transform according to the padding that is applied to the images.
    constant_values (`float` or `Iterable[float]`, *optional*):
        The value to use for the padding if `mode` is `"constant"`.
    return_pixel_mask (`bool`, *optional*, defaults to `True`):
        Whether to return a pixel mask.
    return_tensors (`str` or `TensorType`, *optional*):
        The type of tensors to return. Can be one of:
            - Unset: Return a list of `np.ndarray`.
            - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
            - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
            - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
            - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
    data_format (`str` or `ChannelDimension`, *optional*):
        The channel dimension format of the image. If not provided, it will be the same as the input image.
    input_data_format (`ChannelDimension` or `str`, *optional*):
        The channel dimension format of the input image. If not provided, it will be inferred.
    update_bboxes (`bool`, *optional*, defaults to `True`):
        Whether to update the bounding boxes in the annotations to match the padded images. If the
        bounding boxes have not been converted to relative coordinates and `(centre_x, centre_y, width, height)`
        format, the bounding boxes will not be updated.
    pad_size (`Dict[str, int]`, *optional*):
        The size `{"height": int, "width" int}` to pad the images to. Must be larger than any image size
        provided for preprocessing. If `pad_size` is not provided, images will be padded to the largest
        height and width in the batch.


def __init__(self, cluster_spec, initializer=None, share_gpu=True):
    _active_pool_runners.add(self)
    self._cluster_spec = cluster_spec
    self._initializer = initializer
    self._share_gpu = share_gpu
    self._conn = {}
    self._runner = None

Creates a multi-process pool runner.

Args:
  cluster_spec: Dict for cluster spec. The following is an example of
    cluster with three workers.
    {"worker": ["worker0.example.com:2222",
                "worker1.example.com:2222",
                "worker2.example.com:2222"]}
  initializer: a callable to called at the startup of worker processes.
  share_gpu: Whether to share GPUs among workers. If False, each worker is
    assigned different GPUs in a roundrobin fashion.

Raises:
  RuntimeError: if `multi_process_runner.test_main()` is not called.
  ValueError: if there are more than one chief in the `cluster_spec`.


def create_band_mask_from_inputs(from_blocked_mask, to_blocked_mask):
    exp_blocked_to_pad = torch.cat([to_blocked_mask[:, 1:-3], to_blocked_mask[:, 2:-2], to_blocked_mask[:, 3:-1]], dim=2)
    band_mask = torch.einsum('blq,blk->blqk', from_blocked_mask[:, 2:-2], exp_blocked_to_pad)
    band_mask.unsqueeze_(1)
    return band_mask

Create 3D attention mask from a 2D tensor mask.

Args:
    from_blocked_mask: 2D Tensor of shape [batch_size,
    from_seq_length//from_block_size, from_block_size].
    to_blocked_mask: int32 Tensor of shape [batch_size,
    to_seq_length//to_block_size, to_block_size].

Returns:
    float Tensor of shape [batch_size, 1, from_seq_length//from_block_size-4, from_block_size,
    3*to_block_size].