def run_saved_model_with_feed_dict(saved_model_dir, tag_set, signature_def_key, input_tensor_key_feed_dict, outdir, overwrite_flag, worker=None, init_tpu=False, use_tfrt=False, tf_debug=False):
    meta_graph_def = saved_model_utils.get_meta_graph_def(saved_model_dir, tag_set)
    inputs_tensor_info = _get_inputs_tensor_info_from_meta_graph_def(meta_graph_def, signature_def_key)
    for input_key_name in input_tensor_key_feed_dict.keys():
        if input_key_name not in inputs_tensor_info:
            raise ValueError('"%s" is not a valid input key. Please choose from %s, or use --show option.' % (input_key_name, '"' + '", "'.join(inputs_tensor_info.keys()) + '"'))
    inputs_feed_dict = {inputs_tensor_info[key].name: tensor for key, tensor in input_tensor_key_feed_dict.items()}
    outputs_tensor_info = _get_outputs_tensor_info_from_meta_graph_def(meta_graph_def, signature_def_key)
    output_tensor_keys_sorted = sorted(outputs_tensor_info.keys())
    output_tensor_names_sorted = [outputs_tensor_info[tensor_key].name for tensor_key in output_tensor_keys_sorted]
    config = None
    if use_tfrt:
        logging.info('Using TFRT session.')
        config = config_pb2.ConfigProto(experimental=config_pb2.ConfigProto.Experimental(use_tfrt=True))
    with session.Session(worker, graph=ops_lib.Graph(), config=config) as sess:
        if init_tpu:
            print('Initializing TPU System ...')
            sess.run(tpu.initialize_system())
        loader.load(sess, tag_set.split(','), saved_model_dir)
        if tf_debug:
            sess = local_cli_wrapper.LocalCLIDebugWrapperSession(sess)
        outputs = sess.run(output_tensor_names_sorted, feed_dict=inputs_feed_dict)
        for i, output in enumerate(outputs):
            output_tensor_key = output_tensor_keys_sorted[i]
            print('Result for output key %s:\n%s' % (output_tensor_key, output))
            if outdir:
                if not os.path.isdir(outdir):
                    os.makedirs(outdir)
                output_full_path = os.path.join(outdir, output_tensor_key + '.npy')
                if not overwrite_flag and os.path.exists(output_full_path):
                    raise RuntimeError('Output file %s already exists. Add "--overwrite" to overwrite the existing output files.' % output_full_path)
                np.save(output_full_path, output)
                print('Output %s is saved to %s' % (output_tensor_key, output_full_path))

Runs SavedModel and fetch all outputs.

Runs the input dictionary through the MetaGraphDef within a SavedModel
specified by the given tag_set and SignatureDef. Also save the outputs to file
if outdir is not None.

Args:
  saved_model_dir: Directory containing the SavedModel to execute.
  tag_set: Group of tag(s) of the MetaGraphDef with the SignatureDef map, in
      string format, separated by ','. For tag-set contains multiple tags, all
      tags must be passed in.
  signature_def_key: A SignatureDef key string.
  input_tensor_key_feed_dict: A dictionary maps input keys to numpy ndarrays.
  outdir: A directory to save the outputs to. If the directory doesn't exist,
      it will be created.
  overwrite_flag: A boolean flag to allow overwrite output file if file with
      the same name exists.
  worker: If provided, the session will be run on the worker.  Valid worker
      specification is a bns or gRPC path.
  init_tpu: If true, the TPU system will be initialized after the session
      is created.
  use_tfrt: If true, TFRT session will be used.
  tf_debug: A boolean flag to use TensorFlow Debugger (TFDBG) to observe the
      intermediate Tensor values and runtime GraphDefs while running the
      SavedModel.

Raises:
  ValueError: When any of the input tensor keys is not valid.
  RuntimeError: An error when output file already exists and overwrite is not
  enabled.


class CollapseLongUnions(visitors.Visitor):

    def __init__(self, max_length: int=7):
        super().__init__()
        self.generic_type = pytd.AnythingType()
        self.max_length = max_length

    def VisitUnionType(self, union):
        if len(union.type_list) > self.max_length and (not any((isinstance(t, pytd.Literal) for t in union.type_list))):
            return self.generic_type
        elif self.generic_type in union.type_list:
            return pytd_utils.JoinTypes(union.type_list)
        else:
            return union

Shortens long unions to object (or "?").

Poor man's version of FindCommonSuperClasses. Shorten types like
"str or unicode or int or float or list" to just "object" or "?".

Additionally, if the union already contains at least one "object", we also
potentially replace the entire union with just "object".

Attributes:
  max_length: The maximum number of types to allow in a union. If there are
    more types than this, it is shortened.


class Siglip2Processor(ProcessorMixin):
    attributes = ['image_processor', 'tokenizer']
    image_processor_class = 'AutoImageProcessor'
    tokenizer_class = 'AutoTokenizer'

    def __init__(self, image_processor, tokenizer):
        super().__init__(image_processor, tokenizer)

    def __call__(self, images: Optional[Union[ImageInput, List[ImageInput], List[List[ImageInput]]]]=None, text: Optional[Union[TextInput, 'PreTokenizedInput', List[TextInput], List['PreTokenizedInput']]]=None, audio=None, videos=None, **kwargs: Unpack[Siglip2ProcessorKwargs]) -> BatchFeature:
        """
        Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
        and `kwargs` arguments to GemmaTokenizerFast's [`~GemmaTokenizerFast.__call__`] if `text` is not `None` to encode
        the text. To prepare the image(s), this method forwards the `images` argument to
        Siglip2ImageProcessor's [`~Siglip2ImageProcessor.__call__`] if `images` is not `None`. Please refer to the docstring
        of the above two methods for more information.

        Args:
            images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`):
                The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
                tensor. Both channels-first and channels-last formats are supported.
            text (`str`, `List[str]`, `List[List[str]]`):
                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
            padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `max_length`):
                Select a strategy to pad the returned sequences (according to the model's padding side and padding
                index) among:
                - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
                  acceptable input length for the model if that argument is not provided.
                - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
                  sequence if provided).
                - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different
                  lengths).
            max_length (`int`, *optional*, defaults to 64):
                Maximum length of the returned list and optionally padding length (see above).
            truncation (`bool`, *optional*, defaults to `True`):
                Activates truncation to cut input sequences longer than `max_length` to `max_length`.
            return_tensors (`str` or [`~utils.TensorType`], *optional*, defaults to `'pt'`):
                If set, will return tensors of a particular framework. Acceptable values are:

                - `'tf'`: Return TensorFlow `tf.constant` objects.
                - `'pt'`: Return PyTorch `torch.Tensor` objects.
                - `'np'`: Return NumPy `np.ndarray` objects.
                - `'jax'`: Return JAX `jnp.ndarray` objects.

        Returns:
            [`BatchFeature`]: A [`BatchFeature`] with the following fields:

            - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
              `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
              `None`).
            - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
            - **pixel_attention_mask** -- Attention mask for the pixel values. Returned when `images` is not `None`.
            - **spatial_shapes** -- The number of horizontal and vertical patches per image.
              Returned when `images` is not `None`.
        """
        output_kwargs = self._merge_kwargs(Siglip2ProcessorKwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, **kwargs)
        if text is None and images is None:
            raise ValueError('You have to specify either text or images. Both cannot be none.')
        if text is not None:
            encoding = self.tokenizer(text, **output_kwargs['text_kwargs'])
        if images is not None:
            image_features = self.image_processor(images, **output_kwargs['images_kwargs'])
        if text is not None and images is not None:
            encoding.update(image_features)
            return encoding
        elif text is not None:
            return encoding
        else:
            return_tensors = output_kwargs['common_kwargs']['return_tensors']
            return BatchFeature(data=dict(**image_features), tensor_type=return_tensors)

    def decode(self, *args, **kwargs):
        """
        This method forwards all its arguments to Siglip2Tokenizer's [`~PreTrainedTokenizer.decode`]. Please refer to
        the docstring of this method for more information.
        """
        return self.tokenizer.decode(*args, **kwargs)

    def batch_decode(self, *args, **kwargs):
        """
        This method forwards all its arguments to Siglip2Tokenizer's [`~PreTrainedTokenizer.batch_decode`]. Please
        refer to the docstring of this method for more information.
        """
        return self.tokenizer.batch_decode(*args, **kwargs)

    @property
    def model_input_names(self):
        tokenizer_input_names = self.tokenizer.model_input_names
        image_processor_input_names = self.image_processor.model_input_names
        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))

Constructs a Siglip2 processor which wraps a Siglip2 image processor and a Gemma tokenizer into a single processor.

[`Siglip2Processor`] offers all the functionalities of [`Siglip2ImageProcessor`] and [`GemmaTokenizerFast`]. See the
[`~Siglip2Processor.__call__`] and [`~Siglip2Processor.decode`] for more information.

Args:
    image_processor ([`Siglip2ImageProcessor`]):
        The image processor is a required input.
    tokenizer ([`GemmaTokenizerFast`]):
        The tokenizer is a required input.


def call(self, hidden_states: tf.Tensor, attention_mask: tf.Tensor, causal_attention_mask: tf.Tensor, output_attentions: bool, training: bool=False) -> Tuple[tf.Tensor]:
    residual = hidden_states
    hidden_states = self.layer_norm1(inputs=hidden_states)
    attention_outputs = self.self_attn(hidden_states=hidden_states, attention_mask=attention_mask, causal_attention_mask=causal_attention_mask, output_attentions=output_attentions, training=training)
    hidden_states = attention_outputs[0]
    hidden_states = residual + hidden_states
    residual = hidden_states
    hidden_states = self.layer_norm2(inputs=hidden_states)
    hidden_states = self.mlp(hidden_states=hidden_states)
    hidden_states = residual + hidden_states
    outputs = (hidden_states,) + attention_outputs[1:]
    return outputs

