def einsum(subscripts, *operands, **kwargs):
    if any_symbolic_tensors(operands):
        return Einsum(subscripts).symbolic_call(*operands, **kwargs)
    return backend.numpy.einsum(subscripts, *operands, **kwargs)

Evaluates the Einstein summation convention on the operands.

Args:
    subscripts: Specifies the subscripts for summation as comma separated
        list of subscript labels. An implicit (classical Einstein
        summation) calculation is performed unless the explicit indicator
        `->` is included as well as subscript labels of the precise
        output form.
    operands: The operands to compute the Einstein sum of.

Returns:
    The calculation based on the Einstein summation convention.

Example:
>>> from keras.src import ops
>>> a = ops.arange(25).reshape(5, 5)
>>> b = ops.arange(5)
>>> c = ops.arange(6).reshape(2, 3)

Trace of a matrix:

>>> ops.einsum("ii", a)
60
>>> ops.einsum(a, [0, 0])
60
>>> ops.trace(a)
60

Extract the diagonal:

>>> ops.einsum("ii -> i", a)
array([ 0,  6, 12, 18, 24])
>>> ops.einsum(a, [0, 0], [0])
array([ 0,  6, 12, 18, 24])
>>> ops.diag(a)
array([ 0,  6, 12, 18, 24])

Sum over an axis:

>>> ops.einsum("ij -> i", a)
array([ 10,  35,  60,  85, 110])
>>> ops.einsum(a, [0, 1], [0])
array([ 10,  35,  60,  85, 110])
>>> ops.sum(a, axis=1)
array([ 10,  35,  60,  85, 110])

For higher dimensional tensors summing a single axis can be done
with ellipsis:

>>> ops.einsum("...j -> ...", a)
array([ 10,  35,  60,  85, 110])
>>> np.einsum(a, [..., 1], [...])
array([ 10,  35,  60,  85, 110])

Compute a matrix transpose or reorder any number of axes:

>>> ops.einsum("ji", c)
array([[0, 3],
       [1, 4],
       [2, 5]])
>>> ops.einsum("ij -> ji", c)
array([[0, 3],
       [1, 4],
       [2, 5]])
>>> ops.einsum(c, [1, 0])
array([[0, 3],
       [1, 4],
       [2, 5]])
>>> ops.transpose(c)
array([[0, 3],
       [1, 4],
       [2, 5]])

Matrix vector multiplication:

>>> ops.einsum("ij, j", a, b)
array([ 30,  80, 130, 180, 230])
>>> ops.einsum(a, [0, 1], b, [1])
array([ 30,  80, 130, 180, 230])
>>> ops.einsum("...j, j", a, b)
array([ 30,  80, 130, 180, 230])


def assertAllDifferent(self, tensors):
    values = [array_ops.reshape(t, shape=[-1]) for t in tensors]
    values = array_ops.concat(values, axis=0)
    values = self.evaluate(values)
    values = values.tolist()
    self.assertAllEqual(len(values), len(set(values)))

Checks that there are no duplicate elements anywhere among the tensors.

Args:
  tensors: a list of tensors. They can have different shapes.


def export(self, name=None):
    with ops.name_scope(name, '%s_lookup_table_export_values' % self.name, [self.resource_handle]):
        with ops.colocate_with(self.resource_handle):
            exported_keys, exported_values = gen_lookup_ops.lookup_table_export_v2(self.resource_handle, self._key_dtype, self._value_dtype)
    return (exported_keys, exported_values)

Returns tensors of all keys and values in the table.

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

Returns:
  A pair of tensors with the first tensor containing all keys and the
    second tensors containing all values in the table.


def __init__(self, step_name, transform_id=None):
    self.step_name = step_name
    self.transform_id = transform_id

Creates a new step NameContext.

Args:
  step_name: The name of the step.


class MaxSizePartitioner(Partitioner):

    def __init__(self, max_shard_bytes, max_shards=None, bytes_per_string=16):
        """Creates a new `MaxSizePartitioner`.

    Args:
      max_shard_bytes: The maximum size any given shard is allowed to be.
      max_shards: The maximum number of shards in `int` created taking
        precedence over `max_shard_bytes`.
      bytes_per_string: If the partition value is of type string, this provides
        an estimate of how large each string is.
    """
        if max_shard_bytes < 1:
            raise ValueError(f'Argument `max_shard_bytes` must be positive. Received {max_shard_bytes}')
        if max_shards and max_shards < 1:
            raise ValueError(f'Argument `max_shards` must be positive. Received {max_shards}')
        if bytes_per_string < 1:
            raise ValueError(f'Argument `bytes_per_string` must be positive. Received: {bytes_per_string}')
        self._max_shard_bytes = max_shard_bytes
        self._max_shards = max_shards
        self._bytes_per_string = bytes_per_string

    def __call__(self, shape, dtype, axis=0):
        return partitioned_variables.variable_axis_size_partitioner(max_shard_bytes=self._max_shard_bytes, max_shards=self._max_shards, bytes_per_string_element=self._bytes_per_string, axis=axis)(shape, dtype)

Partitioner that keeps shards below `max_shard_bytes`.

This partitioner ensures each shard has at most `max_shard_bytes`, and tries
to allocate as few shards as possible, i.e., keeping shard size as large
as possible.

If the partitioner hits the `max_shards` limit, then each shard may end up
larger than `max_shard_bytes`. By default `max_shards` equals `None` and no
limit on the number of shards is enforced.

Examples:

>>> partitioner = MaxSizePartitioner(max_shard_bytes=4)
>>> partitions = partitioner(tf.TensorShape([6, 1]), tf.float32)
>>> [6, 1]
>>> partitioner = MaxSizePartitioner(max_shard_bytes=4, max_shards=2)
>>> partitions = partitioner(tf.TensorShape([6, 1]), tf.float32)
>>> [2, 1]
>>> partitioner = MaxSizePartitioner(max_shard_bytes=1024)
>>> partitions = partitioner(tf.TensorShape([6, 1]), tf.float32)
>>> [1, 1]
>>>
>>> # use in ParameterServerStrategy
>>> # strategy = tf.distribute.experimental.ParameterServerStrategy(
>>> #   cluster_resolver=cluster_resolver, variable_partitioner=partitioner)


def forward(self, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.FloatTensor]=None, token_type_ids: Optional[torch.LongTensor]=None, position_ids: Optional[torch.LongTensor]=None, head_mask: Optional[torch.FloatTensor]=None, inputs_embeds: Optional[torch.FloatTensor]=None, labels: Optional[torch.FloatTensor]=None, next_sentence_label: Optional[torch.LongTensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[BigBirdForPreTrainingOutput, Tuple[torch.FloatTensor]]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    outputs = self.bert(input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    sequence_output, pooled_output = outputs[:2]
    prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)
    total_loss = None
    if labels is not None:
        loss_fct = CrossEntropyLoss()
        total_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
    if next_sentence_label is not None and total_loss is not None:
        next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1))
        total_loss = total_loss + next_sentence_loss
    if not return_dict:
        output = (prediction_scores, seq_relationship_score) + outputs[2:]
        return (total_loss,) + output if total_loss is not None else output
    return BigBirdForPreTrainingOutput(loss=total_loss, prediction_logits=prediction_scores, seq_relationship_logits=seq_relationship_score, hidden_states=outputs.hidden_states, attentions=outputs.attentions)

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

    - 0 indicates sequence B is a continuation of sequence A,
    - 1 indicates sequence B is a random sequence.

Example:

```python
>>> from transformers import AutoTokenizer, BigBirdForPreTraining
>>> import torch

>>> tokenizer = AutoTokenizer.from_pretrained("google/bigbird-roberta-base")
>>> model = BigBirdForPreTraining.from_pretrained("google/bigbird-roberta-base")

>>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
>>> outputs = model(**inputs)

>>> prediction_logits = outputs.prediction_logits
>>> seq_relationship_logits = outputs.seq_relationship_logits
```


def create_segment_token_type_ids_from_sequences(self, query_ids: List[int], table_values: List[TableValue]) -> List[int]:
    table_ids = list(zip(*table_values))[0] if table_values else []
    return [0] * (1 + len(query_ids) + 1) + [1] * len(table_ids)

Creates the segment token type IDs according to the query token IDs and a list of table values.

Args:
    query_ids (`List[int]`): list of token IDs corresponding to the ID.
    table_values (`List[TableValue]`): lift of table values, which are named tuples containing the
        token value, the column ID and the row ID of said token.

Returns:
    `List[int]`: List of ints containing the segment token type IDs values.


def diff(self) -> List[str]:
    return set(self.to_track.keys()) - self._seen

This method returns a set difference between the keys in the tracked state dict and the one we have access so far.
This is an effective method to check if we have update all the keys

Returns:
    List[str]: List of keys not yet updated


def _BatchNormWithGlobalNormalizationGrad(op: ops.Operation, grad):
    dx, dm, dv, db, dg = gen_nn_ops.batch_norm_with_global_normalization_grad(op.inputs[0], op.inputs[1], op.inputs[2], op.inputs[4], grad, op.get_attr('variance_epsilon'), op.get_attr('scale_after_normalization'))
    return (dx, dm, dv, db, dg)

Return the gradients for the 5 inputs of BatchNormWithGlobalNormalization.

We do not backprop anything for the mean and var intentionally as they are
not being trained with backprop in the operation.

Args:
  op: The BatchNormOp for which we need to generate gradients.
  grad: Tensor.  The gradients passed to the BatchNormOp.

Returns:
  dx: Backprop for input, which is (grad * (g * rsqrt(v + epsilon)))
  dm: Backprop for mean, which is
      sum_over_rest(grad * g) * (-1 / rsqrt(v + epsilon))
  dv: Backprop for variance, which is
      sum_over_rest(grad * g * (x - m)) * (-1/2) * (v + epsilon) ^ (-3/2)
  db: Backprop for beta, which is grad reduced in all except the
      last dimension.
  dg: Backprop for gamma, which is (grad * ((x - m) * rsqrt(v + epsilon)))


def remove_checkpoint(checkpoint_prefix, checkpoint_format_version=saver_pb2.SaverDef.V2, meta_graph_suffix='meta'):
    _delete_file_if_exists(meta_graph_filename(checkpoint_prefix, meta_graph_suffix))
    if checkpoint_format_version == saver_pb2.SaverDef.V2:
        _delete_file_if_exists(checkpoint_prefix + '.index')
        _delete_file_if_exists(checkpoint_prefix + '.data-?????-of-?????')
    else:
        _delete_file_if_exists(checkpoint_prefix)