Args:
    hidden_states (`tf.Tensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
    attention_mask (`tf.Tensor`): attention mask of size
        `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
    causal_attention_mask (`tf.Tensor`): causal attention mask of size
        `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
    output_attentions (`bool`):
        Whether or not to return the attentions tensors of all attention layers. See `outputs` under returned
        tensors for more detail.


def crowding_distance_sort(frontier: List[pg.DNA]) -> List[pg.DNA]:
    if len(frontier) <= 1:
        return frontier
    individual_num = len(frontier)
    objective_num = len(base.get_fitness(frontier[0]))
    distances = [0.0] * individual_num
    dist = [list(range(individual_num)) for i in range(objective_num)]
    for i in range(objective_num):
        dist[i] = sorted(dist[i], key=lambda idx: base.get_fitness(frontier[idx])[i])
        max_value = base.get_fitness(frontier[dist[i][individual_num - 1]])[i]
        min_value = base.get_fitness(frontier[dist[i][0]])[i]
        for j in range(individual_num):
            if j == 0 or j == individual_num - 1:
                distances[dist[i][j]] = objective_num
            elif max_value > min_value:
                distances[dist[i][j]] += (base.get_fitness(frontier[dist[i][j + 1]])[i] - base.get_fitness(frontier[dist[i][j - 1]])[i]) / (max_value - min_value)
    idx_arr = list(range(individual_num))
    idx_arr = sorted(idx_arr, key=lambda idx: distances[idx], reverse=True)
    return [frontier[idx_arr[i]] for i in range(individual_num)]

Algorithm crowding-distance-assignment implementation.

Check section III B in the original paper.

Args:
  frontier: A list of Individual that need to be sorted.

Returns:
  sorted list of the original list.


def banded_triangular_solve(bands, rhs, lower=True, adjoint=False, name=None):
    with ops.name_scope(name, 'banded_triangular_solve', [bands, rhs]):
        return gen_linalg_ops.banded_triangular_solve(bands, rhs, lower=lower, adjoint=adjoint)

Solve triangular systems of equations with a banded solver.

`bands` is a tensor of shape `[..., K, M]`, where `K` represents the number
of bands stored. This corresponds to a batch of `M` by `M` matrices, whose
`K` subdiagonals (when `lower` is `True`) are stored.

This operator broadcasts the batch dimensions of `bands` and the batch
dimensions of `rhs`.

Examples:

Storing 2 bands of a 3x3 matrix.
Note that first element in the second row is ignored due to
the 'LEFT_RIGHT' padding.

>>> x = [[2., 3., 4.], [1., 2., 3.]]
>>> x2 = [[2., 3., 4.], [10000., 2., 3.]]
>>> y = tf.zeros([3, 3])
>>> z = tf.linalg.set_diag(y, x, align='LEFT_RIGHT', k=(-1, 0))
>>> z
<tf.Tensor: shape=(3, 3), dtype=float32, numpy=
array([[2., 0., 0.],
       [2., 3., 0.],
       [0., 3., 4.]], dtype=float32)>
>>> soln = tf.linalg.banded_triangular_solve(x, tf.ones([3, 1]))
>>> soln
<tf.Tensor: shape=(3, 1), dtype=float32, numpy=
array([[0.5 ],
       [0.  ],
       [0.25]], dtype=float32)>
>>> are_equal = soln == tf.linalg.banded_triangular_solve(x2, tf.ones([3, 1]))
>>> print(tf.reduce_all(are_equal).numpy())
True
>>> are_equal = soln == tf.linalg.triangular_solve(z, tf.ones([3, 1]))
>>> print(tf.reduce_all(are_equal).numpy())
True

Storing 2 superdiagonals of a 4x4 matrix. Because of the 'LEFT_RIGHT' padding
the last element of the first row is ignored.

>>> x = [[2., 3., 4., 5.], [-1., -2., -3., -4.]]
>>> y = tf.zeros([4, 4])
>>> z = tf.linalg.set_diag(y, x, align='LEFT_RIGHT', k=(0, 1))
>>> z
<tf.Tensor: shape=(4, 4), dtype=float32, numpy=
array([[-1.,  2.,  0.,  0.],
       [ 0., -2.,  3.,  0.],
       [ 0.,  0., -3.,  4.],
       [ 0.,  0., -0., -4.]], dtype=float32)>
>>> soln = tf.linalg.banded_triangular_solve(x, tf.ones([4, 1]), lower=False)
>>> soln
<tf.Tensor: shape=(4, 1), dtype=float32, numpy=
array([[-4.       ],
       [-1.5      ],
       [-0.6666667],
       [-0.25     ]], dtype=float32)>
>>> are_equal = (soln == tf.linalg.triangular_solve(
...   z, tf.ones([4, 1]), lower=False))
>>> print(tf.reduce_all(are_equal).numpy())
True

Args:
  bands: A `Tensor` describing the bands of the left hand side, with shape
    `[..., K, M]`. The `K` rows correspond to the diagonal to the `K - 1`-th
    diagonal (the diagonal is the top row) when `lower` is `True` and
    otherwise the `K - 1`-th superdiagonal to the diagonal (the diagonal is
    the bottom row) when `lower` is `False`. The bands are stored with
    'LEFT_RIGHT' alignment, where the superdiagonals are padded on the right
    and subdiagonals are padded on the left. This is the alignment cuSPARSE
    uses.  See  `tf.linalg.set_diag` for more details.
  rhs: A `Tensor` of shape [..., M] or [..., M, N] and with the same dtype as
    `diagonals`. Note that if the shape of `rhs` and/or `diags` isn't known
    statically, `rhs` will be treated as a matrix rather than a vector.
  lower: An optional `bool`. Defaults to `True`. Boolean indicating whether
    `bands` represents a lower or upper triangular matrix.
  adjoint: An optional `bool`. Defaults to `False`. Boolean indicating whether
    to solve with the matrix's block-wise adjoint.
  name:  A name to give this `Op` (optional).

Returns:
  A `Tensor` of shape [..., M] or [..., M, N] containing the solutions.


def compute_output_shape(self, input_shape):
    if context.executing_eagerly():
        self._maybe_build(input_shape)
        with ops.get_default_graph().as_default():
            graph = func_graph.FuncGraph('graph')
            with graph.as_default():
                input_shape = tf_utils.convert_shapes(input_shape, to_tuples=False)
                inputs = nest.map_structure(base_layer_utils.generate_placeholders_from_shape, input_shape)
                try:
                    outputs = self(inputs, training=False)
                except TypeError as e:
                    raise NotImplementedError("We could not automatically infer the static shape of the layer's output. Please implement the `compute_output_shape` method on your layer (%s)." % self.__class__.__name__) from e
        return nest.map_structure(lambda t: t.shape, outputs)
    raise NotImplementedError

Computes the output shape of the layer.

If the layer has not been built, this method will call `build` on the
layer. This assumes that the layer will later be used with inputs that
match the input shape provided here.

Args:
    input_shape: Shape tuple (tuple of integers)
        or list of shape tuples (one per output tensor of the layer).
        Shape tuples can include None for free dimensions,
        instead of an integer.

Returns:
    An input shape tuple.


def num_special_tokens_to_add(self, pair: bool=False) -> int:
    return self._tokenizer.num_special_tokens_to_add(pair)

Returns the number of added tokens when encoding a sequence with special tokens.

<Tip>

This encodes a dummy input and checks the number of added tokens, and is therefore not efficient. Do not put
this inside your training loop.

</Tip>

Args:
    pair (`bool`, *optional*, defaults to `False`):
        Whether the number of added tokens should be computed in the case of a sequence pair or a single
        sequence.

Returns:
    `int`: Number of special tokens added to sequences.


class FocalNetMaskedImageModelingOutput(ModelOutput):
    loss: Optional[torch.FloatTensor] = None
    reconstruction: Optional[torch.FloatTensor] = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    reshaped_hidden_states: Optional[Tuple[torch.FloatTensor]] = None

FocalNet masked image model outputs.

Args:
    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `bool_masked_pos` is provided):
        Masked image modeling (MLM) loss.
    reconstruction (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
        Reconstructed pixel values.
    hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
        shape `(batch_size, sequence_length, hidden_size)`.

        Hidden-states of the model at the output of each layer plus the initial embedding outputs.
    reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
        shape `(batch_size, hidden_size, height, width)`.

        Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
        include the spatial dimensions.


def words_string(fake: Faker, n: int) -> str:
    return ' '.join(fake.words(n))

Provide Faker words as a joined string.

Args:
    * fake: Faker instance
    * n: number of words

Returns:
    * string of n words joined by spaces


def __eq__(self, other: Any) -> bool:
    if self is other:
        return True
    if isinstance(other, MissingValue):
        return self._value_spec == other.value_spec
    return MISSING_VALUE == other

Operator ==.

NOTE: `MissingValue(value_spec) and `utils.MissingValue` are
considered equal, but `MissingValue(value_spec1)` and
`MissingValue(value_spec2)` are considered different. That being said,
the 'eq' operation is not transitive.

However in practice this is not a problem, since user always compare
against `schema.MISSING_VALUE` which is `utils.MissingValue`.
Therefore the `__hash__` function returns the same value with
`utils.MissingValue`.

Args:
  other: the value to compare against.

Returns:
  True if the other value is a general MissingValue or MissingValue of the
    same value spec.


def model_call_event(self) -> asyncio.Event:
    return self._model_call_event

Returns an event that is set when the wrapped processor has all parts.

The event is set when the wrapped processor has all the input parts and
is about to start generating the output.

The event starts in a cleared state when the first part of the input
stream is yielded. It is also cleared at the end of the wrappedprocessor,
when all the output parts have been yielded.

Its default value is unset and this event is set only for a short time
during the call.

Returns:
  An event that is set when the model call is started, that is when all the
  input parts have been sent to the wrapped processor.


class Pager(object):
    HELP_TEXT = '\n  Simple pager commands:\n\n    b, ^B, <PAGE-UP>, <LEFT-ARROW>\n      Back one page.\n    f, ^F, <SPACE>, <PAGE-DOWN>, <RIGHT-ARROW>\n      Forward one page. Does not quit if there are no more lines.\n    g, <HOME>\n      Back to the first page.\n    <number>g\n      Go to <number> lines from the top.\n    G, <END>\n      Forward to the last page.\n    <number>G\n      Go to <number> lines from the bottom.\n    h\n      Print pager command help.\n    j, +, <DOWN-ARROW>\n      Forward one line.\n    k, -, <UP-ARROW>\n      Back one line.\n    /pattern\n      Forward search for pattern.\n    ?pattern\n      Backward search for pattern.\n    n\n      Repeat current search.\n    N\n      Repeat current search in the opposite direction.\n    q, Q, ^C, ^D, ^Z\n      Quit return to the caller.\n    any other character\n      Prompt again.\n\n  Hit any key to continue:'
    PREV_POS_NXT_REPRINT = (-1, -1)

    def __init__(self, contents, out=None, prompt=None):
        """Constructor.

    Args:
      contents: The entire contents of the text lines to page.
      out: The output stream, log.out (effectively) if None.
      prompt: The page break prompt, a default prompt is used if None..
    """
        self._contents = contents
        self._out = out or sys.stdout
        self._search_pattern = None
        self._search_direction = None
        self.prev_pos, self.prev_nxt = self.PREV_POS_NXT_REPRINT
        self._attr = console_attr.GetConsoleAttr()
        self._width, self._height = self._attr.GetTermSize()
        if not prompt:
            prompt = '{bold}--({{percent}}%)--{normal}'.format(bold=self._attr.GetFontCode(bold=True), normal=self._attr.GetFontCode())
        self._clear = '\r{0}\r'.format(' ' * (self._attr.DisplayWidth(prompt) - 6))
        self._prompt = prompt
        self._lines = []
        for line in contents.splitlines():
            self._lines += self._attr.SplitLine(line, self._width)

    def _Write(self, s):
        """Mockable helper that writes s to self._out."""
        self._out.write(s)

    def _GetSearchCommand(self, c):
        """Consumes a search command and returns the equivalent pager command.

    The search pattern is an RE that is pre-compiled and cached for subsequent
    /<newline>, ?<newline>, n, or N commands.

    Args:
      c: The search command char.

    Returns:
      The pager command char.
    """
        self._Write(c)
        buf = ''
        while True:
            p = self._attr.GetRawKey()
            if p in (None, '\n', '\r') or len(p) != 1:
                break
            self._Write(p)
            buf += p
        self._Write('\r' + ' ' * len(buf) + '\r')
        if buf:
            try:
                self._search_pattern = re.compile(buf)
            except re.error:
                self._search_pattern = None
                return ''
        self._search_direction = 'n' if c == '/' else 'N'
        return 'n'

    def _Help(self):
        """Print command help and wait for any character to continue."""
        clear = self._height - (len(self.HELP_TEXT) - len(self.HELP_TEXT.replace('\n', '')))
        if clear > 0:
            self._Write('\n' * clear)
        self._Write(self.HELP_TEXT)
        self._attr.GetRawKey()
        self._Write('\n')

    def Run(self):
        """Run the pager."""
        if len(self._lines) <= self._height:
            self._Write(self._contents)
            return
        reset_prev_values = True
        self._height -= 1
        pos = 0
        while pos < len(self._lines):
            nxt = pos + self._height
            if nxt > len(self._lines):
                nxt = len(self._lines)
                pos = nxt - self._height
            if self.prev_pos < pos < self.prev_nxt:
                self._Write('\n'.join(self._lines[self.prev_nxt:nxt]) + '\n')
            elif pos != self.prev_pos and nxt != self.prev_nxt:
                self._Write('\n'.join(self._lines[pos:nxt]) + '\n')
            percent = self._prompt.format(percent=100 * nxt // len(self._lines))
            digits = ''
            while True:
                if reset_prev_values:
                    self.prev_pos, self.prev_nxt = (pos, nxt)
                    reset_prev_values = False
                self._Write(percent)
                c = self._attr.GetRawKey()
                self._Write(self._clear)
                if c in (None, 'q', 'Q', '\x03', '\x1b'):
                    return
                elif c in ('/', '?'):
                    c = self._GetSearchCommand(c)
                elif c.isdigit():
                    digits += c
                    continue
                if digits:
                    count = int(digits)
                    digits = ''
                else:
                    count = 0
                if c in ('<PAGE-UP>', '<LEFT-ARROW>', 'b', '\x02'):
                    nxt = pos - self._height
                    if nxt < 0:
                        nxt = 0
                elif c in ('<PAGE-DOWN>', '<RIGHT-ARROW>', 'f', '\x06', ' '):
                    if nxt >= len(self._lines):
                        continue
                    nxt = pos + self._height
                    if nxt >= len(self._lines):
                        nxt = pos
                elif c in ('<HOME>', 'g'):
                    nxt = count - 1
                    if nxt > len(self._lines) - self._height:
                        nxt = len(self._lines) - self._height
                    if nxt < 0:
                        nxt = 0
                elif c in ('<END>', 'G'):
                    nxt = len(self._lines) - count
                    if nxt > len(self._lines) - self._height:
                        nxt = len(self._lines) - self._height
                    if nxt < 0:
                        nxt = 0
                elif c == 'h':
                    self._Help()
                    self.prev_pos, self.prev_nxt = self.PREV_POS_NXT_REPRINT
                    nxt = pos
                    break
                elif c in ('<DOWN-ARROW>', 'j', '+', '\n', '\r'):
                    if nxt >= len(self._lines):
                        continue
                    nxt = pos + 1
                    if nxt >= len(self._lines):
                        nxt = pos
                elif c in ('<UP-ARROW>', 'k', '-'):
                    nxt = pos - 1
                    if nxt < 0:
                        nxt = 0
                elif c in ('n', 'N'):
                    if not self._search_pattern:
                        continue
                    nxt = pos
                    i = pos
                    direction = 1 if c == self._search_direction else -1
                    while True:
                        i += direction
                        if i < 0 or i >= len(self._lines):
                            break
                        if self._search_pattern.search(self._lines[i]):
                            nxt = i
                            break
                else:
                    continue
                if nxt != pos:
                    reset_prev_values = True
                    break
            pos = nxt

A simple console text pager.

This pager requires the entire contents to be available. The contents are
written one page of lines at a time. The prompt is written after each page of
lines. A one character response is expected. See HELP_TEXT below for more
info.

The contents are written as is. For example, ANSI control codes will be in
effect. This is different from pagers like more(1) which is ANSI control code
agnostic and miscalculates line lengths, and less(1) which displays control
character names by default.

Attributes:
  _attr: The current ConsoleAttr handle.
  _clear: A string that clears the prompt when written to _out.
  _contents: The entire contents of the text lines to page.
  _height: The terminal height in characters.
  _out: The output stream, log.out (effectively) if None.
  _prompt: The page break prompt.
  _search_direction: The search direction command, n:forward, N:reverse.
  _search_pattern: The current forward/reverse search compiled RE.
  _width: The termonal width in characters.


def rot90(array, k=1, axes=(0, 1)):
    array = convert_to_tensor(array)
    if array.ndim < 2:
        raise ValueError(f'Input array must have at least 2 dimensions. Received: array.ndim={array.ndim}')
    if len(axes) != 2 or axes[0] == axes[1]:
        raise ValueError(f'Invalid axes: {axes}. Axes must be a tuple of two different dimensions.')
    axes = tuple((axis if axis >= 0 else array.ndim + axis for axis in axes))
    if not builtins.all((0 <= axis < array.ndim for axis in axes)):
        raise ValueError(f'Invalid axes {axes} for tensor with {array.ndim} dimensions')
    rotated = torch.rot90(array, k=k, dims=axes)
    if isinstance(array, np.ndarray):
        rotated = rotated.cpu().numpy()
    return rotated

Rotate an array by 90 degrees in the specified plane using PyTorch.

Args:
    array: Input tensor
    k: Number of 90-degree rotations (default=1)
    axes: Tuple of two axes that define the
        plane of rotation (defaults to `(0, 1)`).

Returns:
    Rotated tensor


def bidirectional_dynamic_rnn(cell_fw, cell_bw, inputs, sequence_length=None, initial_state_fw=None, initial_state_bw=None, dtype=None, parallel_iterations=None, swap_memory=False, time_major=False, scope=None):
    rnn_cell_impl.assert_like_rnncell('cell_fw', cell_fw)
    rnn_cell_impl.assert_like_rnncell('cell_bw', cell_bw)
    with vs.variable_scope(scope or 'bidirectional_rnn'):
        with vs.variable_scope('fw') as fw_scope:
            output_fw, output_state_fw = dynamic_rnn(cell=cell_fw, inputs=inputs, sequence_length=sequence_length, initial_state=initial_state_fw, dtype=dtype, parallel_iterations=parallel_iterations, swap_memory=swap_memory, time_major=time_major, scope=fw_scope)
        if not time_major:
            time_axis = 1
            batch_axis = 0
        else:
            time_axis = 0
            batch_axis = 1

        def _reverse(input_, seq_lengths, seq_axis, batch_axis):
            if seq_lengths is not None:
                return array_ops.reverse_sequence(input=input_, seq_lengths=seq_lengths, seq_axis=seq_axis, batch_axis=batch_axis)
            else:
                return array_ops.reverse(input_, axis=[seq_axis])
        with vs.variable_scope('bw') as bw_scope:

            def _map_reverse(inp):
                return _reverse(inp, seq_lengths=sequence_length, seq_axis=time_axis, batch_axis=batch_axis)
            inputs_reverse = nest.map_structure(_map_reverse, inputs)
            tmp, output_state_bw = dynamic_rnn(cell=cell_bw, inputs=inputs_reverse, sequence_length=sequence_length, initial_state=initial_state_bw, dtype=dtype, parallel_iterations=parallel_iterations, swap_memory=swap_memory, time_major=time_major, scope=bw_scope)
    output_bw = _reverse(tmp, seq_lengths=sequence_length, seq_axis=time_axis, batch_axis=batch_axis)
    outputs = (output_fw, output_bw)
    output_states = (output_state_fw, output_state_bw)
    return (outputs, output_states)

Creates a dynamic version of bidirectional recurrent neural network.

Takes input and builds independent forward and backward RNNs. The input_size
of forward and backward cell must match. The initial state for both directions
is zero by default (but can be set optionally) and no intermediate states are
ever returned -- the network is fully unrolled for the given (passed in)
length(s) of the sequence(s) or completely unrolled if length(s) is not
given.

Args:
  cell_fw: An instance of RNNCell, to be used for forward direction.
  cell_bw: An instance of RNNCell, to be used for backward direction.
  inputs: The RNN inputs.
    If time_major == False (default), this must be a tensor of shape:
      `[batch_size, max_time, ...]`, or a nested tuple of such elements.
    If time_major == True, this must be a tensor of shape: `[max_time,
      batch_size, ...]`, or a nested tuple of such elements.
  sequence_length: (optional) An int32/int64 vector, size `[batch_size]`,
    containing the actual lengths for each of the sequences in the batch. If
    not provided, all batch entries are assumed to be full sequences; and time
    reversal is applied from time `0` to `max_time` for each sequence.
  initial_state_fw: (optional) An initial state for the forward RNN. This must
    be a tensor of appropriate type and shape `[batch_size,
    cell_fw.state_size]`. If `cell_fw.state_size` is a tuple, this should be a
    tuple of tensors having shapes `[batch_size, s] for s in
    cell_fw.state_size`.
  initial_state_bw: (optional) Same as for `initial_state_fw`, but using the
    corresponding properties of `cell_bw`.
  dtype: (optional) The data type for the initial states and expected output.
    Required if initial_states are not provided or RNN states have a
    heterogeneous dtype.
  parallel_iterations: (Default: 32).  The number of iterations to run in
    parallel.  Those operations which do not have any temporal dependency and
    can be run in parallel, will be.  This parameter trades off time for
    space.  Values >> 1 use more memory but take less time, while smaller
    values use less memory but computations take longer.
  swap_memory: Transparently swap the tensors produced in forward inference
    but needed for back prop from GPU to CPU.  This allows training RNNs which
    would typically not fit on a single GPU, with very minimal (or no)
    performance penalty.
  time_major: The shape format of the `inputs` and `outputs` Tensors. If true,
    these `Tensors` must be shaped `[max_time, batch_size, depth]`. If false,
    these `Tensors` must be shaped `[batch_size, max_time, depth]`. Using
    `time_major = True` is a bit more efficient because it avoids transposes
    at the beginning and end of the RNN calculation.  However, most TensorFlow
    data is batch-major, so by default this function accepts input and emits
    output in batch-major form.
  scope: VariableScope for the created subgraph; defaults to
    "bidirectional_rnn"

Returns:
  A tuple (outputs, output_states) where:
    outputs: A tuple (output_fw, output_bw) containing the forward and
      the backward rnn output `Tensor`.
      If time_major == False (default),
        output_fw will be a `Tensor` shaped:
        `[batch_size, max_time, cell_fw.output_size]`
        and output_bw will be a `Tensor` shaped:
        `[batch_size, max_time, cell_bw.output_size]`.
      If time_major == True,
        output_fw will be a `Tensor` shaped:
        `[max_time, batch_size, cell_fw.output_size]`
        and output_bw will be a `Tensor` shaped:
        `[max_time, batch_size, cell_bw.output_size]`.
      It returns a tuple instead of a single concatenated `Tensor`, unlike
      in the `bidirectional_rnn`. If the concatenated one is preferred,
      the forward and backward outputs can be concatenated as
      `tf.concat(outputs, 2)`.
    output_states: A tuple (output_state_fw, output_state_bw) containing
      the forward and the backward final states of bidirectional rnn.

Raises:
  TypeError: If `cell_fw` or `cell_bw` is not an instance of `RNNCell`.


The operation that failed, if known.

*N.B.* If the failed op was synthesized at runtime, e.g. a `Send`
or `Recv` op, there will be no corresponding
`tf.Operation`
object.  In that case, this will return `None`, and you should
instead use the `tf.errors.OpError.node_def` to
discover information about the op.

Returns:
  The `Operation` that failed, or None.


def _pytd_constant_to_value(self, pyval: pytd.Node, subst, get_node):
    if isinstance(pyval, pytd.LateType):
        actual = self._load_late_type(pyval)
        return self._constant_to_value(actual, subst, get_node)
    elif isinstance(pyval, pytd.TypeDeclUnit):
        return self._create_module(pyval)
    elif isinstance(pyval, pytd.Module):
        mod = self.ctx.loader.import_name(pyval.module_name)
        return self._create_module(mod)
    elif isinstance(pyval, pytd.Class):
        return self._pytd_class_to_value(pyval, get_node())
    elif isinstance(pyval, pytd.Function):
        f = self.convert_pytd_function(pyval)
        f.is_abstract = pyval.is_abstract
        return f
    elif isinstance(pyval, pytd.ClassType):
        if pyval.cls:
            cls = pyval.cls
        else:
            cls = self.ctx.loader.lookup_pytd(*pyval.name.split('.', 1))
            assert isinstance(cls, pytd.Class)
        return self.constant_to_value(cls, subst)
    elif isinstance(pyval, pytd.NothingType):
        return self.empty
    elif isinstance(pyval, pytd.AnythingType):
        return self.unsolvable
    elif isinstance(pyval, pytd.Constant) and isinstance(pyval.type, pytd.AnythingType):
        return self.unsolvable
    elif isinstance(pyval, pytd.Constant) and isinstance(pyval.type, pytd.GenericType) and (pyval.type.name == 'builtins.type'):
        param, = pyval.type.parameters
        return self.constant_to_value(param, subst)
    elif isinstance(pyval, pytd.UnionType):
        options = [self.constant_to_value(t, subst) for t in pyval.type_list]
        if len(options) > 1:
            return abstract.Union(options, self.ctx)
        else:
            return options[0]
    elif isinstance(pyval, (pytd.TypeParameter, pytd.ParamSpec)):
        constraints = tuple((self.constant_to_value(c, {}) for c in pyval.constraints))
        bound = pyval.bound and self.constant_to_value(pyval.bound, {})
        if isinstance(pyval, pytd.ParamSpec):
            cls = abstract.ParamSpec
        else:
            cls = abstract.TypeParameter
        return cls(pyval.name, self.ctx, constraints=constraints, bound=bound, scope=pyval.scope)
    elif isinstance(pyval, (pytd.ParamSpecArgs, pytd.ParamSpecKwargs)):
        return self.unsolvable
    elif isinstance(pyval, pytd.Concatenate):
        params = [self.constant_to_value(p, subst) for p in pyval.parameters]
        return abstract.Concatenate(params, self.ctx)
    elif isinstance(pyval, pytd.GenericType) and pyval.name == 'typing.ClassVar':
        param, = pyval.parameters
        return self.constant_to_value(param, subst)
    elif isinstance(pyval, pytd.GenericType):
        return self._pytd_generic_type_to_value(pyval, subst, get_node)
    elif isinstance(pyval, pytd.Literal):
        value = self._get_literal_value(pyval.value, subst)
        return abstract.LiteralClass(value, self.ctx)
    elif isinstance(pyval, pytd.Annotated):
        typ = self.constant_to_value(pyval.base_type, subst)
        return self._apply_metadata_annotations(typ, pyval.annotations)
    else:
        raise NotImplementedError(f"Can't convert pytd constant {type(pyval)} {pyval!r}")

Convert a pytd type to an abstract value.

Args:
  pyval: The PyTD value to convert.
  subst: The current type parameters.
  get_node: A getter function for the current node.

Returns:
  A Value that represents the constant, or None if we couldn't convert.
Raises:
  NotImplementedError: If we don't know how to convert a value.
  TypeParameterError: If we can't find a substitution for a type parameter.


def _print(self, *args):

    def _format(name, arr):
        """Prints compatibility check results with a format.

      Args:
        name: String that is the title representing list `arr`.
        arr: List of items to be printed in a certain format.
      """
        title = '### All Compatibility %s ###' % str(name)
        tlen = len(title)
        print('-' * tlen)
        print(title)
        print('-' * tlen)
        print(' Total # of %s: %s\n' % (str(name), str(len(arr))))
        if arr:
            for item in arr:
                detail = ''
                if isinstance(item[1], list):
                    for itm in item[1]:
                        detail += str(itm) + ', '
                    detail = detail[:-2]
                else:
                    detail = str(item[1])
                print("  %s ('%s')\n" % (str(item[0]), detail))
        else:
            print('  No %s' % name)
        print('\n')
    for p_item in args:
        if p_item == 'failures':
            _format('Failures', self.failures)
        elif p_item == 'successes':
            _format('Successes', self.successes)
        elif p_item == 'failure_msgs':
            _format('Failure Messages', self.error_msg)
        elif p_item == 'warning_msgs':
            _format('Warning Messages', self.warning_msg)
        else:
            raise Exception('[Error] Wrong input provided for %s.' % _get_func_name())

Prints compatibility check status and failure or warning messages.

Prints to console without using `logging`.

Args:
  *args: String(s) that is one of:
          [`failures`,       # all failures
           `successes`,      # all successes
           `failure_msgs`,   # failure message(s) recorded upon failure(s)
           `warning_msgs`]   # warning message(s) recorded upon warning(s)
Raises:
  Exception: If *args not in:
               [`failures`, `successes`, `failure_msgs`, `warning_msg`]


def min(x, axis=None, keepdims=False):
    return math_ops.reduce_min(x, axis, keepdims)

Minimum value in a tensor.

Args:
    x: A tensor or variable.
    axis: An integer, the axis to find minimum values.
    keepdims: A boolean, whether to keep the dimensions or not.
        If `keepdims` is `False`, the rank of the tensor is reduced
        by 1. If `keepdims` is `True`,
        the reduced dimension is retained with length 1.

Returns:
    A tensor with minimum values of `x`.


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

Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
adding special tokens. A Big Bird sequence has the following format:

- single sequence: `[CLS] X [SEP]`
- pair of sequences: `[CLS] A [SEP] B [SEP]`

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

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


def forward(self, hidden_states):
    forwarded_states = self.mlp(hidden_states)
    output = hidden_states + self.norm(forwarded_states)
    return output

Args:
    hidden_states (`torch.Tensor`) :
        [num_groups, tokens_per_group, hidden_dim] inputs to send to experts.
Returns:
    torch.Tensor[num_groups, tokens_per_group, hidden_dim]


def get_completions(self, context_word, prefix):
    if context_word not in self._comp_dict:
        return (None, None)
    comp_items = self._comp_dict[context_word]
    comp_items = sorted([item for item in comp_items if item.startswith(prefix)])
    return (comp_items, self._common_prefix(comp_items))

Get the tab completions given a context word and a prefix.

Args:
  context_word: The context word.
  prefix: The prefix of the incomplete word.

Returns:
  (1) None if no registered context matches the context_word.
      A list of str for the matching completion items. Can be an empty list
      of a matching context exists, but no completion item matches the
      prefix.
  (2) Common prefix of all the words in the first return value. If the
      first return value is None, this return value will be None, too. If
      the first return value is not None, i.e., a list, this return value
      will be a str, which can be an empty str if there is no common
      prefix among the items of the list.


def quantize_and_dequantize(input, input_min, input_max, signed_input=True, num_bits=8, range_given=False, round_mode='HALF_TO_EVEN', name=None, narrow_range=False, axis=None):
    if axis is None:
        axis = -1
    elif axis < 0:
        if input.shape.ndims is None:
            raise ValueError('input should have known rank to use negative axis.')
        axis %= input.shape.ndims
    return gen_array_ops.quantize_and_dequantize_v2(input, input_min=input_min, input_max=input_max, signed_input=signed_input, num_bits=num_bits, range_given=range_given, round_mode=round_mode, narrow_range=narrow_range, axis=axis, name=name)

Quantizes then dequantizes a tensor.

Args:
  input: A `Tensor` to quantize and dequantize.
  input_min: If range_given=True, the minimum input value, that needs to be
    represented in the quantized representation. If axis is specified, this
    should be a vector of minimum values for each slice along axis.
  input_max: If range_given=True, the maximum input value that needs to be
    represented in the quantized representation. If axis is specified, this
    should be a vector of maximum values for each slice along axis.
  signed_input: True if the quantization is signed or unsigned.
  num_bits: The bitwidth of the quantization.
  range_given: If true use `input_min` and `input_max` for the range of the
    input, otherwise determine min and max from the input `Tensor`.
  round_mode: Rounding mode when rounding from float values to quantized ones.
    one of ['HALF_TO_EVEN', 'HALF_UP']
  name: Optional name for the operation.
  narrow_range: If true, then the absolute value of the quantized minimum
    value is the same as the quantized maximum value, instead of 1 greater.
    i.e. for 8 bit quantization, the minimum value is -127 instead of -128.
  axis: Integer. If specified, refers to a dimension of the input tensor, such
    that quantization will be per slice along that dimension.

Returns:
  A `Tensor`. Each element is the result of quantizing and dequantizing the
  corresponding element of `input`.


def histogram(name, data, step=None, buckets=None, description=None):
    try:
        from tensorboard.summary.v2 import histogram as histogram_v2
    except ImportError as exc:
        raise TBNotInstalledError('tf.summary.histogram') from exc
    return histogram_v2(name=name, data=data, step=step, buckets=buckets, description=description)

Write a histogram summary.

See also `tf.summary.scalar`, `tf.summary.SummaryWriter`.

Writes a histogram to the current default summary writer, for later analysis
in TensorBoard's 'Histograms' and 'Distributions' dashboards (data written
using this API will appear in both places). Like `tf.summary.scalar` points,
each histogram is associated with a `step` and a `name`. All the histograms
with the same `name` constitute a time series of histograms.

The histogram is calculated over all the elements of the given `Tensor`
without regard to its shape or rank.

This example writes 2 histograms:

```python
w = tf.summary.create_file_writer('test/logs')
with w.as_default():
    tf.summary.histogram("activations", tf.random.uniform([100, 50]), step=0)
    tf.summary.histogram("initial_weights", tf.random.normal([1000]), step=0)
```

A common use case is to examine the changing activation patterns (or lack
thereof) at specific layers in a neural network, over time.

```python
w = tf.summary.create_file_writer('test/logs')
with w.as_default():
for step in range(100):
    # Generate fake "activations".
    activations = [
        tf.random.normal([1000], mean=step, stddev=1),
        tf.random.normal([1000], mean=step, stddev=10),
        tf.random.normal([1000], mean=step, stddev=100),
    ]

    tf.summary.histogram("layer1/activate", activations[0], step=step)
    tf.summary.histogram("layer2/activate", activations[1], step=step)
    tf.summary.histogram("layer3/activate", activations[2], step=step)
```

Arguments:
  name: A name for this summary. The summary tag used for TensorBoard will be
    this name prefixed by any active name scopes.
  data: A `Tensor` of any shape. The histogram is computed over its elements,
    which must be castable to `float64`.
  step: Explicit `int64`-castable monotonic step value for this summary. If
    omitted, this defaults to `tf.summary.experimental.get_step()`, which must
    not be None.
  buckets: Optional positive `int`. The output will have this many buckets,
    except in two edge cases. If there is no data, then there are no buckets.
    If there is data but all points have the same value, then all buckets'
    left and right endpoints are the same and only the last bucket has nonzero
    count. Defaults to 30 if not specified.
  description: Optional long-form description for this summary, as a constant
    `str`. Markdown is supported. Defaults to empty.

Returns:
  True on success, or false if no summary was emitted because no default
  summary writer was available.

Raises:
  ValueError: if a default writer exists, but no step was provided and
    `tf.summary.experimental.get_step()` is None.


Linear activation function (pass-through).

A "linear" activation is an identity function:
it returns the input, unmodified.

Args:
    x: Input tensor.


def get_output_shape_at(self, node_index):
    return self._get_node_attribute_at_index(node_index, 'output_shapes', 'output shape')

Retrieves the output shape(s) of a layer at a given node.

Args:
    node_index: Integer, index of the node
        from which to retrieve the attribute.
        E.g. `node_index=0` will correspond to the
        first time the layer was called.

Returns:
    A shape tuple
    (or list of shape tuples if the layer has multiple outputs).

Raises:
  RuntimeError: If called in Eager mode.


def _create_uninitialized_mirrored_tpu_replicated_variables(**kwargs):
    dtype = kwargs.get('dtype', None)
    shape = kwargs.get('shape', None)
    initial_value = kwargs.get('initial_value', None)
    if initial_value is None:
        return _create_mirrored_tpu_replicated_variables(**kwargs)
    with maybe_init_scope():
        if initial_value is not None:
            if callable(initial_value):
                initial_value = initial_value()
            initial_value = ops.convert_to_tensor(initial_value, dtype=dtype)
            kwargs['initial_value'] = initial_value
            if dtype is None:
                kwargs['dtype'] = kwargs['initial_value'].dtype
            if shape is None:
                kwargs['shape'] = kwargs['initial_value'].shape
    mirrored_replicated_var_list = []
    for replica_id in range(num_replicas):
        replicated_var_list = []
        for logic_core_id in range(num_cores_per_replica):
            with ops.device(self._tpu_devices[replica_id][logic_core_id]):
                v = uninitialized_variable_creator(**kwargs)
            replicated_var_list.append(v)
        replica_name = '{}/r:{}'.format(kwargs['name'], replica_id)
        tpu_replicated_var = tpu_replicated_variable.TPUReplicatedVariable(variables=replicated_var_list, name=replica_name)
        mirrored_replicated_var_list.append(tpu_replicated_var)
    return mirrored_replicated_var_list

Returns a list of `TPUReplicatedVariable`s.

The list consists of `num_replicas` `TPUReplicatedVariable`s and can be
used to initialize a `TPUMirroredVariable`. Each `TPUReplicatedVariable`
contains a list of `tf.Variable`s which are replicated to
`num_cores_per_replica` logical cores to enable XLA SPMD compilation.

Args:
  **kwargs: the keyword arguments for creating a variable


class GotOcr2CausalLMOutputWithPast(ModelOutput):
    loss: Optional[torch.FloatTensor] = None
    logits: Optional[torch.FloatTensor] = None
    past_key_values: Optional[List[torch.FloatTensor]] = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None
    image_hidden_states: Optional[torch.FloatTensor] = None

Base class for GotOcr2 causal language model (or autoregressive) outputs.

Args:
    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
        Language modeling loss (for next-token prediction).
    logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
        Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
    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)`)

        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
        `past_key_values` input) to speed up sequential decoding.
    hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
        one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.

        Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
    attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
        Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
        sequence_length)`.

        Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
        heads.
    image_hidden_states (`torch.FloatTensor`, *optional*):
        A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.
        image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.


class AdditiveAttention(Attention):

    def __init__(self, use_scale=True, dropout=0.0, **kwargs):
        super().__init__(use_scale=use_scale, dropout=dropout, **kwargs)

    def build(self, input_shape):
        self._validate_inputs(input_shape)
        dim = input_shape[0][-1]
        self.scale = None
        if self.use_scale:
            self.scale = self.add_weight(name='scale', shape=[dim], initializer='glorot_uniform', dtype=self.dtype, trainable=True)

    def _calculate_scores(self, query, key):
        """Calculates attention scores as a nonlinear sum of query and key.

        Args:
            query: Query tensor of shape `(batch_size, Tq, dim)`.
            key: Key tensor of shape `(batch_size, Tv, dim)`.

        Returns:
            Tensor of shape `(batch_size, Tq, Tv)`.
        """
        q_reshaped = ops.expand_dims(query, axis=-2)
        k_reshaped = ops.expand_dims(key, axis=-3)
        scale = self.scale if self.use_scale else 1.0
        return ops.sum(scale * ops.tanh(q_reshaped + k_reshaped), axis=-1)

    def get_config(self):
        base_config = super().get_config()
        del base_config['score_mode']
        return base_config

Additive attention layer, a.k.a. Bahdanau-style attention.

Inputs are a list with 2 or 3 elements:
1. A `query` tensor of shape `(batch_size, Tq, dim)`.
2. A `value` tensor of shape `(batch_size, Tv, dim)`.
3. A optional `key` tensor of shape `(batch_size, Tv, dim)`. If none
    supplied, `value` will be used as `key`.