Removes a checkpoint given by `checkpoint_prefix`.

Args:
  checkpoint_prefix: The prefix of a V1 or V2 checkpoint. Typically the result
    of `Saver.save()` or that of `tf.train.latest_checkpoint()`, regardless of
    sharded/non-sharded or V1/V2.
  checkpoint_format_version: `SaverDef.CheckpointFormatVersion`, defaults to
    `SaverDef.V2`.
  meta_graph_suffix: Suffix for `MetaGraphDef` file. Defaults to 'meta'.


def GetGroupMap(self, since=None):
    return GroupUpdateGetter().GetUpdates(self, self.conf['group_url'], since)

Return the group map from this source.

Args:
  since: Get data only changed since this timestamp (inclusive) or None
  for all data.

Returns:
  instance of group.GroupMap


def values(self):
    return self._values

The concatenated rows for this ragged tensor.

`rt.values` is a potentially ragged tensor formed by flattening the two
outermost dimensions of `rt` into a single dimension.

`rt.values.shape = [nvals] + rt.shape[2:]` (where `nvals` is the
number of items in the outer two dimensions of `rt`).

`rt.ragged_rank = self.ragged_rank - 1`

Returns:
  A potentially ragged tensor.

#### Example:

>>> rt = tf.ragged.constant([[3, 1, 4, 1], [], [5, 9, 2], [6], []])
>>> print(rt.values)
tf.Tensor([3 1 4 1 5 9 2 6], shape=(8,), dtype=int32)


def add_variable(self, feature_column, var):
    del feature_column, var
    raise NotImplementedError('StateManager.add_variable')

Adds an existing variable to the state.

Args:
  feature_column: A `FeatureColumn` object to associate this variable with.
  var: The variable.


class UdopConfig(PretrainedConfig):
    model_type = 'udop'
    keys_to_ignore_at_inference = ['past_key_values']
    attribute_map = {'hidden_size': 'd_model', 'num_attention_heads': 'num_heads', 'num_hidden_layers': 'num_layers'}

    def __init__(self, vocab_size=33201, d_model=1024, d_kv=64, d_ff=4096, num_layers=24, num_decoder_layers=None, num_heads=16, relative_attention_num_buckets=32, relative_attention_max_distance=128, relative_bias_args=[{'type': '1d'}, {'type': 'horizontal'}, {'type': 'vertical'}], dropout_rate=0.1, layer_norm_epsilon=1e-06, initializer_factor=1.0, feed_forward_proj='relu', is_encoder_decoder=True, use_cache=True, pad_token_id=0, eos_token_id=1, max_2d_position_embeddings=1024, image_size=224, patch_size=16, num_channels=3, **kwargs):
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.d_kv = d_kv
        self.d_ff = d_ff
        self.num_layers = num_layers
        self.num_decoder_layers = num_decoder_layers if num_decoder_layers is not None else self.num_layers
        self.num_heads = num_heads
        self.relative_attention_num_buckets = relative_attention_num_buckets
        self.relative_attention_max_distance = relative_attention_max_distance
        self.dropout_rate = dropout_rate
        self.layer_norm_epsilon = layer_norm_epsilon
        self.initializer_factor = initializer_factor
        self.feed_forward_proj = feed_forward_proj
        self.use_cache = use_cache
        self.max_2d_position_embeddings = max_2d_position_embeddings
        self.image_size = image_size
        self.patch_size = patch_size
        self.num_channels = num_channels
        if not isinstance(relative_bias_args, list):
            raise TypeError('`relative_bias_args` should be a list of dictionaries.')
        self.relative_bias_args = relative_bias_args
        act_info = self.feed_forward_proj.split('-')
        self.dense_act_fn = act_info[-1]
        self.is_gated_act = act_info[0] == 'gated'
        if len(act_info) > 1 and act_info[0] != 'gated' or len(act_info) > 2:
            raise ValueError(f"`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer.Please make sure `feed_forward_proj` is of the format `gated-{{ACT_FN}}` or `{{ACT_FN}}`, e.g. 'gated-gelu' or 'relu'")
        super().__init__(pad_token_id=pad_token_id, eos_token_id=eos_token_id, is_encoder_decoder=is_encoder_decoder, **kwargs)