The calculation follows the steps:
1. Calculate attention scores using `query` and `key` with shape
    `(batch_size, Tq, Tv)` as a non-linear sum
    `scores = reduce_sum(tanh(query + key), axis=-1)`.
2. Use scores to calculate a softmax distribution with shape
    `(batch_size, Tq, Tv)`.
3. Use the softmax distribution to create a linear combination of `value`
    with shape `(batch_size, Tq, dim)`.

Args:
    use_scale: If `True`, will create a scalar variable to scale the
        attention scores.
    dropout: Float between 0 and 1. Fraction of the units to drop for the
        attention scores. Defaults to `0.0`.

Call arguments:
    inputs: List of the following tensors:
        - `query`: Query tensor of shape `(batch_size, Tq, dim)`.
        - `value`: Value tensor of shape `(batch_size, Tv, dim)`.
        - `key`: Optional key tensor of shape `(batch_size, Tv, dim)`. If
            not given, will use `value` for both `key` and `value`, which is
            the most common case.
    mask: List of the following tensors:
        - `query_mask`: A boolean mask tensor of shape `(batch_size, Tq)`.
            If given, the output will be zero at the positions where
            `mask==False`.
        - `value_mask`: A boolean mask tensor of shape `(batch_size, Tv)`.
            If given, will apply the mask such that values at positions
             where `mask==False` do not contribute to the result.
    return_attention_scores: bool, it `True`, returns the attention scores
        (after masking and softmax) as an additional output argument.
    training: Python boolean indicating whether the layer should behave in
        training mode (adding dropout) or in inference mode (no dropout).
    use_causal_mask: Boolean. Set to `True` for decoder self-attention. Adds
        a mask such that position `i` cannot attend to positions `j > i`.
        This prevents the flow of information from the future towards the
        past. Defaults to `False`.

Output:
    Attention outputs of shape `(batch_size, Tq, dim)`.
    (Optional) Attention scores after masking and softmax with shape
        `(batch_size, Tq, Tv)`.


def tanh(x):
    if any_symbolic_tensors((x,)):
        return Tanh().symbolic_call(x)
    return backend.numpy.tanh(x)

Hyperbolic tangent, element-wise.

Arguments:
    x: Input tensor.

Returns:
    Output tensor of same shape as `x`.


def add_tensor_filter(self, filter_name, tensor_filter):
    if self._session_wrapper:
        self._session_wrapper.add_tensor_filter(filter_name, tensor_filter)
    else:
        self._pending_tensor_filters[filter_name] = tensor_filter

Add a tensor filter.

See doc of `LocalCLIDebugWrapperSession.add_tensor_filter()` for details.
Override default behavior to accommodate the possibility of this method
being
called prior to the initialization of the underlying
`LocalCLIDebugWrapperSession` object.

Args:
  filter_name: See doc of `LocalCLIDebugWrapperSession.add_tensor_filter()`
    for details.
  tensor_filter: See doc of
    `LocalCLIDebugWrapperSession.add_tensor_filter()` for details.


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

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

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

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

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


def temp_sqlite_database(prefix='yaml_jdbc_it_'):
    conn = cursor = None
    try:
        db_name = f'{prefix}{uuid.uuid4().hex}.db'
        conn = sqlite3.connect(db_name)
        cursor = conn.cursor()
        cursor.execute('\n      CREATE TABLE tmp_table (\n        value INTEGER PRIMARY KEY,\n        rank INTEGER\n      )\n    ')
        conn.commit()
        yield f'jdbc:sqlite:{db_name}'
    except (sqlite3.Error, Exception) as err:
        logging.error('Error interacting with temporary SQLite DB: %s', err)
        raise err
    finally:
        if cursor:
            cursor.close()
        if conn:
            conn.close()
        try:
            if os.path.exists(db_name):
                os.remove(db_name)
        except Exception as err:
            logging.error('Error deleting temporary SQLite DB: %s', err)
            raise err