This is the configuration class to store the configuration of a [`UdopForConditionalGeneration`]. It is used to
instantiate a UDOP 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 UDOP
[microsoft/udop-large](https://huggingface.co/microsoft/udop-large) architecture.

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

Arguments:
    vocab_size (`int`, *optional*, defaults to 33201):
        Vocabulary size of the UDOP model. Defines the number of different tokens that can be represented by the
        `inputs_ids` passed when calling [`UdopForConditionalGeneration`].
    d_model (`int`, *optional*, defaults to 1024):
        Size of the encoder layers and the pooler layer.
    d_kv (`int`, *optional*, defaults to 64):
        Size of the key, query, value projections per attention head. The `inner_dim` of the projection layer will
        be defined as `num_heads * d_kv`.
    d_ff (`int`, *optional*, defaults to 4096):
        Size of the intermediate feed forward layer in each `UdopBlock`.
    num_layers (`int`, *optional*, defaults to 24):
        Number of hidden layers in the Transformer encoder and decoder.
    num_decoder_layers (`int`, *optional*):
        Number of hidden layers in the Transformer decoder. Will use the same value as `num_layers` if not set.
    num_heads (`int`, *optional*, defaults to 16):
        Number of attention heads for each attention layer in the Transformer encoder and decoder.
    relative_attention_num_buckets (`int`, *optional*, defaults to 32):
        The number of buckets to use for each attention layer.
    relative_attention_max_distance (`int`, *optional*, defaults to 128):
        The maximum distance of the longer sequences for the bucket separation.
    relative_bias_args (`List[dict]`, *optional*, defaults to `[{'type': '1d'}, {'type': 'horizontal'}, {'type': 'vertical'}]`):
        A list of dictionaries containing the arguments for the relative bias layers.
    dropout_rate (`float`, *optional*, defaults to 0.1):
        The ratio for all dropout layers.
    layer_norm_epsilon (`float`, *optional*, defaults to 1e-06):
        The epsilon used by the layer normalization layers.
    initializer_factor (`float`, *optional*, defaults to 1.0):
        A factor for initializing all weight matrices (should be kept to 1, used internally for initialization
        testing).
    feed_forward_proj (`string`, *optional*, defaults to `"relu"`):
        Type of feed forward layer to be used. Should be one of `"relu"` or `"gated-gelu"`. Udopv1.1 uses the
        `"gated-gelu"` feed forward projection. Original Udop uses `"relu"`.
    is_encoder_decoder (`bool`, *optional*, defaults to `True`):
        Whether the model should behave as an encoder/decoder or not.
    use_cache (`bool`, *optional*, defaults to `True`):
        Whether or not the model should return the last key/values attentions (not used by all models).
    pad_token_id (`int`, *optional*, defaults to 0):
        The id of the padding token in the vocabulary.
    eos_token_id (`int`, *optional*, defaults to 1):
        The id of the end-of-sequence token in the vocabulary.
    max_2d_position_embeddings (`int`, *optional*, defaults to 1024):
        The maximum absolute position embeddings for relative position encoding.
    image_size (`int`, *optional*, defaults to 224):
        The size of the input images.
    patch_size (`int`, *optional*, defaults to 16):
        The patch size used by the vision encoder.
    num_channels (`int`, *optional*, defaults to 3):
        The number of channels in the input images.


class StatelessScope:

    def __init__(self, state_mapping=None, collect_losses=False, initialize_variables=True):
        from keras.src import backend
        from keras.src.backend.common.variables import Variable
        self.collect_losses = collect_losses
        self.initialize_variables = initialize_variables
        self.losses = []
        self.state_mapping = {}
        state_mapping = state_mapping or {}
        for k, v in state_mapping:
            if not isinstance(k, Variable):
                raise ValueError(f'Invalid reference variable in StatelessScope: all keys in argument `mapping` must be Variable instances. Received instead: {k}')
            if isinstance(v, Variable):
                v = backend.cast(v.value, dtype=k.dtype)
            else:
                v = backend.convert_to_tensor(v, dtype=k.dtype)
            if k.shape != v.shape:
                raise ValueError(f'Invalid variable value in StatelessScope: all values in argument `mapping` must be tensors with a shape that matches the corresponding variable shape. For variable {k}, received invalid value {v} with shape {v.shape}.')
            self.state_mapping[id(k)] = v

    def __enter__(self):
        self.original_scope = get_stateless_scope()
        global_state.set_global_attribute('stateless_scope', self)
        return self

    def add_loss(self, loss):
        self.losses.append(loss)

    def add_update(self, update):
        variable, value = update
        self.state_mapping[id(variable)] = value

    def get_current_value(self, variable):
        return self.state_mapping.get(id(variable), None)

    def __exit__(self, *args, **kwargs):
        global_state.set_global_attribute('stateless_scope', self.original_scope)
        if self.original_scope is None and self.initialize_variables:
            from keras.src.backend.common.variables import initialize_all_variables
            initialize_all_variables()

Scope to prevent any update to Keras Variables.

The values of variables to be used inside the scope
should be passed via the `state_mapping` argument, a
list of tuples `(k, v)` where `k` is a `Variable`
and `v` is the intended value for this variable
(a backend tensor).

Updated values can be collected on scope exit via
`value = scope.get_current_value(variable)`. No updates
will be applied in-place to any variables for the duration
of the scope.

Example:

```python
state_mapping = [(k, ops.ones(k.shape, k.dtype)) for k in model.weights]
with keras.StatelessScope(state_mapping) as scope:
    outputs = model.some_function(inputs)

# All model variables remain unchanged. Their new values can be
# collected via:
for k in model.weights:
    new_value = scope.get_current_value(k)
    print(f"New value for {k}: {new_value})
```


def flip_back(output_flipped, flip_pairs, target_type='gaussian-heatmap'):
    if target_type not in ['gaussian-heatmap', 'combined-target']:
        raise ValueError('target_type should be gaussian-heatmap or combined-target')
    if output_flipped.ndim != 4:
        raise ValueError('output_flipped should be [batch_size, num_keypoints, height, width]')
    batch_size, num_keypoints, height, width = output_flipped.shape
    channels = 1
    if target_type == 'combined-target':
        channels = 3
        output_flipped[:, 1::3, ...] = -output_flipped[:, 1::3, ...]
    output_flipped = output_flipped.reshape(batch_size, -1, channels, height, width)
    output_flipped_back = output_flipped.clone()
    for left, right in flip_pairs.tolist():
        output_flipped_back[:, left, ...] = output_flipped[:, right, ...]
        output_flipped_back[:, right, ...] = output_flipped[:, left, ...]
    output_flipped_back = output_flipped_back.reshape((batch_size, num_keypoints, height, width))
    output_flipped_back = output_flipped_back.flip(-1)
    return output_flipped_back

Flip the flipped heatmaps back to the original form.

Args:
    output_flipped (`torch.tensor` of shape `(batch_size, num_keypoints, height, width)`):
        The output heatmaps obtained from the flipped images.
    flip_pairs (`torch.Tensor` of shape `(num_keypoints, 2)`):
        Pairs of keypoints which are mirrored (for example, left ear -- right ear).
    target_type (`str`, *optional*, defaults to `"gaussian-heatmap"`):
        Target type to use. Can be gaussian-heatmap or combined-target.
        gaussian-heatmap: Classification target with gaussian distribution.
        combined-target: The combination of classification target (response map) and regression target (offset map).
        Paper ref: Huang et al. The Devil is in the Details: Delving into Unbiased Data Processing for Human Pose Estimation (CVPR 2020).

Returns:
    torch.Tensor: heatmaps that flipped back to the original image


def canonicalize_without_job_and_task(d):
    canonicalized_device = canonicalize(d)
    spec = tf_device.DeviceSpec.from_string(canonicalized_device)
    spec = spec.replace(job=None, task=None, replica=0)
    return spec.to_string()

Partially canonicalize device string.

This returns device string from `d` without including job and task.
This is most useful for parameter server strategy where the device strings are
generated on the chief, but executed on workers.

 For example:
  If d = '/cpu:0', default='/job:worker/task:1', it returns
    '/replica:0/device:CPU:0'.
  If d = '/cpu:0', default='/job:worker', it returns
    '/replica:0/device:CPU:0'.
  If d = '/gpu:0', default=None, it returns
    '/replica:0/device:GPU:0'.

Note: This uses "job:localhost" as the default if executing eagerly.

Args:
  d: a device string or tf.config.LogicalDevice

Returns:
  a partially canonicalized device string.


def get_git_version(git_base_path, git_tag_override):
    unknown_label = b'unknown'
    try:
        val = bytes(subprocess.check_output(['git', str('--git-dir=%s/.git' % git_base_path), str('--work-tree=%s' % git_base_path), 'describe', '--long', '--tags']).strip())
        version_separator = b'-'
        if git_tag_override and val:
            split_val = val.split(version_separator)
            if len(split_val) < 3:
                raise Exception("Expected git version in format 'TAG-COMMITS AFTER TAG-HASH' but got '%s'" % val)
            abbrev_commit = split_val[-1]
            val = version_separator.join([bytes(git_tag_override, 'utf-8'), b'0', abbrev_commit])
        return val if val else unknown_label
    except (subprocess.CalledProcessError, OSError):
        return unknown_label

Get the git version from the repository.

This function runs `git describe ...` in the path given as `git_base_path`.
This will return a string of the form:
<base-tag>-<number of commits since tag>-<shortened sha hash>

For example, 'v0.10.0-1585-gbb717a6' means v0.10.0 was the last tag when
compiled. 1585 commits are after that commit tag, and we can get back to this
version by running `git checkout gbb717a6`.

Args:
  git_base_path: where the .git directory is located
  git_tag_override: Override the value for the git tag. This is useful for
    releases where we want to build the release before the git tag is
    created.
Returns:
  A bytestring representing the git version


def make_connection_with_forwarded_port(self, host_port, device_port, uid=UNKNOWN_UID, cmd=ConnectionHandshakeCommand.INIT):
    self.host_port = host_port
    self.device_port = device_port
    self._counter = self._id_counter()
    self.create_socket_connection()
    self.send_handshake_request(uid, cmd)

Makes a connection to the server with the given forwarded port.

This process assumes that a device port has already been forwarded to a
host port, and it only makes a connection to the snippet server based on
the forwarded port. This is typically used by clients that share the same
snippet server, e.g. the snippet client and its event client.

Args:
  host_port: int, the host port which has already been forwarded.
  device_port: int, the device port listened by the snippet server.
  uid: int, the uid of the server session to continue. It will be ignored
    if the `cmd` requires the server to create a new session.
  cmd: ConnectionHandshakeCommand, the handshake command Enum for the
    server, which requires the server to create a new session or use the
    current session.


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

Tokenize and prepare for the model a list of sequences or a list of pairs of sequences.

Args:
    batch_text_or_text_pairs (`List[str]`, `List[Tuple[str, str]]`, `List[List[str]]`, `List[Tuple[List[str], List[str]]]`, and for not-fast tokenizers, also `List[List[int]]`, `List[Tuple[List[int], List[int]]]`):
        Batch of sequences or pair of sequences to be encoded. This can be a list of
        string/string-sequences/int-sequences or a list of pair of string/string-sequences/int-sequence (see
        details in `encode_plus`).


def _weighted_flat_map(input_datasets: Sequence[dataset_ops.DatasetV2], weights: Optional[Sequence[Union[float, tensor.Tensor]]]=None, name: Optional[str]=None) -> dataset_ops.DatasetV2:
    return _WeightedFlatMap(input_datasets, weights, name=name)

A `Dataset` that fetches elements from `input_datasets` and flattens them.

This operation combines elements from multiple datasets into a flattened
dataset. Elements are read in proportion to the `weights` assigned to each
input dataset. All requested elements from a dataset are read before reading
the elements from the next dataset.

For example, suppose we have 2 datasets:

# TODO(wilsin): Make the following code testable after the API is released.
dataset1 = tf.data.Dataset.range(0, 10)
dataset2 = tf.data.Dataset.range(10, 20),

Suppose that we call `weighted_flat_map` from these 2 datasets with the
following weights:

dataset = tf.data.Dataset.weighted_flat_map([dataset1, dataset2], [0.5, 1.0])

Then, the outcome of the elements is:
# [0, 1, 2, 3, 4, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19]

Args:
  input_datasets: A non-empty list of `tf.data.Dataset` objects with
    compatible structure.
  weights: (Optional.) A list or Tensor of `len(datasets)` non-zero
    floating-point values where `weights[i]` represents the probability to
    sample from `datasets[i]`, or a `tf.data.Dataset` object where each
    element is such a list. Defaults to a uniform distribution across
    `datasets`.
  name: (Optional.) A name for the tf.data operation.

Returns:
  A dataset that reads elements from all its inputs, reading the requested
  elements from an input according to the weight before proceeding to the next
  input. The number of elements read from an input is in proportion to its
  weight given in `weights`.

Raises:
  TypeError: if the `datasets` or `weights` arguments have the wrong type.
  ValueError:
    - if `input_datasets` has less than 2 datasets.
    - if `weights` is specified and does not match the length of
      `input_datasets`.
  InvalidArgumentError:
    - if any of the `input_datasets` has an unknown or infinite cardinality.
    - if any of the `weights` has a value that is less than or equal to 0.0


def squeeze_batch_dims(inp, op, inner_rank, name=None):
    with ops.name_scope(name, 'squeeze_batch_dims', [inp]):
        inp = ops.convert_to_tensor(inp, name='input')
        shape = inp.shape
        inner_shape = shape[-inner_rank:]
        if not inner_shape.is_fully_defined():
            inner_shape = array_ops.shape(inp)[-inner_rank:]
        batch_shape = shape[:-inner_rank]
        if not batch_shape.is_fully_defined():
            batch_shape = array_ops.shape(inp)[:-inner_rank]
        if isinstance(inner_shape, tensor_shape.TensorShape):
            inp_reshaped = array_ops.reshape(inp, [-1] + inner_shape.as_list())
        else:
            inp_reshaped = array_ops.reshape(inp, array_ops.concat(([-1], inner_shape), axis=-1))
        out_reshaped = op(inp_reshaped)
        out_inner_shape = out_reshaped.shape[-inner_rank:]
        if not out_inner_shape.is_fully_defined():
            out_inner_shape = array_ops.shape(out_reshaped)[-inner_rank:]
        out = array_ops.reshape(out_reshaped, array_ops.concat((batch_shape, out_inner_shape), axis=-1))
        out.set_shape(inp.shape[:-inner_rank] + out.shape[-inner_rank:])
        return out

Returns `unsqueeze_batch(op(squeeze_batch(inp)))`.

Where `squeeze_batch` reshapes `inp` to shape
`[prod(inp.shape[:-inner_rank])] + inp.shape[-inner_rank:]`
and `unsqueeze_batch` does the reverse reshape but on the output.

Args:
  inp: A tensor with dims `batch_shape + inner_shape` where `inner_shape`
    is length `inner_rank`.
  op: A callable that takes a single input tensor and returns a single.
    output tensor.
  inner_rank: A python integer.
  name: A string.

Returns:
  `unsqueeze_batch_op(squeeze_batch(inp))`.


def preprocess_input(x, data_format=None):
    return x

A placeholder method for backward compatibility.

The preprocessing logic has been included in the efficientnet model
implementation. Users are no longer required to call this method to
normalize the input data. This method does nothing and only kept as a
placeholder to align the API surface between old and new version of model.

Args:
    x: A floating point `numpy.array` or a tensor.
    data_format: Optional data format of the image tensor/array. `None`
        means the global setting `keras.backend.image_data_format()`
        is used (unless you changed it, it uses `"channels_last"`).
        Defaults to `None`.

Returns:
    Unchanged `numpy.array` or tensor.


def wav_to_features(sample_rate, clip_duration_ms, window_size_ms, window_stride_ms, feature_bin_count, quantize, preprocess, input_wav, output_c_file):
    sess = tf.compat.v1.InteractiveSession()
    model_settings = models.prepare_model_settings(0, sample_rate, clip_duration_ms, window_size_ms, window_stride_ms, feature_bin_count, preprocess)
    audio_processor = input_data.AudioProcessor(None, None, 0, 0, '', 0, 0, model_settings, None)
    results = audio_processor.get_features_for_wav(input_wav, model_settings, sess)
    features = results[0]
    variable_base = os.path.splitext(os.path.basename(input_wav).lower())[0]
    with gfile.GFile(output_c_file, 'w') as f:
        f.write('/* File automatically created by\n')
        f.write(' * tensorflow/examples/speech_commands/wav_to_features.py \\\n')
        f.write(' * --sample_rate=%d \\\n' % sample_rate)
        f.write(' * --clip_duration_ms=%d \\\n' % clip_duration_ms)
        f.write(' * --window_size_ms=%d \\\n' % window_size_ms)
        f.write(' * --window_stride_ms=%d \\\n' % window_stride_ms)
        f.write(' * --feature_bin_count=%d \\\n' % feature_bin_count)
        if quantize:
            f.write(' * --quantize=1 \\\n')
        f.write(' * --preprocess="%s" \\\n' % preprocess)
        f.write(' * --input_wav="%s" \\\n' % input_wav)
        f.write(' * --output_c_file="%s" \\\n' % output_c_file)
        f.write(' */\n\n')
        f.write('const int g_%s_width = %d;\n' % (variable_base, model_settings['fingerprint_width']))
        f.write('const int g_%s_height = %d;\n' % (variable_base, model_settings['spectrogram_length']))
        if quantize:
            features_min, features_max = input_data.get_features_range(model_settings)
            f.write('const unsigned char g_%s_data[] = {' % variable_base)
            i = 0
            for value in features.flatten():
                quantized_value = int(round(255 * (value - features_min) / (features_max - features_min)))
                if quantized_value < 0:
                    quantized_value = 0
                if quantized_value > 255:
                    quantized_value = 255
                if i == 0:
                    f.write('\n  ')
                f.write('%d, ' % quantized_value)
                i = (i + 1) % 10
        else:
            f.write('const float g_%s_data[] = {\n' % variable_base)
            i = 0
            for value in features.flatten():
                if i == 0:
                    f.write('\n  ')
                f.write('%f, ' % value)
                i = (i + 1) % 10
        f.write('\n};\n')

Converts an audio file into its corresponding feature map.

Args:
  sample_rate: Expected sample rate of the wavs.
  clip_duration_ms: Expected duration in milliseconds of the wavs.
  window_size_ms: How long each spectrogram timeslice is.
  window_stride_ms: How far to move in time between spectrogram timeslices.
  feature_bin_count: How many bins to use for the feature fingerprint.
  quantize: Whether to train the model for eight-bit deployment.
  preprocess: Spectrogram processing mode; "mfcc", "average" or "micro".
  input_wav: Path to the audio WAV file to read.
  output_c_file: Where to save the generated C source file.


def rgb_to_grayscale(images, name=None):
    with ops.name_scope(name, 'rgb_to_grayscale', [images]) as name:
        images = ops.convert_to_tensor(images, name='images')
        orig_dtype = images.dtype
        flt_image = convert_image_dtype(images, dtypes.float32)
        rgb_weights = [0.2989, 0.587, 0.114]
        gray_float = math_ops.tensordot(flt_image, rgb_weights, [-1, -1])
        gray_float = array_ops.expand_dims(gray_float, -1)
        return convert_image_dtype(gray_float, orig_dtype, name=name)

Converts one or more images from RGB to Grayscale.

Outputs a tensor of the same `DType` and rank as `images`.  The size of the
last dimension of the output is 1, containing the Grayscale value of the
pixels.

>>> original = tf.constant([[[1.0, 2.0, 3.0]]])
>>> converted = tf.image.rgb_to_grayscale(original)
>>> print(converted.numpy())
[[[1.81...]]]

Args:
  images: The RGB tensor to convert. The last dimension must have size 3 and
    should contain RGB values.
  name: A name for the operation (optional).

Returns:
  The converted grayscale image(s).


def _run_contained(task_type, task_id, fn, args, kwargs):
    is_successful = False
    return_value = None
    exc_info = None
    try:
        return_value = fn(*args, **kwargs)
        is_successful = True
        return _ProcessStatusInfo(task_type=task_type, task_id=task_id, is_successful=is_successful, exc_info=exc_info, return_value=return_value)
    except Exception:
        exc_info = sys.exc_info()
        return _ProcessStatusInfo(task_type=task_type, task_id=task_id, is_successful=is_successful, exc_info=exc_info, return_value=return_value)

Runs `fn` with `args` and `kwargs`.

The function returns _ProcessStatusInfo which captures the return value and
the exception.

Args:
  task_type: the task type.
  task_id: the task index.
  fn: the function to be run.
  args: optional positional arguments to be supplied in `fn`.
  kwargs: optional keyword arguments to be supplied in `fn`.

Returns:
  a _ProcessStatusInfo.


def get_image_features(self, pixel_values: torch.FloatTensor, pixel_mask: Optional[torch.FloatTensor]=None, vision_feature_layer: int=-1):
    vision_feature_layer = vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer
    patch_attention_mask = self._create_patch_attention_mask(pixel_mask)
    image_outputs = self.vision_tower(pixel_values, patch_attention_mask=patch_attention_mask, output_hidden_states=True)
    image_attn_mask = None
    if patch_attention_mask is not None:
        flattened_mask = patch_attention_mask.flatten(1)
        image_attn_mask = torch.logical_not(flattened_mask)
    selected_image_feature = image_outputs.hidden_states[vision_feature_layer]
    image_features = self.multi_modal_projector(selected_image_feature, attn_mask=image_attn_mask)
    return image_features

Obtains image last hidden states from the vision tower and apply multimodal projection.

Args:
    pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`):
       The tensors corresponding to the input images.
    pixel_mask (`torch.FloatTensor]`, *optional*):
        The tensors corresponding to the input image mask.
    vision_feature_layer (`Union[int, List[int]]`, *optional*):
        The index of the layer to select the vision feature. If multiple indices are provided,
        the vision feature of the corresponding indices will be concatenated to form the
        vision features.
Returns:
    image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`).


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

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

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