Context manager to provide a temporary SQLite database via JDBC for
testing.

This function creates a temporary SQLite database file on the local
filesystem. It establishes a connection using 'sqlite3', creates a predefined
'tmp_table', and then yields a JDBC connection string suitable for use in
tests that require a generic JDBC connection (specifically configured for
SQLite in this case).

The SQLite database file is automatically cleaned up (closed and deleted)
when the context manager exits.

Args:
    prefix (str): A prefix to use for the temporary database file name.

Yields:
    str: A JDBC connection string for the temporary SQLite database.
         Example format: "jdbc:sqlite:<path_to_db_file>"

Raises:
    sqlite3.Error: If there's an error connecting to or interacting with
                   the SQLite database during setup.
    Exception: Any other exception encountered during the setup or cleanup
               process.


def default_sequence_length(self) -> int:
    return OnnxConfig.default_fixed_sequence

The default sequence length to use if no other indication

Returns:
    Integer > 0


def rematerialized_call(self, layer_call, *args, **kwargs):

    def compute_size(x):
        return math.prod([d or 1 for d in x.shape]) if isinstance(x, KerasTensor) else 0
    if self._remat_mode.mode == 'full':
        return remat.remat(layer_call)
    elif self._remat_mode.mode == 'list_of_layers' and self.name in self._remat_mode.layer_names:
        return remat.remat(layer_call)
    elif self._remat_mode.mode == 'larger_than':
        output_spec = self.compute_output_spec(*args, **kwargs)
        output_size = sum(tree.flatten(tree.map_structure(compute_size, output_spec)))
        if output_size and output_size > self._remat_mode.output_size_threshold:
            return remat.remat(layer_call)
    elif self._remat_mode.mode == 'activations':
        has_activation = hasattr(self, 'activation') and self.activation is not None
        if has_activation:

            @functools.wraps(layer_call)
            def rematerialized_activation_call_wrapper(*args, **kwargs):
                original_activation = self.activation
                self.activation = remat.remat(original_activation)
                try:
                    return layer_call(*args, **kwargs)
                finally:
                    self.activation = original_activation
            return rematerialized_activation_call_wrapper
    return layer_call

Enable rematerialization dynamically for layer's call method.

Args:
    layer_call: The original `call` method of a layer.

Returns:
    Rematerialized layer's `call` method.


def post_process_object_detection(self, outputs, threshold: float=0.5, target_sizes: Union[TensorType, List[Tuple]]=None):
    out_logits, out_bbox = (outputs.logits, outputs.pred_boxes)
    if target_sizes is not None:
        if len(out_logits) != len(target_sizes):
            raise ValueError('Make sure that you pass in as many target sizes as the batch dimension of the logits')
    prob = nn.functional.softmax(out_logits, -1)
    scores, labels = prob[..., :-1].max(-1)
    boxes = center_to_corners_format(out_bbox)
    if target_sizes is not None:
        if isinstance(target_sizes, List):
            img_h = torch.Tensor([i[0] for i in target_sizes])
            img_w = torch.Tensor([i[1] for i in target_sizes])
        else:
            img_h, img_w = target_sizes.unbind(1)
        scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1).to(boxes.device)
        boxes = boxes * scale_fct[:, None, :]
    results = []
    for s, l, b in zip(scores, labels, boxes):
        score = s[s > threshold]
        label = l[s > threshold]
        box = b[s > threshold]
        results.append({'scores': score, 'labels': label, 'boxes': box})
    return results

Converts the raw output of [`DetrForObjectDetection`] into final bounding boxes in (top_left_x, top_left_y,
bottom_right_x, bottom_right_y) format. Only supports PyTorch.

Args:
    outputs ([`DetrObjectDetectionOutput`]):
        Raw outputs of the model.
    threshold (`float`, *optional*):
        Score threshold to keep object detection predictions.
    target_sizes (`torch.Tensor` or `List[Tuple[int, int]]`, *optional*):
        Tensor of shape `(batch_size, 2)` or list of tuples (`Tuple[int, int]`) containing the target size
        `(height, width)` of each image in the batch. If unset, predictions will not be resized.
Returns:
    `List[Dict]`: A list of dictionaries, each dictionary containing the scores, labels and boxes for an image
    in the batch as predicted by the model.


def _read_file(filename):
    graph_def = graph_pb2.GraphDef()
    if not file_io.file_exists(filename):
        raise IOError(f'File {filename} does not exist.')
    with file_io.FileIO(filename, 'rb') as f:
        file_content = f.read()
    try:
        graph_def.ParseFromString(file_content)
        return graph_def
    except Exception:
        pass
    try:
        text_format.Merge(file_content, graph_def)
    except text_format.ParseError as e:
        raise IOError(f'Cannot parse file {filename}: {str(e)}.')
    return graph_def

Reads a file containing `GraphDef` and returns the protocol buffer.

Args:
  filename: `graph_def` filename including the path.

Returns:
  A `GraphDef` protocol buffer.

Raises:
  IOError: If the file doesn't exist, or cannot be successfully parsed.


def diag(x, k=0):
    if any_symbolic_tensors((x,)):
        return Diag(k=k).symbolic_call(x)
    return backend.numpy.diag(x, k=k)

Extract a diagonal or construct a diagonal array.

Args:
    x: Input tensor. If `x` is 2-D, returns the k-th diagonal of `x`.
        If `x` is 1-D, return a 2-D tensor with `x` on the k-th diagonal.
    k: The diagonal to consider. Defaults to `0`. Use `k > 0` for diagonals
        above the main diagonal, and `k < 0` for diagonals below
        the main diagonal.

Returns:
    The extracted diagonal or constructed diagonal tensor.

Examples:
>>> from keras.src import ops
>>> x = ops.arange(9).reshape((3, 3))
>>> x
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])

>>> ops.diag(x)
array([0, 4, 8])
>>> ops.diag(x, k=1)
array([1, 5])
>>> ops.diag(x, k=-1)
array([3, 7])

>>> ops.diag(ops.diag(x)))
array([[0, 0, 0],
       [0, 4, 0],
       [0, 0, 8]])


def add_to_optionally_restored(self, var):
    self._optionally_restored.append(var)

Add a variable to the list of optionally restored variables.

There are situations where certain variables should be ignored in assertions
such as assert_existing_objects_matched(). One example is that of a
checkpoint saved with train.Saver(), and restored with train.Checkpoint():
it is possible for the train.Saver() checkpoint to be missing the internal
`save_counter` variable, which we want to ignore on restore.

Args:
  var: The variable to treat as optionally restored.