def _merge_run_options(self, options, incoming_options):
    options.trace_level = max(options.trace_level, incoming_options.trace_level)
    options.timeout_in_ms = max(options.timeout_in_ms, incoming_options.timeout_in_ms)
    options.inter_op_thread_pool = max(options.inter_op_thread_pool, incoming_options.inter_op_thread_pool)
    options.output_partition_graphs = max(options.output_partition_graphs, incoming_options.output_partition_graphs)
    options.debug_options.debug_tensor_watch_opts.extend(incoming_options.debug_options.debug_tensor_watch_opts)
    options.debug_options.reset_disk_byte_usage = options.debug_options.reset_disk_byte_usage or incoming_options.debug_options.reset_disk_byte_usage
    options.report_tensor_allocations_upon_oom = options.report_tensor_allocations_upon_oom or incoming_options.report_tensor_allocations_upon_oom

Merge two instances of RunOptions into the first one.

During the merger, the numerical fields including trace_level,
timeout_in_ms, inter_op_thread_pool are set to the larger one of the two.
The boolean value is set to the logical OR of the two.
debug_tensor_watch_opts of the original options is extended with that from
the incoming one.

Args:
  options: The options to merge into.
  incoming_options: The options to be merged into the first argument.


def postprocess_qa_predictions(examples, features, predictions: tuple[np.ndarray, np.ndarray], version_2_with_negative: bool=False, n_best_size: int=20, max_answer_length: int=30, null_score_diff_threshold: float=0.0, output_dir: Optional[str]=None, prefix: Optional[str]=None, log_level: Optional[int]=logging.WARNING):
    if len(predictions) != 2:
        raise ValueError('`predictions` should be a tuple with two elements (start_logits, end_logits).')
    all_start_logits, all_end_logits = predictions
    if len(predictions[0]) != len(features):
        raise ValueError(f'Got {len(predictions[0])} predictions and {len(features)} features.')
    example_id_to_index = {k: i for i, k in enumerate(examples['id'])}
    features_per_example = collections.defaultdict(list)
    for i, feature in enumerate(features):
        features_per_example[example_id_to_index[feature['example_id']]].append(i)
    all_predictions = collections.OrderedDict()
    all_nbest_json = collections.OrderedDict()
    if version_2_with_negative:
        scores_diff_json = collections.OrderedDict()
    logger.setLevel(log_level)
    logger.info(f'Post-processing {len(examples)} example predictions split into {len(features)} features.')
    for example_index, example in enumerate(tqdm(examples)):
        feature_indices = features_per_example[example_index]
        min_null_prediction = None
        prelim_predictions = []
        for feature_index in feature_indices:
            start_logits = all_start_logits[feature_index]
            end_logits = all_end_logits[feature_index]
            offset_mapping = features[feature_index]['offset_mapping']
            token_is_max_context = features[feature_index].get('token_is_max_context', None)
            feature_null_score = start_logits[0] + end_logits[0]
            if min_null_prediction is None or min_null_prediction['score'] > feature_null_score:
                min_null_prediction = {'offsets': (0, 0), 'score': feature_null_score, 'start_logit': start_logits[0], 'end_logit': end_logits[0]}
            start_indexes = np.argsort(start_logits)[-1:-n_best_size - 1:-1].tolist()
            end_indexes = np.argsort(end_logits)[-1:-n_best_size - 1:-1].tolist()
            for start_index in start_indexes:
                for end_index in end_indexes:
                    if start_index >= len(offset_mapping) or end_index >= len(offset_mapping) or offset_mapping[start_index] is None or (len(offset_mapping[start_index]) < 2) or (offset_mapping[end_index] is None) or (len(offset_mapping[end_index]) < 2):
                        continue
                    if end_index < start_index or end_index - start_index + 1 > max_answer_length:
                        continue
                    if token_is_max_context is not None and (not token_is_max_context.get(str(start_index), False)):
                        continue
                    prelim_predictions.append({'offsets': (offset_mapping[start_index][0], offset_mapping[end_index][1]), 'score': start_logits[start_index] + end_logits[end_index], 'start_logit': start_logits[start_index], 'end_logit': end_logits[end_index]})
        if version_2_with_negative and min_null_prediction is not None:
            prelim_predictions.append(min_null_prediction)
            null_score = min_null_prediction['score']
        predictions = sorted(prelim_predictions, key=lambda x: x['score'], reverse=True)[:n_best_size]
        if version_2_with_negative and min_null_prediction is not None and (not any((p['offsets'] == (0, 0) for p in predictions))):
            predictions.append(min_null_prediction)
        context = example['context']
        for pred in predictions:
            offsets = pred.pop('offsets')
            pred['text'] = context[offsets[0]:offsets[1]]
        if len(predictions) == 0 or (len(predictions) == 1 and predictions[0]['text'] == ''):
            predictions.insert(0, {'text': 'empty', 'start_logit': 0.0, 'end_logit': 0.0, 'score': 0.0})
        scores = np.array([pred.pop('score') for pred in predictions])
        exp_scores = np.exp(scores - np.max(scores))
        probs = exp_scores / exp_scores.sum()
        for prob, pred in zip(probs, predictions):
            pred['probability'] = prob
        if not version_2_with_negative:
            all_predictions[example['id']] = predictions[0]['text']
        else:
            i = 0
            while predictions[i]['text'] == '':
                i += 1
            best_non_null_pred = predictions[i]
            score_diff = null_score - best_non_null_pred['start_logit'] - best_non_null_pred['end_logit']
            scores_diff_json[example['id']] = float(score_diff)
            if score_diff > null_score_diff_threshold:
                all_predictions[example['id']] = ''
            else:
                all_predictions[example['id']] = best_non_null_pred['text']
        all_nbest_json[example['id']] = [{k: float(v) if isinstance(v, (np.float16, np.float32, np.float64)) else v for k, v in pred.items()} for pred in predictions]
    if output_dir is not None:
        if not os.path.isdir(output_dir):
            raise OSError(f'{output_dir} is not a directory.')
        prediction_file = os.path.join(output_dir, 'predictions.json' if prefix is None else f'{prefix}_predictions.json')
        nbest_file = os.path.join(output_dir, 'nbest_predictions.json' if prefix is None else f'{prefix}_nbest_predictions.json')
        if version_2_with_negative:
            null_odds_file = os.path.join(output_dir, 'null_odds.json' if prefix is None else f'{prefix}_null_odds.json')
        logger.info(f'Saving predictions to {prediction_file}.')
        with open(prediction_file, 'w') as writer:
            writer.write(json.dumps(all_predictions, indent=4) + '\n')
        logger.info(f'Saving nbest_preds to {nbest_file}.')
        with open(nbest_file, 'w') as writer:
            writer.write(json.dumps(all_nbest_json, indent=4) + '\n')
        if version_2_with_negative:
            logger.info(f'Saving null_odds to {null_odds_file}.')
            with open(null_odds_file, 'w') as writer:
                writer.write(json.dumps(scores_diff_json, indent=4) + '\n')
    return all_predictions

Post-processes the predictions of a question-answering model to convert them to answers that are substrings of the
original contexts. This is the base postprocessing functions for models that only return start and end logits.

Args:
    examples: The non-preprocessed dataset (see the main script for more information).
    features: The processed dataset (see the main script for more information).
    predictions (:obj:`Tuple[np.ndarray, np.ndarray]`):
        The predictions of the model: two arrays containing the start logits and the end logits respectively. Its
        first dimension must match the number of elements of :obj:`features`.
    version_2_with_negative (:obj:`bool`, `optional`, defaults to :obj:`False`):
        Whether or not the underlying dataset contains examples with no answers.
    n_best_size (:obj:`int`, `optional`, defaults to 20):
        The total number of n-best predictions to generate when looking for an answer.
    max_answer_length (:obj:`int`, `optional`, defaults to 30):
        The maximum length of an answer that can be generated. This is needed because the start and end predictions
        are not conditioned on one another.
    null_score_diff_threshold (:obj:`float`, `optional`, defaults to 0):
        The threshold used to select the null answer: if the best answer has a score that is less than the score of
        the null answer minus this threshold, the null answer is selected for this example (note that the score of
        the null answer for an example giving several features is the minimum of the scores for the null answer on
        each feature: all features must be aligned on the fact they `want` to predict a null answer).

        Only useful when :obj:`version_2_with_negative` is :obj:`True`.
    output_dir (:obj:`str`, `optional`):
        If provided, the dictionaries of predictions, n_best predictions (with their scores and logits) and, if
        :obj:`version_2_with_negative=True`, the dictionary of the scores differences between best and null
        answers, are saved in `output_dir`.
    prefix (:obj:`str`, `optional`):
        If provided, the dictionaries mentioned above are saved with `prefix` added to their names.
    log_level (:obj:`int`, `optional`, defaults to ``logging.WARNING``):
        ``logging`` log level (e.g., ``logging.WARNING``)


Merge several accumulators to a single accumulator.

This method takes the partial values in several accumulators and combines
them into a single accumulator. This computation must not be order-specific
(that is, merge([a, b]) must return the same result as merge([b, a]).

Args:
  accumulators: the accumulators to merge, as a list.

Returns:
  A merged accumulator.


def get_word_index(path='reuters_word_index.json'):
    origin_folder = 'https://storage.googleapis.com/tensorflow/tf-keras-datasets/'
    path = get_file(path, origin=origin_folder + 'reuters_word_index.json', file_hash='4d44cc38712099c9e383dc6e5f11a921')
    with open(path) as f:
        return json.load(f)

Retrieves a dict mapping words to their index in the Reuters dataset.

Actual word indices starts from 3, with 3 indices reserved for:
0 (padding), 1 (start), 2 (oov).

E.g. word index of 'the' is 1, but the in the actual training data, the
index of 'the' will be 1 + 3 = 4. Vice versa, to translate word indices in
training data back to words using this mapping, indices need to subtract 3.

Args:
    path: where to cache the data (relative to `~/.keras/dataset`).

Returns:
    The word index dictionary. Keys are word strings, values are their
    index.


def model_custom_sink(simplekv, KVs, final_table_name_no_ptransform, final_table_name_with_ptransform):
    final_table_name = final_table_name_no_ptransform
    with beam.Pipeline(options=PipelineOptions()) as pipeline:
        kvs = pipeline | 'CreateKVs' >> beam.Create(KVs)
        kvs | 'WriteToSimpleKV' >> beam.io.Write(SimpleKVSink(simplekv, 'http://url_to_simple_kv/', final_table_name))
    final_table_name = final_table_name_with_ptransform
    with beam.Pipeline(options=PipelineOptions()) as pipeline:
        kvs = pipeline | 'CreateKVs' >> beam.core.Create(KVs)
        kvs | 'WriteToSimpleKV' >> WriteToKVSink(simplekv, 'http://url_to_simple_kv/', final_table_name)

Demonstrates creating a new custom sink and using it in a pipeline.

Uses the new sink in an example pipeline.

Additionally demonstrates how a sink should be implemented using a
``PTransform``. This is the recommended way to develop sinks that are to be
distributed to a large number of end users.

This method runs two pipelines.

(1) A pipeline that uses ``SimpleKVSink`` directly using the ``df.Write``
    transform.
(2) A pipeline that uses a custom ``PTransform`` that wraps
    ``SimpleKVSink``.

Args:
  simplekv: an object that mocks the key-value storage.

  KVs: the set of key-value pairs to be written in the example pipeline.

  final_table_name_no_ptransform: the prefix of final set of tables to be
                                  created by the example pipeline that uses
                                  ``SimpleKVSink`` directly.

  final_table_name_with_ptransform: the prefix of final set of tables to be
                                    created by the example pipeline that uses
                                    a ``PTransform`` that wraps
                                    ``SimpleKVSink``.


class ZeroPadding3D(Layer):

    def __init__(self, padding=((1, 1), (1, 1), (1, 1)), data_format=None, **kwargs):
        super().__init__(**kwargs)
        self.data_format = backend.standardize_data_format(data_format)
        if isinstance(padding, int):
            self.padding = ((padding, padding), (padding, padding), (padding, padding))
        elif hasattr(padding, '__len__'):
            if len(padding) != 3:
                raise ValueError(f'`padding` should have 3 elements. Received: {padding}.')
            dim1_padding = argument_validation.standardize_tuple(padding[0], 2, '1st entry of padding', allow_zero=True)
            dim2_padding = argument_validation.standardize_tuple(padding[1], 2, '2nd entry of padding', allow_zero=True)
            dim3_padding = argument_validation.standardize_tuple(padding[2], 2, '3rd entry of padding', allow_zero=True)
            self.padding = (dim1_padding, dim2_padding, dim3_padding)
        else:
            raise ValueError(f'`padding` should be either an int, a tuple of 3 ints (symmetric_dim1_pad, symmetric_dim2_pad, symmetric_dim3_pad), or a tuple of 3 tuples of 2 ints ((left_dim1_pad, right_dim1_pad), (left_dim2_pad, right_dim2_pad), (left_dim3_pad, right_dim2_pad)). Received: padding={padding}.')
        self.input_spec = InputSpec(ndim=5)

    def compute_output_shape(self, input_shape):
        output_shape = list(input_shape)
        spatial_dims_offset = 2 if self.data_format == 'channels_first' else 1
        for index in range(0, 3):
            if output_shape[index + spatial_dims_offset] is not None:
                output_shape[index + spatial_dims_offset] += self.padding[index][0] + self.padding[index][1]
        return tuple(output_shape)

    def call(self, inputs):
        if self.data_format == 'channels_first':
            all_dims_padding = ((0, 0), (0, 0), *self.padding)
        else:
            all_dims_padding = ((0, 0), *self.padding, (0, 0))
        return ops.pad(inputs, all_dims_padding)

    def get_config(self):
        config = {'padding': self.padding, 'data_format': self.data_format}
        base_config = super().get_config()
        return {**base_config, **config}

Zero-padding layer for 3D data (spatial or spatio-temporal).

Example:

>>> input_shape = (1, 1, 2, 2, 3)
>>> x = np.arange(np.prod(input_shape)).reshape(input_shape)
>>> y = keras.layers.ZeroPadding3D(padding=2)(x)
>>> y.shape
(1, 5, 6, 6, 3)

Args:
    padding: Int, or tuple of 3 ints, or tuple of 3 tuples of 2 ints.
        - If int: the same symmetric padding is applied to depth, height,
          and width.
        - If tuple of 3 ints: interpreted as three different symmetric
          padding values for depth, height, and width:
          `(symmetric_dim1_pad, symmetric_dim2_pad, symmetric_dim3_pad)`.
        - If tuple of 3 tuples of 2 ints: interpreted as
          `((left_dim1_pad, right_dim1_pad), (left_dim2_pad,
          right_dim2_pad), (left_dim3_pad, right_dim3_pad))`.
    data_format: A string, one of `"channels_last"` (default) or
        `"channels_first"`. The ordering of the dimensions in the inputs.
        `"channels_last"` corresponds to inputs with shape
        `(batch_size, spatial_dim1, spatial_dim2, spatial_dim3, channels)`
        while `"channels_first"` corresponds to inputs with shape
        `(batch_size, channels, spatial_dim1, spatial_dim2, spatial_dim3)`.
        When unspecified, uses `image_data_format` value found in your Keras
        config file at `~/.keras/keras.json` (if exists). Defaults to
        `"channels_last"`.

Input shape:
    5D tensor with shape:
    - If `data_format` is `"channels_last"`:
      `(batch_size, first_axis_to_pad, second_axis_to_pad,
      third_axis_to_pad, depth)`
    - If `data_format` is `"channels_first"`:
      `(batch_size, depth, first_axis_to_pad, second_axis_to_pad,
      third_axis_to_pad)`

Output shape:
    5D tensor with shape:
    - If `data_format` is `"channels_last"`:
      `(batch_size, first_padded_axis, second_padded_axis,
      third_axis_to_pad, depth)`
    - If `data_format` is `"channels_first"`:
      `(batch_size, depth, first_padded_axis, second_padded_axis,
      third_axis_to_pad)`


def tokenize(self, text, never_split=None):
    never_split = self.never_split.union(set(never_split)) if never_split else self.never_split
    text = self._clean_text(text)
    if self.tokenize_chinese_chars:
        text = self._tokenize_chinese_chars(text)
    unicode_normalized_text = unicodedata.normalize('NFC', text)
    orig_tokens = whitespace_tokenize(unicode_normalized_text)
    split_tokens = []
    for token in orig_tokens:
        if token not in never_split:
            if self.do_lower_case:
                token = token.lower()
                if self.strip_accents is not False:
                    token = self._run_strip_accents(token)
            elif self.strip_accents:
                token = self._run_strip_accents(token)
        split_tokens.extend(self._run_split_on_punc(token, never_split))
    output_tokens = whitespace_tokenize(' '.join(split_tokens))
    return output_tokens

Basic Tokenization of a piece of text. For sub-word tokenization, see WordPieceTokenizer.

Args:
    never_split (`List[str]`, *optional*)
        Kept for backward compatibility purposes. Now implemented directly at the base class level (see
        [`PreTrainedTokenizer.tokenize`]) List of token not to split.


def fuse_awq_modules(model, quantization_config):
    if isinstance(quantization_config, dict):
        quantization_config = AwqConfig.from_dict(quantization_config)
    backend = quantization_config.backend
    modules_to_fuse = get_modules_to_fuse(model, quantization_config)
    modules_to_not_convert = getattr(quantization_config, 'modules_to_not_convert', None)
    if backend == AwqBackendPackingMethod.AUTOAWQ:
        from awq.modules.fused.attn import QuantAttentionFused
        from awq.modules.fused.mlp import QuantFusedMLP
        from awq.modules.fused.norm import FasterTransformerRMSNorm
    else:
        raise ValueError('Fusing is only supported for the AutoAWQ backend')
    fused_attention_modules = []
    for name, module in model.named_modules():
        if modules_to_not_convert is not None:
            if any((module_name_to_not_convert in name for module_name_to_not_convert in modules_to_not_convert)):
                continue
        _fuse_awq_layernorm(modules_to_fuse['layernorm'], module, FasterTransformerRMSNorm)
        if quantization_config.version != 'ipex':
            _fuse_awq_mlp(model, name, modules_to_fuse['mlp'], module, QuantFusedMLP)
        else:
            logger.info('The IPEX version AWQ does not support fuse mlp for now.')
        attention_has_been_fused = _fuse_awq_attention_layers(model, module, modules_to_fuse, name, QuantAttentionFused)
        if attention_has_been_fused:
            fused_attention_modules.append(name.split('.')[0])
    if len(fused_attention_modules) > 0:
        for module_name, module in model.named_modules():
            if any((module_name in fused_attention_modules for fused_attention_parent_module in fused_attention_modules)):
                if hasattr(module, 'config') and hasattr(module.config, '_attn_implementation'):
                    module.config._attn_implementation = 'custom'
    return model

Optionally fuse some modules in the model to speedup inference.

Args:
    model (`~PreTrainedModel`):
        The model to fuse - note this model should have been converted into AWQ format beforehand.
    quantization_config (`Union[AwqConfig, dict]`):
        The quantization configuration to use.


def from_proto(context_def, import_scope=None):
    ret = WhileContext(context_def=context_def, import_scope=import_scope)
    ret.Enter()
    for nested_def in context_def.nested_contexts:
        from_control_flow_context_def(nested_def, import_scope=import_scope)
    ret.Exit()
    return ret

Returns a `WhileContext` object created from `context_def`.

Args:
  context_def: A `WhileContextDef` protocol buffer.
  import_scope: Optional `string`. Name scope to add.

Returns:
  A `WhileContext` Python object.


def filter_devices(ads, func):
    results = []
    for ad in ads:
        if func(ad):
            results.append(ad)
    return results

Finds the AndroidDevice instances from a list that match certain
conditions.

Args:
  ads: A list of AndroidDevice instances.
  func: A function that takes an AndroidDevice object and returns True
    if the device satisfies the filter condition.

Returns:
  A list of AndroidDevice instances that satisfy the filter condition.


def backward(ctx, grad_at_output: torch.Tensor):
    multiplier, selected_experts, masked_gates = ctx.saved_tensors
    grad_at_output = grad_at_output * multiplier
    grad_at_scores_expanded = masked_gates * grad_at_output.mul(-1)
    grad_at_scores_expanded.scatter_add_(dim=-1, index=selected_experts, src=grad_at_output)
    return (grad_at_scores_expanded, None, None, None, None)

Backward pass for the custom autograd function.

Args:
    ctx: Context object with saved tensors from the forward pass.
    grad_at_output (torch.Tensor): Gradient at the output.

Returns:
    Tuple[torch.Tensor, None, None, None, None]: Gradients for the inputs.


def get_multimodal_features(self, input_ids: Optional[torch.LongTensor]=None, pixel_values: Optional[torch.FloatTensor]=None, attention_mask: Optional[torch.Tensor]=None, return_dict: Optional[bool]=None, interpolate_pos_encoding: bool=False) -> torch.FloatTensor:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    vision_outputs = self.vision_model(pixel_values=pixel_values, output_attentions=True, output_hidden_states=True, return_dict=return_dict, interpolate_pos_encoding=interpolate_pos_encoding)
    image_embeds = vision_outputs[0]
    image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long)
    text_outputs = self.text_model(input_ids=input_ids, attention_mask=attention_mask, encoder_hidden_states=image_embeds, encoder_attention_mask=image_atts, return_dict=return_dict)
    pooled_output = text_outputs[1]
    multimodal_features = self.text_projection(pooled_output)
    return multimodal_features