def forward(self, attention_mask: Optional[torch.Tensor]=None, head_mask: Optional[torch.Tensor]=None, inputs_embeds: Optional[torch.FloatTensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[Tuple, BaseModelOutput]:
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    hidden_states = self.value_embedding(inputs_embeds)
    embed_pos = self.embed_positions(inputs_embeds.size())
    hidden_states = self.layernorm_embedding(hidden_states + embed_pos)
    hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
    if attention_mask is not None:
        attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype)
    encoder_states = () if output_hidden_states else None
    all_attentions = () if output_attentions else None
    if head_mask is not None:
        if head_mask.size()[0] != len(self.layers):
            raise ValueError(f'The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}.')
    for idx, (encoder_layer, conv_layer) in enumerate(zip(self.layers, self.conv_layers)):
        if output_hidden_states:
            encoder_states = encoder_states + (hidden_states,)
        to_drop = False
        if self.training:
            dropout_probability = torch.rand([])
            if dropout_probability < self.layerdrop:
                to_drop = True
        if to_drop:
            layer_outputs = (None, None)
        else:
            if self.gradient_checkpointing and self.training:
                layer_outputs = self._gradient_checkpointing_func(encoder_layer.__call__, hidden_states, attention_mask, head_mask[idx] if head_mask is not None else None, output_attentions)
                if conv_layer is not None:
                    output = self._gradient_checkpointing_func(conv_layer, layer_outputs[0])
                    layer_outputs = (output,) + layer_outputs[1:]
            else:
                layer_outputs = encoder_layer(hidden_states, attention_mask, layer_head_mask=head_mask[idx] if head_mask is not None else None, output_attentions=output_attentions)
                if conv_layer is not None:
                    output = conv_layer(layer_outputs[0])
                    layer_outputs = (output,) + layer_outputs[1:]
            hidden_states = layer_outputs[0]
        if output_attentions:
            all_attentions = all_attentions + (layer_outputs[1],)
    if output_hidden_states:
        encoder_states = encoder_states + (hidden_states,)
    if not return_dict:
        return tuple((v for v in [hidden_states, encoder_states, all_attentions] if v is not None))
    return BaseModelOutput(last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions)

Args:
    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)
    head_mask (`torch.Tensor` of shape `(encoder_layers, encoder_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**.

    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.


def build(self, var_list):
    if self.built:
        return
    super().build(var_list)
    self._momentums = self.add_optimizer_variables(var_list, 'momentum')

Initialize optimizer variables.

Lion optimizer has one variable `momentums`.

Args:
    var_list: list of model variables to build Lion variables on.


def resize(self, image: np.ndarray, size: Dict[str, int], resample: PILImageResampling=PILImageResampling.BICUBIC, data_format: Optional[Union[str, ChannelDimension]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None, **kwargs) -> np.ndarray:
    size = get_size_dict(size)
    if 'height' not in size or 'width' not in size:
        raise ValueError(f'The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}')
    output_size = (size['height'], size['width'])
    return resize(image, size=output_size, resample=resample, data_format=data_format, input_data_format=input_data_format, **kwargs)

Resize an image to `(size["height"], size["width"])`.

Args:
    image (`np.ndarray`):
        Image to resize.
    size (`Dict[str, int]`):
        Dictionary in the format `{"height": int, "width": int}` specifying the size of the output image.
    resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BICUBIC`):
        `PILImageResampling` filter to use when resizing the image e.g. `PILImageResampling.BICUBIC`.
    data_format (`ChannelDimension` or `str`, *optional*):
        The channel dimension format for the output image. If unset, the channel dimension format of the input
        image is used. Can be one of:
        - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
        - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
        - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
    input_data_format (`ChannelDimension` or `str`, *optional*):
        The channel dimension format for the input image. If unset, the channel dimension format is inferred
        from the input image. Can be one of:
        - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
        - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
        - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.

Returns:
    `np.ndarray`: The resized image.


def get_rel_pos(q_size, k_size, rel_pos):
    max_rel_dist = int(2 * max(q_size, k_size) - 1)
    if rel_pos.shape[0] != max_rel_dist:
        rel_pos_resized = nn.functional.interpolate(rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1), size=max_rel_dist, mode='linear')
        rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0)
    else:
        rel_pos_resized = rel_pos
    q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0)
    k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0)
    relative_coords = q_coords - k_coords + (k_size - 1) * max(q_size / k_size, 1.0)
    return rel_pos_resized[relative_coords.long()]

Get relative positional embeddings according to the relative positions of query and key sizes.

Args:
    q_size (`int`):
        Size of query q.
    k_size (`int`):
        Size of key k.
    rel_pos (`torch.Tensor`):
        Relative position embeddings (num_embeddings, num_channels).

Returns:
    Extracted positional embeddings according to relative positions.


def _to_boolean(operand: List[WorkSpaceMessage]) -> Optional[bool]:
    if not operand:
        return None
    if len(operand) > 1:
        raise ValueError('Expected a single boolean result but got multiple items.')
    if not fhir_types.is_boolean(operand[0].message):
        raise ValueError('Expected a boolean but got a non-boolean value.')
    return proto_utils.get_value_at_field(operand[0].message, 'value')

Converts an evaluation result to a boolean value or None.

Args:
  operand: an expression operand result to convert to boolean.

Returns:
  the boolean value, or None if the operand was empty.

Raises:
  ValueError if it is not an empty result or a single, boolean value.


def compute_context_repetition_mask(self, input_ids: torch.LongTensor) -> torch.LongTensor:
    self._check_input_ids_shape(input_ids)
    batch_size, _ = input_ids.shape
    state = SynthIDTextWatermarkState(batch_size=batch_size, ngram_len=self.ngram_len, context_history_size=self.context_history_size, device=self.device)
    contexts = input_ids[:, :-1].unfold(dimension=1, size=self.ngram_len - 1, step=1)
    _, num_contexts, _ = contexts.shape
    are_repeated_contexts = []
    for i in range(num_contexts):
        context = contexts[:, i, :]
        hash_result = torch.ones(batch_size, device=self.device, dtype=torch.long)
        context_hash = self.accumulate_hash(hash_result, context)[:, None]
        is_repeated_context = (state.context_history == context_hash).any(dim=1, keepdim=True)
        are_repeated_contexts.append(is_repeated_context)
        state.context_history = torch.concat((context_hash, state.context_history), dim=1)[:, :-1]
    are_repeated_contexts = torch.concat(are_repeated_contexts, dim=1)
    return torch.logical_not(are_repeated_contexts)

Computes repetition mask.

0 and 1 stand for repeated and not repeated context n-1 grams respectively.

Args:
    input_ids (`torch.LongTensor`):
        Input token ids (batch_size, input_len).

Returns:
    Repetitions mask (batch_size, input_len - (ngram_len - 1)).


def assert_scalar(tensor, name=None, message=None):
    with ops.name_scope(name, 'assert_scalar', [tensor]) as name_scope:
        tensor = ops.convert_to_tensor(tensor, name=name_scope)
        shape = tensor.get_shape()
        message = _message_prefix(message)
        if shape.ndims != 0:
            if context.executing_eagerly():
                raise ValueError('%sExpected scalar shape, saw shape: %s.' % (message, shape))
            else:
                raise ValueError('%sExpected scalar shape for %s, saw shape: %s.' % (message, tensor.name, shape))
        return tensor

Asserts that the given `tensor` is a scalar (i.e. zero-dimensional).

This function raises `ValueError` unless it can be certain that the given
`tensor` is a scalar. `ValueError` is also raised if the shape of `tensor` is
unknown.

Args:
  tensor: A `Tensor`.
  name:  A name for this operation. Defaults to "assert_scalar"
  message: A string to prefix to the default message.

Returns:
  The input tensor (potentially converted to a `Tensor`).

Raises:
  ValueError: If the tensor is not scalar (rank 0), or if its shape is
    unknown.


def _get_num_audio_features(self, audio_lengths: Sequence[int]) -> Sequence[int]:
    hop_length = self.melspec_kwargs['hop_length']
    effective_window_size = self.projector_window_size // self.projector_downsample_rate
    projector_lengths = []
    for raw_length in audio_lengths:
        mel_length = raw_length // hop_length + 1
        encoder_length = mel_length // 2
        nblocks = math.ceil(encoder_length / self.projector_window_size)
        projector_length = nblocks * effective_window_size
        projector_lengths.append(projector_length)
    return projector_lengths

Gets the (variable length) number of features (i.e., projector output) for the sequences
being considered.

Args:
    audio_lengths (`Sequence[int]`):
        Sequence of one or more raw audio lengths.


def __init__(self, hooks=None, scaffold=None, master='', config=None, checkpoint_dir=None, stop_grace_period_secs=120, checkpoint_filename_with_path=None):
    session_creator = ChiefSessionCreator(scaffold=scaffold, master=master, config=config, checkpoint_dir=checkpoint_dir, checkpoint_filename_with_path=checkpoint_filename_with_path)
    super(SingularMonitoredSession, self).__init__(session_creator, hooks, should_recover=False, stop_grace_period_secs=stop_grace_period_secs)

Creates a SingularMonitoredSession.