Returns:
    multimodal_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The multimodal embeddings
    obtained by applying the image embeddings to the text encoder using the cross-attention mechanism.

Examples:
```python
>>> from PIL import Image
>>> import requests
>>> from transformers import AutoProcessor, BlipModel

>>> model = BlipModel.from_pretrained("Salesforce/blip-image-captioning-base")
>>> processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-base")

>>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
>>> image = Image.open(requests.get(url, stream=True).raw)
>>> texts = ["a photo of a cat", "a photo of a dog"]
>>> inputs = processor(images=image, text=texts, padding=True, return_tensors="pt")

>>> multimodal_features = model.get_multimodal_features(**inputs)
```


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

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

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

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


def And(exprs):
    return simplify_exprs(exprs, _And, FALSE, TRUE)

Create a conjunction or its simplified equivalent.

This will ensure that, when an _And is returned, none of its immediate
subterms is TRUE, FALSE, or another conjunction.

Args:
  exprs: An iterable. The subterms.

Returns:
  A BooleanTerm.


def res_call(self, ns, types_ns, node, f_type, args, keywords):
    raise NotImplementedError('subclasses must implement')

Resolves the return type an external function or method call.

Args:
  ns: namespace
  types_ns: types namespace
  node: str, the function name
  f_type: types of the actual function being called, if known
  args: types of each respective argument in node.args
  keywords: types of each respective argument in node.keywords

Returns:
  Tuple (return_type, side_effect_types). The first element is just the
  return types of the function. The second element is a map from
  argument names to sets of types, and allow modelling side effects of
  functions (for example via global or nonlocal).


class TFGreedySearchEncoderDecoderOutput(ModelOutput):
    sequences: Optional[tf.Tensor] = None
    scores: Optional[Tuple[tf.Tensor]] = None
    encoder_attentions: Optional[Tuple[tf.Tensor]] = None
    encoder_hidden_states: Optional[Tuple[tf.Tensor]] = None
    decoder_attentions: Optional[Tuple[Tuple[tf.Tensor]]] = None
    cross_attentions: Optional[Tuple[Tuple[tf.Tensor]]] = None
    decoder_hidden_states: Optional[Tuple[Tuple[tf.Tensor]]] = None

Base class for outputs of encoder-decoder generation models using greedy search. Hidden states and attention
weights of the decoder (respectively the encoder) can be accessed via the encoder_attentions and the
encoder_hidden_states attributes (respectively the decoder_attentions and the decoder_hidden_states attributes)

Args:
    sequences (`tf.Tensor` of shape `(batch_size, sequence_length)`):
        The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter
        if all batches finished early due to the `eos_token_id`.
    scores (`tuple(tf.Tensor)` *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`):
        Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax)
        at each generation step. Tuple of `tf.Tensor` with up to `max_new_tokens` elements (one element for each
        generated token), with each tensor of shape `(batch_size, config.vocab_size)`.
    encoder_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):
        Tuple of `tf.Tensor` (one for each layer of the decoder) of shape `(batch_size, num_heads, sequence_length,
        sequence_length)`.
    encoder_hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape
        `(batch_size, sequence_length, hidden_size)`.
    decoder_attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):
        Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of
        `tf.Tensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`.
    cross_attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):
        Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of
        `tf.Tensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`.
    decoder_hidden_states (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of
        `tf.Tensor` of shape `(batch_size, generated_length, hidden_size)`.


def where_v2(condition, x=None, y=None, name=None):
    if x is None and y is None:
        with ops.name_scope(name, 'Where', [condition]) as name:
            condition = ops.convert_to_tensor(condition, preferred_dtype=dtypes.bool, name='condition')
            return gen_array_ops.where(condition=condition, name=name)
    elif x is not None and y is not None:
        return gen_math_ops.select_v2(condition=condition, t=x, e=y, name=name)
    else:
        raise ValueError('x and y must both be non-None or both be None.')

Returns the indices of non-zero elements, or multiplexes `x` and `y`.

This operation has two modes:

1. **Return the indices of non-zero elements** - When only
   `condition` is provided the result is an `int64` tensor where each row is
   the index of a non-zero element of `condition`. The result's shape
   is `[tf.math.count_nonzero(condition), tf.rank(condition)]`.
2. **Multiplex `x` and `y`** - When both `x` and `y` are provided the
   result has the shape of `x`, `y`, and `condition` broadcast together. The
   result is taken from `x` where `condition` is non-zero
   or `y` where `condition` is zero.

#### 1. Return the indices of non-zero elements

Note: In this mode `condition` can have a dtype of `bool` or any numeric
dtype.

If `x` and `y` are not provided (both are None):

`tf.where` will return the indices of `condition` that are non-zero,
in the form of a 2-D tensor with shape `[n, d]`, where `n` is the number of
non-zero elements in `condition` (`tf.count_nonzero(condition)`), and `d` is
the number of axes of `condition` (`tf.rank(condition)`).

Indices are output in row-major order. The `condition` can have a `dtype` of
`tf.bool`, or any numeric `dtype`.

Here `condition` is a 1-axis `bool` tensor with 2 `True` values. The result
has a shape of `[2,1]`

>>> tf.where([True, False, False, True]).numpy()
array([[0],
       [3]])

Here `condition` is a 2-axis integer tensor, with 3 non-zero values. The
result has a shape of `[3, 2]`.

>>> tf.where([[1, 0, 0], [1, 0, 1]]).numpy()
array([[0, 0],
       [1, 0],
       [1, 2]])

Here `condition` is a 3-axis float tensor, with 5 non-zero values. The output
shape is `[5, 3]`.

>>> float_tensor = [[[0.1, 0], [0, 2.2], [3.5, 1e6]],
...                 [[0,   0], [0,   0], [99,    0]]]
>>> tf.where(float_tensor).numpy()
array([[0, 0, 0],
       [0, 1, 1],
       [0, 2, 0],
       [0, 2, 1],
       [1, 2, 0]])

These indices are the same that `tf.sparse.SparseTensor` would use to
represent the condition tensor:

>>> sparse = tf.sparse.from_dense(float_tensor)
>>> sparse.indices.numpy()
array([[0, 0, 0],
       [0, 1, 1],
       [0, 2, 0],
       [0, 2, 1],
       [1, 2, 0]])

A complex number is considered non-zero if either the real or imaginary
component is non-zero:

>>> tf.where([complex(0.), complex(1.), 0+1j, 1+1j]).numpy()
array([[1],
       [2],
       [3]])

#### 2. Multiplex `x` and `y`

Note: In this mode `condition` must have a dtype of `bool`.

If `x` and `y` are also provided (both have non-None values) the `condition`
tensor acts as a mask that chooses whether the corresponding
element / row in the output should be taken from `x` (if the element in
`condition` is `True`) or `y` (if it is `False`).

The shape of the result is formed by
[broadcasting](https://docs.scipy.org/doc/numpy/reference/ufuncs.html)
together the shapes of `condition`, `x`, and `y`.

When all three inputs have the same size, each is handled element-wise.

>>> tf.where([True, False, False, True],
...          [1, 2, 3, 4],
...          [100, 200, 300, 400]).numpy()
array([  1, 200, 300,   4], dtype=int32)

There are two main rules for broadcasting:

1. If a tensor has fewer axes than the others, length-1 axes are added to the
   left of the shape.
2. Axes with length-1 are streched to match the coresponding axes of the other
   tensors.

A length-1 vector is streched to match the other vectors:

>>> tf.where([True, False, False, True], [1, 2, 3, 4], [100]).numpy()
array([  1, 100, 100,   4], dtype=int32)

A scalar is expanded to match the other arguments:

>>> tf.where([[True, False], [False, True]], [[1, 2], [3, 4]], 100).numpy()
array([[  1, 100], [100,   4]], dtype=int32)
>>> tf.where([[True, False], [False, True]], 1, 100).numpy()
array([[  1, 100], [100,   1]], dtype=int32)

A scalar `condition` returns the complete `x` or `y` tensor, with
broadcasting applied.

>>> tf.where(True, [1, 2, 3, 4], 100).numpy()
array([1, 2, 3, 4], dtype=int32)
>>> tf.where(False, [1, 2, 3, 4], 100).numpy()
array([100, 100, 100, 100], dtype=int32)

For a non-trivial example of broadcasting, here `condition` has a shape of
`[3]`, `x` has a shape of `[3,3]`, and `y` has a shape of `[3,1]`.
Broadcasting first expands the shape of `condition` to `[1,3]`. The final
broadcast shape is `[3,3]`. `condition` will select columns from `x` and `y`.
Since `y` only has one column, all columns from `y` will be identical.

>>> tf.where([True, False, True],
...          x=[[1, 2, 3],
...             [4, 5, 6],
...             [7, 8, 9]],
...          y=[[100],
...             [200],
...             [300]]
... ).numpy()
array([[ 1, 100, 3],
       [ 4, 200, 6],
       [ 7, 300, 9]], dtype=int32)

Note that if the gradient of either branch of the `tf.where` generates
a `NaN`, then the gradient of the entire `tf.where` will be `NaN`. This is
because the gradient calculation for `tf.where` combines the two branches, for
performance reasons.

A workaround is to use an inner `tf.where` to ensure the function has
no asymptote, and to avoid computing a value whose gradient is `NaN` by
replacing dangerous inputs with safe inputs.

Instead of this,

>>> x = tf.constant(0., dtype=tf.float32)
>>> with tf.GradientTape() as tape:
...   tape.watch(x)
...   y = tf.where(x < 1., 0., 1. / x)
>>> print(tape.gradient(y, x))
tf.Tensor(nan, shape=(), dtype=float32)

Although, the `1. / x` values are never used, its gradient is a `NaN` when
`x = 0`. Instead, we should guard that with another `tf.where`

>>> x = tf.constant(0., dtype=tf.float32)
>>> with tf.GradientTape() as tape:
...   tape.watch(x)
...   safe_x = tf.where(tf.equal(x, 0.), 1., x)
...   y = tf.where(x < 1., 0., 1. / safe_x)
>>> print(tape.gradient(y, x))
tf.Tensor(0.0, shape=(), dtype=float32)

See also:

* `tf.sparse` - The indices returned by the first form of `tf.where` can be
   useful in `tf.sparse.SparseTensor` objects.
* `tf.gather_nd`, `tf.scatter_nd`, and related ops - Given the
  list of indices returned from `tf.where` the `scatter` and `gather` family
  of ops can be used fetch values or insert values at those indices.
* `tf.strings.length` - `tf.string` is not an allowed dtype for the
  `condition`. Use the string length instead.

Args:
  condition: A `tf.Tensor` of dtype bool, or any numeric dtype. `condition`
    must have dtype `bool` when `x` and `y` are provided.
  x: If provided, a Tensor which is of the same type as `y`, and has a shape
    broadcastable with `condition` and `y`.
  y: If provided, a Tensor which is of the same type as `x`, and has a shape
    broadcastable with `condition` and `x`.
  name: A name of the operation (optional).

Returns:
  If `x` and `y` are provided:
    A `Tensor` with the same type as `x` and `y`, and shape that
    is broadcast from `condition`, `x`, and `y`.
  Otherwise, a `Tensor` with shape `[tf.math.count_nonzero(condition),
  tf.rank(condition)]`.

Raises:
  ValueError: When exactly one of `x` or `y` is non-None, or the shapes
    are not all broadcastable.


def get_next(self, name=None):
    self._get_next_call_count += 1
    if self._get_next_call_count > GET_NEXT_CALL_WARNING_THRESHOLD:
        warnings.warn(GET_NEXT_CALL_WARNING_MESSAGE)
    with ops.colocate_with(self._iterator_resource):
        flat_ret = gen_dataset_ops.iterator_get_next(self._iterator_resource, output_types=self._flat_tensor_types, output_shapes=self._flat_tensor_shapes, name=name)
        return structure.from_tensor_list(self._element_spec, flat_ret)

Returns the next element.

In graph mode, you should typically call this method *once* and use its
result as the input to another computation. A typical loop will then call
`tf.Session.run` on the result of that computation. The loop will terminate
when the `Iterator.get_next()` operation raises
`tf.errors.OutOfRangeError`. The following skeleton shows how to use
this method when building a training loop:

```python
dataset = ...  # A `tf.data.Dataset` object.
iterator = dataset.make_initializable_iterator()
next_element = iterator.get_next()

# Build a TensorFlow graph that does something with each element.
loss = model_function(next_element)
optimizer = ...  # A `tf.compat.v1.train.Optimizer` object.
train_op = optimizer.minimize(loss)

with tf.compat.v1.Session() as sess:
  try:
    while True:
      sess.run(train_op)
  except tf.errors.OutOfRangeError:
    pass
```

NOTE: It is legitimate to call `Iterator.get_next()` multiple times, e.g.
when you are distributing different elements to multiple devices in a single
step. However, a common pitfall arises when users call `Iterator.get_next()`
in each iteration of their training loop. `Iterator.get_next()` adds ops to
the graph, and executing each op allocates resources (including threads); as
a consequence, invoking it in every iteration of a training loop causes
slowdown and eventual resource exhaustion. To guard against this outcome, we
log a warning when the number of uses crosses a fixed threshold of
suspiciousness.

Args:
  name: (Optional.) A name for the created operation.

Returns:
  A (nested) structure of values matching `tf.data.Iterator.element_spec`.


class Adam(Optimizer):

    def __init__(self, lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=None, decay=0.0, amsgrad=False, **kwargs):
        super(Adam, self).__init__(**kwargs)
        with backend.name_scope(self.__class__.__name__):
            self.iterations = backend.variable(0, dtype='int64', name='iterations')
            self.lr = backend.variable(lr, name='lr')
            self.beta_1 = backend.variable(beta_1, name='beta_1')
            self.beta_2 = backend.variable(beta_2, name='beta_2')
            self.decay = backend.variable(decay, name='decay')
        if epsilon is None:
            epsilon = backend.epsilon()
        self.epsilon = epsilon
        self.initial_decay = decay
        self.amsgrad = amsgrad

    def _create_all_weights(self, params):
        ms = [backend.zeros(backend.int_shape(p), dtype=backend.dtype(p)) for p in params]
        vs = [backend.zeros(backend.int_shape(p), dtype=backend.dtype(p)) for p in params]
        if self.amsgrad:
            vhats = [backend.zeros(backend.int_shape(p), dtype=backend.dtype(p)) for p in params]
        else:
            vhats = [backend.zeros(1) for _ in params]
        self.weights = [self.iterations] + ms + vs + vhats
        return (ms, vs, vhats)

    def get_updates(self, loss, params):
        grads = self.get_gradients(loss, params)
        self.updates = []
        lr = self.lr
        if self.initial_decay > 0:
            lr = lr * (1.0 / (1.0 + self.decay * math_ops.cast(self.iterations, backend.dtype(self.decay))))
        with ops.control_dependencies([state_ops.assign_add(self.iterations, 1)]):
            t = math_ops.cast(self.iterations, backend.floatx())
        lr_t = lr * (backend.sqrt(1.0 - math_ops.pow(self.beta_2, t)) / (1.0 - math_ops.pow(self.beta_1, t)))
        ms, vs, vhats = self._create_all_weights(params)
        for p, g, m, v, vhat in zip(params, grads, ms, vs, vhats):
            m_t = self.beta_1 * m + (1.0 - self.beta_1) * g
            v_t = self.beta_2 * v + (1.0 - self.beta_2) * math_ops.square(g)
            if self.amsgrad:
                vhat_t = math_ops.maximum(vhat, v_t)
                p_t = p - lr_t * m_t / (backend.sqrt(vhat_t) + self.epsilon)
                self.updates.append(state_ops.assign(vhat, vhat_t))
            else:
                p_t = p - lr_t * m_t / (backend.sqrt(v_t) + self.epsilon)
            self.updates.append(state_ops.assign(m, m_t))
            self.updates.append(state_ops.assign(v, v_t))
            new_p = p_t
            if getattr(p, 'constraint', None) is not None:
                new_p = p.constraint(new_p)
            self.updates.append(state_ops.assign(p, new_p))
        return self.updates

    def get_config(self):
        config = {'lr': float(backend.get_value(self.lr)), 'beta_1': float(backend.get_value(self.beta_1)), 'beta_2': float(backend.get_value(self.beta_2)), 'decay': float(backend.get_value(self.decay)), 'epsilon': self.epsilon, 'amsgrad': self.amsgrad}
        base_config = super(Adam, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

Adam optimizer.

Default parameters follow those provided in the original paper.

Args:
  lr: float >= 0. Learning rate.
  beta_1: float, 0 < beta < 1. Generally close to 1.
  beta_2: float, 0 < beta < 1. Generally close to 1.
  epsilon: float >= 0. Fuzz factor.
    If `None`, defaults to `backend.epsilon()`.
  decay: float >= 0. Learning rate decay over each update.
  amsgrad: boolean. Whether to apply the AMSGrad variant of this algorithm
    from the paper "On the Convergence of Adam and Beyond".


def _html_tree_view_render(self, *, view: 'HtmlTreeView', name: Optional[str]=None, parent: Any=None, root_path: Optional[KeyPath]=None, **kwargs) -> Html:
    return self._html_tree_view(view=view, name=name, parent=parent, root_path=root_path, **view.get_kwargs(kwargs, self._html_tree_view_config(), root_path or KeyPath())).add_style(*self._html_tree_view_css_styles())

The entrypoint of rendering the subtree represented by this extension.

Args:
  view: The view to render the object.
  name: The name of the object.
  parent: The parent of the object.
  root_path: The key path of the object relative to the root.
  **kwargs: kwargs to pass to `view.render()` on this extension.

Returns:
  The rendered HTML.


class TableTransformerModelOutput(Seq2SeqModelOutput):
    intermediate_hidden_states: Optional[torch.FloatTensor] = None

Base class for outputs of the TABLE_TRANSFORMER encoder-decoder model. This class adds one attribute to Seq2SeqModelOutput,
namely an optional stack of intermediate decoder activations, i.e. the output of each decoder layer, each of them
gone through a layernorm. This is useful when training the model with auxiliary decoding losses.

Args:
    last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
        Sequence of hidden-states at the output of the last layer of the decoder of the model.
    decoder_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 layer) of
        shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the decoder at the output of each
        layer plus the initial embedding outputs.
    decoder_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 of the decoder, after the attention softmax, used to compute the
        weighted average in the self-attention heads.
    cross_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 of the decoder's cross-attention layer, after the attention softmax,
        used to compute the weighted average in the cross-attention heads.
    encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
        Sequence of hidden-states at the output of the last layer of the encoder of the model.
    encoder_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 layer) of
        shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the encoder at the output of each
        layer plus the initial embedding outputs.
    encoder_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 of the encoder, after the attention softmax, used to compute the
        weighted average in the self-attention heads.
    intermediate_hidden_states (`torch.FloatTensor` of shape `(config.decoder_layers, batch_size, sequence_length, hidden_size)`, *optional*, returned when `config.auxiliary_loss=True`):
        Intermediate decoder activations, i.e. the output of each decoder layer, each of them gone through a
        layernorm.


def __call__(self, images: ImageInput=None, text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]]=None, audio=None, videos=None, **kwargs: Unpack[PaliGemmaProcessorKwargs]) -> BatchFeature:
    images, text = _validate_images_text_input_order(images, text)
    output_kwargs = self._merge_kwargs(PaliGemmaProcessorKwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, **kwargs)
    suffix = output_kwargs['text_kwargs'].pop('suffix', None)
    return_token_type_ids = True if suffix is not None else False
    if images is None:
        raise ValueError('`images` are expected as arguments to a `PaliGemmaProcessor` instance.')
    if text is None:
        logger.warning_once('You are using PaliGemma without a text prefix. It will perform as a picture-captioning model.')
        text = ''
    if _is_str_or_image(text):
        text = [text]
    elif isinstance(text, list) and _is_str_or_image(text[0]):
        pass
    if text is not None and images is not None:
        if not any((IMAGE_TOKEN in sample for sample in text)):
            logger.warning('You are passing both `text` and `images` to `PaliGemmaProcessor`. The processor expects special image tokens in the text, as many tokens as there are images per each text. It is recommended to add `<image>` tokens in the very beginning of your text. For this call, we will infer how many images each text has and add special tokens.')
            if isinstance(text, List) and isinstance(images, List):
                if len(images) != len(text):
                    raise ValueError(f'Received {len(images)} images for {len(text)} prompts. Each prompt should be associated with an image or list of images.')
            if is_valid_image(images):
                images = [[images]]
            elif isinstance(images, (list, tuple)) and is_valid_image(images[0]):
                images = [[image] for image in images]
            elif not (isinstance(images, (list, tuple)) and isinstance(images[0], (list, tuple)) and is_valid_image(images[0][0])):
                raise ValueError('images must be an image, list of images or list of list of images')
            input_strings = [build_string_from_input(prompt=prompt, bos_token=self.tokenizer.bos_token, image_seq_len=self.image_seq_length, image_token=IMAGE_TOKEN, num_images=len(image_list) if isinstance(image_list, list) else 1) for prompt, image_list in zip(text, images)]
            images = make_flat_list_of_images(images)
        else:
            expanded_samples = []
            for sample in text:
                expanded_sample = sample.replace(IMAGE_TOKEN, IMAGE_TOKEN * self.image_seq_length)
                bos_rfind_index = expanded_sample.rfind(IMAGE_TOKEN)
                bos_index = bos_rfind_index + len(IMAGE_TOKEN) if bos_rfind_index != -1 else 0
                expanded_sample = expanded_sample[:bos_index] + self.tokenizer.bos_token + expanded_sample[bos_index:]
                expanded_samples.append(expanded_sample)
            input_strings = [f'{sample}\n' for sample in expanded_samples]
    if suffix is not None and _is_str_or_image(suffix):
        suffix = [suffix]
    if suffix is not None:
        suffix = [sfx + self.tokenizer.eos_token for sfx in suffix]
    pixel_values = self.image_processor(images, **output_kwargs['images_kwargs'])['pixel_values']
    return_tensors = output_kwargs['text_kwargs'].pop('return_tensors', None)
    return_mm_token_type_ids = output_kwargs['text_kwargs'].pop('return_mm_token_type_ids', None)
    inputs = self.tokenizer(input_strings, text_pair=suffix, return_token_type_ids=return_token_type_ids, **output_kwargs['text_kwargs'])
    self._check_special_mm_tokens(input_strings, inputs, modalities=['image'])
    return_data = {**inputs, 'pixel_values': pixel_values}
    if return_token_type_ids:
        labels = np.array(inputs['input_ids'])
        labels[np.array(inputs['token_type_ids']) == 0] = -100
        return_data.update({'labels': labels})
    if return_mm_token_type_ids:
        array_ids = np.array(return_data['input_ids'])
        mm_token_type_ids = np.zeros_like(return_data['input_ids'])
        mm_token_type_ids[array_ids == self.image_token_id] = 1
        return_data['mm_token_type_ids'] = mm_token_type_ids.tolist()
    return BatchFeature(data=return_data, tensor_type=return_tensors)

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` and `kwrags` arguments to
SiglipImageProcessor's [`~SiglipImageProcessor.__call__`] if `images` is not `None`. Please refer to the docstring
of the above two methods for more information.

The usage for PaliGemma fine-tuning preparation is slightly different than usual. suffix passed are suffixes to
the prompt in `text`, and will be placed after the prompt. This is because attention is handled differently for
the prefix and the suffix. For instance,
```python
image = PIL_cow_image
prompt = "answer en Where is the cow standing?"
suffix = "on the beach"
inputs = processor(text=prompt, images=image, suffix=suffix)
```
Here `inputs` will contain the `input_ids` and `token_type_ids` that follow
```python
inputs["input_ids"][:, 256:]
# tensor([[     2,   6006,    603,    573,  13910,   9980, 235336,    108,    477,   573,   8318]])
inputs["token_type_ids"][:, 256:]
tensor([[0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1]])
```
Meaning the last three tokens are of "label" ("suffix") type while the other ones are of "prefix" type.

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. In case of a NumPy array/PyTorch tensor, each image should be of shape (C, H, W), where C is a
        number of channels, H and W are image height and width.
    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).
    return_tensors (`str` or [`~utils.TensorType`], *optional*):
        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.
    suffix (`str`, `List[str]`, `List[List[str]]`):
        The suffixes or batch of suffixes to be encoded. Only necessary for finetuning. See https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/paligemma/README.md
        for more information. If your prompt is "<image> What is on the image", the suffix corresponds to the expected prediction "a cow sitting on a bench".

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`. If `suffix`
      is provided, the `input_ids` will also contain the suffix input ids.
    - **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`.
    - **labels** -- Labels compatible with training if `suffix` is not None