Args:
  hooks: An iterable of `SessionRunHook' objects.
  scaffold: A `Scaffold` used for gathering or building supportive ops. If
    not specified a default one is created. It's used to finalize the graph.
  master: `String` representation of the TensorFlow master to use.
  config: `ConfigProto` proto used to configure the session.
  checkpoint_dir: A string.  Optional path to a directory where to restore
    variables.
  stop_grace_period_secs: Number of seconds given to threads to stop after
    `close()` has been called.
  checkpoint_filename_with_path: A string. Optional path to a checkpoint
    file from which to restore variables.


def get_config(self):
    raise NotImplementedError(f'{self} does not implement get_config()')

Returns the config of the quantizer.

A quantizer config is a Python dictionary (serializable)
containing all configuration parameters of the quantizer.
The same quantizer can be reinstantiated later
(without any saved state) from this configuration.

This method is optional if you are just training and executing models,
exporting to and from SavedModels, or using weight checkpoints.

This method is required for Keras `model_to_estimator`, saving and
loading models to HDF5 formats, Keras model cloning, some visualization
utilities, and exporting models to and from JSON.

Returns:
    Python dictionary.


def __init__(self, session, run_with_hooks_fn):
    self._session = session
    self._run_with_hooks_fn = run_with_hooks_fn

Initializes the `step_context` argument for a `step_fn` invocation.

Args:
  session: An instance of `tf.compat.v1.Session`.
  run_with_hooks_fn: A function for running fetches and hooks.


def _compute_linear_scaling_rope_parameters(config: Optional[PretrainedConfig]=None, device: Optional['torch.device']=None, seq_len: Optional[int]=None, **rope_kwargs) -> tuple['torch.Tensor', float]:
    if config is not None and len(rope_kwargs) > 0:
        raise ValueError(f'Unexpected arguments: `**rope_kwargs` and `config` are mutually exclusive in `_compute_linear_scaling_rope_parameters`, got `rope_kwargs`={rope_kwargs} and `config`={config}')
    if len(rope_kwargs) > 0:
        factor = rope_kwargs['factor']
    elif config is not None:
        factor = config.rope_scaling['factor']
    inv_freq, attention_factor = _compute_default_rope_parameters(config, device, seq_len, **rope_kwargs)
    inv_freq /= factor
    return (inv_freq, attention_factor)

Computes the inverse frequencies with linear scaling. Credits to the Reddit user /u/kaiokendev
Args:
    config ([`~transformers.PretrainedConfig`]):
        The model configuration.
    device (`torch.device`):
        The device to use for initialization of the inverse frequencies.
    seq_len (`int`, *optional*):
        The current sequence length. Unused for this type of RoPE.
    rope_kwargs (`Dict`, *optional*):
        BC compatibility with the previous RoPE class instantiation, will be removed in v4.45.
Returns:
    Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
    post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).


class MixtralConfig(PretrainedConfig):
    model_type = 'mixtral'
    keys_to_ignore_at_inference = ['past_key_values']
    base_model_tp_plan = {'layers.*.self_attn.q_proj': 'colwise', 'layers.*.self_attn.k_proj': 'colwise', 'layers.*.self_attn.v_proj': 'colwise', 'layers.*.self_attn.o_proj': 'rowwise', 'layers.*.block_sparse_moe.gate': 'colwise_rep', 'layers.*.block_sparse_moe.experts.*.w1': 'colwise', 'layers.*.block_sparse_moe.experts.*.w2': 'rowwise', 'layers.*.block_sparse_moe.experts.*.w3': 'colwise'}
    base_model_pp_plan = {'embed_tokens': (['input_ids'], ['inputs_embeds']), 'layers': (['hidden_states', 'attention_mask'], ['hidden_states']), 'norm': (['hidden_states'], ['hidden_states'])}

    def __init__(self, vocab_size=32000, hidden_size=4096, intermediate_size=14336, num_hidden_layers=32, num_attention_heads=32, num_key_value_heads=8, head_dim=None, hidden_act='silu', max_position_embeddings=4096 * 32, initializer_range=0.02, rms_norm_eps=1e-05, use_cache=True, pad_token_id=None, bos_token_id=1, eos_token_id=2, tie_word_embeddings=False, rope_theta=1000000.0, sliding_window=None, attention_dropout=0.0, num_experts_per_tok=2, num_local_experts=8, output_router_logits=False, router_aux_loss_coef=0.001, router_jitter_noise=0.0, **kwargs):
        self.vocab_size = vocab_size
        self.max_position_embeddings = max_position_embeddings
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.sliding_window = sliding_window
        if num_key_value_heads is None:
            num_key_value_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.hidden_act = hidden_act
        self.initializer_range = initializer_range
        self.rms_norm_eps = rms_norm_eps
        self.use_cache = use_cache
        self.rope_theta = rope_theta
        self.attention_dropout = attention_dropout
        self.head_dim = head_dim
        self.num_experts_per_tok = num_experts_per_tok
        self.num_local_experts = num_local_experts
        self.output_router_logits = output_router_logits
        self.router_aux_loss_coef = router_aux_loss_coef
        self.router_jitter_noise = router_jitter_noise
        super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs)

This is the configuration class to store the configuration of a [`MixtralModel`]. It is used to instantiate an
Mixtral model according to the specified arguments, defining the model architecture. Instantiating a configuration
with the defaults will yield a similar configuration to that of the Mixtral-7B-v0.1 or Mixtral-7B-Instruct-v0.1.

[mixtralai/Mixtral-8x7B](https://huggingface.co/mixtralai/Mixtral-8x7B)
[mixtralai/Mixtral-7B-Instruct-v0.1](https://huggingface.co/mixtralai/Mixtral-7B-Instruct-v0.1)

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

Args:
    vocab_size (`int`, *optional*, defaults to 32000):
        Vocabulary size of the Mixtral model. Defines the number of different tokens that can be represented by the
        `inputs_ids` passed when calling [`MixtralModel`]
    hidden_size (`int`, *optional*, defaults to 4096):
        Dimension of the hidden representations.
    intermediate_size (`int`, *optional*, defaults to 14336):
        Dimension of the MLP representations.
    num_hidden_layers (`int`, *optional*, defaults to 32):
        Number of hidden layers in the Transformer encoder.
    num_attention_heads (`int`, *optional*, defaults to 32):
        Number of attention heads for each attention layer in the Transformer encoder.
    num_key_value_heads (`int`, *optional*, defaults to 8):
        This is the number of key_value heads that should be used to implement Grouped Query Attention. If
        `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
        `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
        converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
        by meanpooling all the original heads within that group. For more details, check out [this
        paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to `8`.
    head_dim (`int`, *optional*, defaults to `hidden_size // num_attention_heads`):
        The attention head dimension.
    hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
        The non-linear activation function (function or string) in the decoder.
    max_position_embeddings (`int`, *optional*, defaults to `4096*32`):
        The maximum sequence length that this model might ever be used with. Mixtral's sliding window attention
        allows sequence of up to 4096*32 tokens.
    initializer_range (`float`, *optional*, defaults to 0.02):
        The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
    rms_norm_eps (`float`, *optional*, defaults to 1e-05):
        The epsilon used by the rms normalization layers.
    use_cache (`bool`, *optional*, defaults to `True`):
        Whether or not the model should return the last key/values attentions (not used by all models). Only
        relevant if `config.is_decoder=True`.
    pad_token_id (`int`, *optional*):
        The id of the padding token.
    bos_token_id (`int`, *optional*, defaults to 1):
        The id of the "beginning-of-sequence" token.
    eos_token_id (`int`, *optional*, defaults to 2):
        The id of the "end-of-sequence" token.
    tie_word_embeddings (`bool`, *optional*, defaults to `False`):
        Whether the model's input and output word embeddings should be tied.
    rope_theta (`float`, *optional*, defaults to 1000000.0):
        The base period of the RoPE embeddings.
    sliding_window (`int`, *optional*):
        Sliding window attention window size. If not specified, will default to `4096`.
    attention_dropout (`float`, *optional*, defaults to 0.0):
        The dropout ratio for the attention probabilities.
    num_experts_per_tok (`int`, *optional*, defaults to 2):
        The number of experts to route per-token, can be also interpreted as the `top-k` routing
        parameter
    num_local_experts (`int`, *optional*, defaults to 8):
        Number of experts per Sparse MLP layer.
    output_router_logits (`bool`, *optional*, defaults to `False`):
        Whether or not the router logits should be returned by the model. Enabling this will also
        allow the model to output the auxiliary loss. See [here]() for more details
    router_aux_loss_coef (`float`, *optional*, defaults to 0.001):
        The aux loss factor for the total loss.
    router_jitter_noise (`float`, *optional*, defaults to 0.0):
        Amount of noise to add to the router.

```python
>>> from transformers import MixtralModel, MixtralConfig

>>> # Initializing a Mixtral 7B style configuration
>>> configuration = MixtralConfig()

>>> # Initializing a model from the Mixtral 7B style configuration
>>> model = MixtralModel(configuration)

>>> # Accessing the model configuration
>>> configuration = model.config
```


def call(self, input_ids=None, position_ids=None, token_type_ids=None, inputs_embeds=None, training=False):
    assert not (input_ids is None and inputs_embeds is None)
    if input_ids is not None:
        check_embeddings_within_bounds(input_ids, self.config.vocab_size)
        inputs_embeds = tf.gather(params=self.weight, indices=input_ids)
    input_shape = shape_list(inputs_embeds)[:-1]
    if token_type_ids is None:
        token_type_ids = tf.fill(dims=input_shape, value=0)
    if self.trigram_input:
        inputs_embeds = tf.concat([tf.pad(inputs_embeds[:, 1:], ((0, 0), (0, 1), (0, 0))), inputs_embeds, tf.pad(inputs_embeds[:, :-1], ((0, 0), (1, 0), (0, 0)))], axis=2)
    if self.trigram_input or self.embedding_size != self.hidden_size:
        inputs_embeds = self.embedding_transformation(inputs_embeds)
    if position_ids is None:
        position_ids = tf.expand_dims(tf.range(start=0, limit=input_shape[-1]), axis=0)
    position_embeds = tf.gather(params=self.position_embeddings, indices=position_ids)
    token_type_embeds = tf.gather(params=self.token_type_embeddings, indices=token_type_ids)
    final_embeddings = inputs_embeds + position_embeds + token_type_embeds
    final_embeddings = self.LayerNorm(inputs=final_embeddings)
    final_embeddings = self.dropout(inputs=final_embeddings, training=training)
    return final_embeddings

Applies embedding based on inputs tensor.

Returns:
    final_embeddings (`tf.Tensor`): output embedding tensor.


class CodeGenerator(NodeStateTracker, gast.NodeVisitor):

    def __init__(self, ctx):
        super(CodeGenerator, self).__init__(ctx)
        self._output_code = ''
        self.source_map = {}

    def emit(self, code):
        self._output_code += code

    @property
    def code_buffer(self):
        return self._output_code

    def visit(self, node):
        if anno.hasanno(node, anno.Basic.SKIP_PROCESSING):
            return
        parent_origin = self.ctx.current_origin
        eof_before = len(self._output_code)
        if anno.hasanno(node, anno.Basic.ORIGIN):
            self.ctx.current_origin = anno.getanno(node, anno.Basic.ORIGIN)
        try:
            ret = super(CodeGenerator, self).visit(node)
            eof_after = len(self._output_code)
            if eof_before - eof_after:
                inherited_origin = anno.getanno(node, anno.Basic.ORIGIN, default=parent_origin)
                if inherited_origin is not None:
                    self.source_map[eof_before, eof_after] = inherited_origin
            return ret
        finally:
            self.ctx.current_origin = parent_origin

Base class for general-purpose Python-to-string code transformation.

Similar to Base, but outputs arbitrary strings instead of a Python AST.

This uses the same visitor mechanism that the standard NodeVisitor uses,
meaning that subclasses write handlers for the different kinds of nodes.
New code is generated using the emit method, which appends to a code buffer
that can be afterwards obtained from code_buffer.

Example:

  class SimpleCodeGen(CodeGenerator):

    def visitIf(self, node):
      self.emit('if ')
      self.visit(node.test)
      self.emit(' { ')
      self.visit(node.body)
      self.emit(' } else { ')
      self.visit(node.orelse)
      self.emit(' } ')

  node = ast.parse(...)
  gen = SimpleCodeGen()
  gen.visit(node)
  # gen.code_buffer contains the resulting code


class AriaGroupedExpertsGemm(nn.Module):

    def __init__(self, in_features, out_features, groups):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.groups = groups
        self.weight = nn.Parameter(torch.empty(groups, in_features, out_features))

    def forward(self, input, tokens_per_expert):
        """
        Perform grouped matrix multiplication.

        Args:
            input (`torch.Tensor`):
                Input tensor of shape (num_tokens, in_features).
            tokens_per_expert (`torch.Tensor`):
                Number of tokens assigned to each expert.

        Returns:
            torch.Tensor: Output tensor of shape (num_tokens, out_features).
        """
        return sequential_experts_gemm(input, self.weight, tokens_per_expert.cpu())

Grouped GEMM (General Matrix Multiplication) module for efficient expert computation.
This module utilizes the grouped_gemm library (https://github.com/fanshiqing/grouped_gemm)
for optimized performance. If the grouped_gemm library is not installed, it gracefully
falls back to a sequential GEMM implementation, which may be slower but ensures
functionality.

Args:
    in_features (`int`):
        Number of input features.
    out_features (`int`):
        Number of output features.
    groups (`int`):
        Number of expert groups.