def _truncate(self, processed_features: Union[dict[str, np.ndarray], BatchFeature], max_length: Optional[int]=None, pad_to_multiple_of: Optional[int]=None, truncation: Optional[bool]=None):
    if not truncation:
        return processed_features
    elif truncation and max_length is None:
        raise ValueError('When setting ``truncation=True``, make sure that ``max_length`` is defined.')
    required_input = processed_features[self.model_input_names[0]]
    if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
        max_length = (max_length // pad_to_multiple_of + 1) * pad_to_multiple_of
    needs_to_be_truncated = len(required_input) > max_length
    if needs_to_be_truncated:
        processed_features[self.model_input_names[0]] = processed_features[self.model_input_names[0]][:max_length]
        if 'attention_mask' in processed_features:
            processed_features['attention_mask'] = processed_features['attention_mask'][:max_length]
    return processed_features

Truncate inputs to predefined length or max length in the batch

Args:
    processed_features(`Union[Dict[str, np.ndarray], BatchFeature]`):
        Dictionary of input values (`np.ndarray[float]`) / input vectors (`List[np.ndarray[float]]`) or batch
        of inputs values (`List[np.ndarray[int]]`) / input vectors (`List[np.ndarray[int]]`)
    max_length (`int`, *optional*):
        maximum length of the returned list and optionally padding length (see below)
    pad_to_multiple_of (`int`, *optional*) :
        Integer if set will pad the sequence to a multiple of the provided value. This is especially useful to
        enable the use of Tensor Core on NVIDIA hardware with compute capability `>= 7.5` (Volta), or on TPUs
        which benefit from having sequence lengths be a multiple of 128.
    truncation (`bool`, *optional*):
        Activates truncation to cut input sequences longer than `max_length` to `max_length`.


def enable_output_all_intermediates(fn: _F) -> _F:

    def wrapper(*args, **kwargs):
        output_all_intermediates_old = control_flow_util_v2._EXPERIMENTAL_OUTPUT_ALL_INTERMEDIATES_OVERRIDE
        control_flow_util_v2._EXPERIMENTAL_OUTPUT_ALL_INTERMEDIATES_OVERRIDE = True
        try:
            return fn(*args, **kwargs)
        finally:
            control_flow_util_v2._EXPERIMENTAL_OUTPUT_ALL_INTERMEDIATES_OVERRIDE = output_all_intermediates_old
    return wrapper

Force-enable outputting all intermediates from functional control flow ops.

Args:
  fn: the function to be wrapped

Returns:
  The wrapped function


def initialize(job_addresses=None, num_processes=None, process_id=None):
    if job_addresses is None and 'KERAS_DISTRIBUTION_JOB_ADDRESSES' in os.environ:
        job_addresses = os.environ['KERAS_DISTRIBUTION_JOB_ADDRESSES']
    if num_processes is None and 'KERAS_DISTRIBUTION_NUM_PROCESSES' in os.environ:
        num_processes = int(os.environ['KERAS_DISTRIBUTION_NUM_PROCESSES'])
    if process_id is None and 'KERAS_DISTRIBUTION_PROCESS_ID' in os.environ:
        process_id = int(os.environ['KERAS_DISTRIBUTION_PROCESS_ID'])
    distribution_lib.initialize(job_addresses, num_processes, process_id)

Initialize the distribution system for multi-host/process setting.

Calling `initialize` will prepare the backend for execution on multi-host
GPU or TPUs. It should be called before any computations.

Note that the parameters can also be injected via environment variables,
which can be better controlled by the launch script at startup time.
For certain backend that also rely on the environment variables to
configure, Keras will properly forward them.

Args:
    job_addresses: string. Comma separated IP addresses for all the jobs
        that will form the whole computation cluster. Note that for JAX
        backend, only the address for job 0 (coodinator) is needed. For
        certain runtime like cloud TPU, this value can be `None`, and the
        backend will figure it out with the TPU environment variables. You
        can also config this value via environment variable
        `KERAS_DISTRIBUTION_JOB_ADDRESSES`.
    num_processes: int. The number of worker/processes that will form the
        whole computation cluster. For certain runtime like cloud TPU, this
        value can be `None`, and the backend will figure it out with the TPU
        environment variables. You can also configure this value via
        environment variable `KERAS_DISTRIBUTION_NUM_PROCESSES`.
    process_id: int. The ID number of the current worker/process. The value
        should be ranged from `0` to `num_processes - 1`. `0` will indicate
        the current worker/process is the master/coordinate job. You can
        also configure this value via environment variable
        `KERAS_DISTRIBUTION_PROCESS_ID`.

    Example:
        Suppose there are two GPU processes, and process 0 is running at
        address `10.0.0.1:1234`, and process 1 is running at address
        `10.0.0.2:2345`. To configure such cluster, you can run

    On process 0:
    ```python
    keras.distribute.initialize(
        job_addresses="10.0.0.1:1234,10.0.0.2:2345",
        num_processes=2,
        process_id=0)
    ```

    On process 1:
    ```python
    keras.distribute.initialize(
        job_addresses="10.0.0.1:1234,10.0.0.2:2345",
        num_processes=2,
        process_id=1)
    ```

    or via the environment variables:
    On process 0:
    ```python
    os.environ[
        "KERAS_DISTRIBUTION_JOB_ADDRESSES"] = "10.0.0.1:1234,10.0.0.2:2345"
    os.environ["KERAS_DISTRIBUTION_NUM_PROCESSES"] = "2"
    os.environ["KERAS_DISTRIBUTION_PROCESS_ID"] = "0"
    keras.distribute.initialize()
    ```

    On process 1:
    ```python
    os.environ[
        "KERAS_DISTRIBUTION_JOB_ADDRESSES"] = "10.0.0.1:1234,10.0.0.2:2345"
    os.environ["KERAS_DISTRIBUTION_NUM_PROCESSES"] = "2"
    os.environ["KERAS_DISTRIBUTION_PROCESS_ID"] = "1"
    keras.distribute.initialize()
    ```

    Also note that for JAX backend, the `job_addresses` can be further
    reduced to just the master/coordinator address, which is
    `10.0.0.1:1234`.


def enter_cond_section(self, section_id):
    assert section_id not in self.cond_entry
    assert section_id not in self.cond_leaves
    self.cond_leaves[section_id] = []

Enters a conditional section.

Conditional sections define an entry node, and one or more branches.

Args:
  section_id: Hashable, the same node that will be used in calls to the
    section_id arg passed to new_cond_branch


def compute(i, tas):
    elems_value_batchable = [ta.read(i) for ta in elems_batchable_ta]
    elems_value_flat = _elems_value_batchable_to_flat(elems_value_batchable, elems_flat_signature)
    elems_value = elems_unflatten(elems_value_flat)
    ag_ctx = autograph_ctx.control_status_ctx()
    autographed_fn = autograph.tf_convert(fn, ag_ctx)
    result_value = autographed_fn(elems_value)
    nest.assert_same_structure(fn_output_signature or elems, result_value)
    result_value_flat = nest.flatten(result_value)
    result_value_batchable = _result_value_flat_to_batchable(result_value_flat, result_flat_signature)
    tas = [ta.write(i, value) for ta, value in zip(tas, result_value_batchable)]
    return (i + 1, tas)

The loop body of map_fn.

Args:
  i: the loop counter
  tas: the flat TensorArray accumulator list

Returns:
  (i + 1, tas): the updated counter + updated TensorArrays

Raises:
  TypeError: if fn_output_signature and result_value structure don't match
  ValueType: if fn_output_signature and result_value lengths don't match


class TFBaseModelOutputWithNoAttention(ModelOutput):
    last_hidden_state: Optional[tf.Tensor] = None
    hidden_states: Optional[Tuple[tf.Tensor, ...]] = None

Base class for model's outputs, with potential hidden states.

Args:
    last_hidden_state (`tf.Tensor` shape `(batch_size, num_channels, height, width)`):
        Sequence of hidden-states at the output of the last layer of the model.
    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, if the model has an embedding layer, + one for
        the output of each layer) of shape `(batch_size, num_channels, height, width)`.

        Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.


def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, output_attentions: bool=False, query_states: Optional[torch.Tensor]=None, relative_pos: Optional[torch.Tensor]=None, rel_embeddings: Optional[torch.Tensor]=None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
    if query_states is None:
        qp = self.in_proj(hidden_states)
        query_layer, key_layer, value_layer = self.transpose_for_scores(qp).chunk(3, dim=-1)
    else:
        ws = self.in_proj.weight.chunk(self.num_attention_heads * 3, dim=0)
        qkvw = [torch.cat([ws[i * 3 + k] for i in range(self.num_attention_heads)], dim=0) for k in range(3)]
        q = torch.matmul(qkvw[0], query_states.t().to(dtype=qkvw[0].dtype))
        k = torch.matmul(qkvw[1], hidden_states.t().to(dtype=qkvw[1].dtype))
        v = torch.matmul(qkvw[2], hidden_states.t().to(dtype=qkvw[2].dtype))
        query_layer, key_layer, value_layer = [self.transpose_for_scores(x) for x in [q, k, v]]
    query_layer = query_layer + self.transpose_for_scores(self.q_bias[None, None, :])
    value_layer = value_layer + self.transpose_for_scores(self.v_bias[None, None, :])
    rel_att: int = 0
    scale_factor = 1 + len(self.pos_att_type)
    scale = scaled_size_sqrt(query_layer, scale_factor)
    query_layer = query_layer / scale.to(dtype=query_layer.dtype)
    attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
    if self.relative_attention and rel_embeddings is not None and (relative_pos is not None):
        rel_embeddings = self.pos_dropout(rel_embeddings)
        rel_att = self.disentangled_att_bias(query_layer, key_layer, relative_pos, rel_embeddings, scale_factor)
    if rel_att is not None:
        attention_scores = attention_scores + rel_att
    if self.head_logits_proj is not None:
        attention_scores = self.head_logits_proj(attention_scores.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
    attention_mask = attention_mask.bool()
    attention_scores = attention_scores.masked_fill(~attention_mask, torch.finfo(query_layer.dtype).min)
    attention_probs = nn.functional.softmax(attention_scores, dim=-1)
    attention_probs = self.dropout(attention_probs)
    if self.head_weights_proj is not None:
        attention_probs = self.head_weights_proj(attention_probs.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
    context_layer = torch.matmul(attention_probs, value_layer)
    context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
    new_context_layer_shape = context_layer.size()[:-2] + (-1,)
    context_layer = context_layer.view(new_context_layer_shape)
    if not output_attentions:
        return (context_layer, None)
    return (context_layer, attention_probs)

Call the module

Args:
    hidden_states (`torch.FloatTensor`):
        Input states to the module usually the output from previous layer, it will be the Q,K and V in
        *Attention(Q,K,V)*

    attention_mask (`torch.BoolTensor`):
        An attention mask matrix of shape [*B*, *N*, *N*] where *B* is the batch size, *N* is the maximum
        sequence length in which element [i,j] = *1* means the *i* th token in the input can attend to the *j*
        th token.

    output_attentions (`bool`, *optional*):
        Whether return the attention matrix.

    query_states (`torch.FloatTensor`, *optional*):
        The *Q* state in *Attention(Q,K,V)*.

    relative_pos (`torch.LongTensor`):
        The relative position encoding between the tokens in the sequence. It's of shape [*B*, *N*, *N*] with
        values ranging in [*-max_relative_positions*, *max_relative_positions*].

    rel_embeddings (`torch.FloatTensor`):
        The embedding of relative distances. It's a tensor of shape [\(2 \times
        \text{max_relative_positions}\), *hidden_size*].


def load_library(library_location):
    if os.path.exists(library_location):
        if os.path.isdir(library_location):
            directory_contents = os.listdir(library_location)
            kernel_libraries = [os.path.join(library_location, f) for f in directory_contents if _is_shared_object(f)]
        else:
            kernel_libraries = [library_location]
        for lib in kernel_libraries:
            py_tf.TF_LoadLibrary(lib)
    else:
        raise OSError(errno.ENOENT, 'The file or folder to load kernel libraries from does not exist.', library_location)

Loads a TensorFlow plugin.

"library_location" can be a path to a specific shared object, or a folder.
If it is a folder, all shared objects that are named "libtfkernel*" will be
loaded. When the library is loaded, kernels registered in the library via the
`REGISTER_*` macros are made available in the TensorFlow process.

Args:
  library_location: Path to the plugin or the folder of plugins.
    Relative or absolute filesystem path to a dynamic library file or folder.

Returns:
  None

Raises:
  OSError: When the file to be loaded is not found.
  RuntimeError: when unable to load the library.


def _GetMaxSizeFromNestedMaximumIterations(value, while_ctxt):
    value_name = value.name
    curr_ctxt = ops.get_default_graph()._get_control_flow_context()
    curr_ctxt_name = curr_ctxt.name if curr_ctxt is not None else ''
    max_size = constant_op.constant(1)
    while while_ctxt not in (None, curr_ctxt):
        max_iter = while_ctxt.maximum_iterations
        if max_iter is None:
            raise ValueError("Cannot create a gradient accumulator for tensor '%s' inside XLA while_loop because maximum_iterations was not passed to the tf.while_loop call ('%s')." % (value_name, while_ctxt.name))
        max_iter_ctxt = max_iter.op._get_control_flow_context()
        if util.IsContainingContext(curr_ctxt, max_iter_ctxt):
            max_size *= max_iter
        else:
            const_max_iter = tensor_util.constant_value(max_iter)
            if const_max_iter is None:
                raise ValueError("Cannot create a gradient accumulator for tensor '%s' inside XLA while_loop. maximum_iterations tensor '%s' for while_loop context '%s' must be statically known (e.g. a constant value or known shape dimension), or be defined at or outside the while loop context '%s' (currently defined in '%s')." % (value_name, max_iter.name, while_ctxt.name, curr_ctxt_name, max_iter_ctxt.name))
            max_size *= const_max_iter
        while_ctxt = util.GetContainingWhileContext(while_ctxt.outer_context, stop_ctxt=curr_ctxt)
    return max_size

Calculate a max_size for use by stack ops inside an XLA while_loop.

Args:
  value: The value inside the while_loop forward context.  Used for printing
    error messages.
  while_ctxt: The forward context inside which value resides.  This does not
    always match the value's immediate context, as `value` may be inside e.g.
    a cond context inside the while_loop.

Returns:
  A tensor containing the `max_size` to feed to a Stack initializer.

Raises:
  ValueError: If `value` is nested inside a `while_loop` that either
    lacks a `maximum_iterations` parameter, or the `maximum_iterations`
    parameter:

      - is inside a `while_loop` that is a parent of the calling context, and
      - cannot be evaluated at graph build time to a constant.


def _get_column_values(table, col_index):
    index_to_values = {}
    for row_index, row in table.iterrows():
        text = normalize_for_match(row[col_index].text)
        index_to_values[row_index] = list(_get_numeric_values(text))
    return index_to_values

Parses text in column and returns a dict mapping row_index to values. This is the _get_column_values function from
number_annotation_utils.py of the original implementation

Args:
  table: Pandas dataframe
  col_index: integer, indicating the index of the column to get the numeric values of


def reset_from_seed(self, seed):
    state = create_rng_state(seed, self.algorithm)
    self._state_var.assign(state)

Resets the generator by a new seed.

See `from_seed` for the meaning of "seed".

Args:
  seed: the new seed.


class FlaxTopPLogitsWarper(FlaxLogitsWarper):

    def __init__(self, top_p: float, filter_value: float=-float('Inf'), min_tokens_to_keep: int=1):
        if not isinstance(top_p, float) or (top_p < 0 or top_p > 1.0):
            raise ValueError(f'`top_p` has to be a float > 0 and < 1, but is {top_p}')
        if not isinstance(min_tokens_to_keep, int) or min_tokens_to_keep < 1:
            raise ValueError(f'`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}')
        self.top_p = top_p
        self.filter_value = filter_value
        self.min_tokens_to_keep = min_tokens_to_keep

    def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray:
        topk_scores, topk_indices = lax.top_k(scores, scores.shape[-1])
        mask_scores = jnp.full_like(scores, self.filter_value)
        cumulative_probs = jax.nn.softmax(topk_scores, axis=-1).cumsum(axis=-1)
        score_mask = cumulative_probs < self.top_p
        score_mask = jnp.roll(score_mask, 1)
        score_mask |= score_mask.at[:, 0].set(True)
        score_mask = score_mask.at[:, :self.min_tokens_to_keep].set(True)
        topk_next_scores = jnp.where(score_mask, topk_scores, mask_scores)
        next_scores = jax.lax.sort_key_val(topk_indices, topk_next_scores)[-1]
        return next_scores

[`FlaxLogitsWarper`] that performs top-p, i.e. restricting to top tokens summing to prob_cut_off <= prob_cut_off.

Args:
    top_p (`float`):
        If set to < 1, only the smallest set of most probable tokens with probabilities that add up to `top_p` or
        higher are kept for generation.
    filter_value (`float`, *optional*, defaults to -inf):
        All filtered values will be set to this float value.
    min_tokens_to_keep (`int`, *optional*, defaults to 1):
        Minimum number of tokens that cannot be filtered.


def string_length(input, name=None, unit='BYTE'):
    return gen_string_ops.string_length(input, unit=unit, name=name)

Computes the length of each string given in the input tensor.

>>> strings = tf.constant(['Hello','TensorFlow', '🙂'])
>>> tf.strings.length(strings).numpy() # default counts bytes
array([ 5, 10, 4], dtype=int32)
>>> tf.strings.length(strings, unit="UTF8_CHAR").numpy()
array([ 5, 10, 1], dtype=int32)

Args:
  input: A `Tensor` of type `string`. The strings for which to compute the
    length for each element.
  name: A name for the operation (optional).
  unit: An optional `string` from: `"BYTE", "UTF8_CHAR"`. Defaults to
    `"BYTE"`. The unit that is counted to compute string length.  One of:
      `"BYTE"` (for the number of bytes in each string) or `"UTF8_CHAR"` (for
      the number of UTF-8 encoded Unicode code points in each string). Results
      are undefined if `unit=UTF8_CHAR` and the `input` strings do not contain
      structurally valid UTF-8.

Returns:
  A `Tensor` of type `int32`, containing the length of the input string in
  the same element of the input tensor.


def get_dataset(categories: list, split: str='train'):
    labels = ['sadness', 'joy', 'love', 'anger', 'fear', 'surprise']
    label_map = {class_name: class_id for class_id, class_name in enumerate(labels)}
    labels_subset = np.array([label_map[class_name] for class_name in categories])
    emotion_dataset = load_dataset('emotion', download_mode='force_redownload')
    X, y = (np.array(emotion_dataset[split]['text']), np.array(emotion_dataset[split]['label']))
    subclass_idxs = [idx for idx, label in enumerate(y) if label in labels_subset]
    X_subset, y_subset = (X[subclass_idxs], y[subclass_idxs])
    return (X_subset.tolist(), y_subset.tolist())

Takes a list of categories and a split (train/test/dev) and returns the
corresponding subset of the dataset

Args:
  categories (list): list of emotion categories to use
  split (str): The split of the dataset to use. Can be either "train", "dev", or "test".
  Defaults to train

Returns:
  A list of text and a list of labels


class EncoderRepetitionPenaltyLogitsProcessor(LogitsProcessor):

    def __init__(self, penalty: float, encoder_input_ids: torch.LongTensor):
        if not isinstance(penalty, float) or not penalty > 0:
            raise ValueError(f'`penalty` has to be a strictly positive float, but is {penalty}')
        self.penalty = 1 / penalty
        self.encoder_input_ids = encoder_input_ids

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        score = torch.gather(scores, 1, self.encoder_input_ids)
        score = torch.where(score < 0, score * self.penalty, score / self.penalty)
        scores_processed = scores.scatter(1, self.encoder_input_ids, score)
        return scores_processed

[`LogitsProcessor`] that works similarly to [`RepetitionPenaltyLogitsProcessor`], but with an *inverse* penalty
that is applied to the tokens present in the prompt. In other words, a penalty above 1.0 increases the odds of
selecting tokens that were present in the prompt.

It was designed to avoid hallucination in input-grounded tasks, like summarization. Although originally intended
for encoder-decoder models, it can also be used with decoder-only models like LLMs.

Args:
    penalty (`float`):
        The parameter for repetition penalty. 1.0 means no penalty. Above 1.0 rewards prompt tokens. Between 0.0
        and 1.0 penalizes prompt tokens.
    encoder_input_ids (`torch.LongTensor`):
        The encoder_input_ids that should be repeated within the decoder ids.

Examples:

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

>>> tokenizer = AutoTokenizer.from_pretrained("bigscience/bloomz-560m")
>>> model = AutoModelForCausalLM.from_pretrained("bigscience/bloomz-560m")

>>> inputs = tokenizer(["Alice and Bob. The third member's name was"], return_tensors="pt")
>>> gen_out = model.generate(**inputs)
>>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0])
Alice and Bob. The third member's name was not mentioned.

>>> # With the `encoder_repetition_penalty` argument we can trigger this logits processor in `generate`, which can
>>> # promote the use of prompt tokens ("Bob" in this example)
>>> gen_out = model.generate(**inputs, encoder_repetition_penalty=1.2)
>>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0])
Alice and Bob. The third member's name was Bob. The third member's name was Bob.
```


def dropout_v2(x, rate, noise_shape=None, seed=None, name=None):
    uniform_sampler = functools.partial(random_ops.random_uniform, seed=seed)

    def dummy_rng_step():
        random_seed.get_seed(seed)
    return _dropout(x=x, rate=rate, noise_shape=noise_shape, uniform_sampler=uniform_sampler, dummy_rng_step=dummy_rng_step, name=name, default_name='dropout')

Computes dropout: randomly sets elements to zero to prevent overfitting.

Warning: You should consider using
`tf.nn.experimental.stateless_dropout` instead of this function. The
difference between `tf.nn.experimental.stateless_dropout` and this
function is analogous to the difference between
`tf.random.stateless_uniform` and `tf.random.uniform`. Please see
[Random number
generation](https://www.tensorflow.org/guide/random_numbers) guide
for a detailed description of the various RNG systems in TF. As the
guide states, legacy stateful RNG ops like `tf.random.uniform` and
`tf.nn.dropout` are not deprecated yet but highly discouraged,
because their states are hard to control.

Note: The behavior of dropout has changed between TensorFlow 1.x and 2.x.
When converting 1.x code, please use named arguments to ensure behavior stays
consistent.

See also: `tf.keras.layers.Dropout` for a dropout layer.

[Dropout](https://arxiv.org/abs/1207.0580) is useful for regularizing DNN
models. Inputs elements are randomly set to zero (and the other elements are
rescaled). This encourages each node to be independently useful, as it cannot
rely on the output of other nodes.

More precisely: With probability `rate` elements of `x` are set to `0`.
The remaining elements are scaled up by `1.0 / (1 - rate)`, so that the
expected value is preserved.

>>> tf.random.set_seed(0)
>>> x = tf.ones([3,5])
>>> tf.nn.dropout(x, rate = 0.5, seed = 1).numpy()
array([[2., 0., 0., 2., 2.],
     [2., 2., 2., 2., 2.],
     [2., 0., 2., 0., 2.]], dtype=float32)

>>> tf.random.set_seed(0)
>>> x = tf.ones([3,5])
>>> tf.nn.dropout(x, rate = 0.8, seed = 1).numpy()
array([[0., 0., 0., 5., 5.],
     [0., 5., 0., 5., 0.],
     [5., 0., 5., 0., 5.]], dtype=float32)

>>> tf.nn.dropout(x, rate = 0.0) == x
<tf.Tensor: shape=(3, 5), dtype=bool, numpy=
  array([[ True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True]])>

By default, each element is kept or dropped independently.  If `noise_shape`
is specified, it must be
[broadcastable](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html)
to the shape of `x`, and only dimensions with `noise_shape[i] == shape(x)[i]`
will make independent decisions. This is useful for dropping whole
channels from an image or sequence. For example:

>>> tf.random.set_seed(0)
>>> x = tf.ones([3,10])
>>> tf.nn.dropout(x, rate = 2/3, noise_shape=[1,10], seed=1).numpy()
array([[0., 0., 0., 3., 3., 0., 3., 3., 3., 0.],
     [0., 0., 0., 3., 3., 0., 3., 3., 3., 0.],
     [0., 0., 0., 3., 3., 0., 3., 3., 3., 0.]], dtype=float32)

Args:
  x: A floating point tensor.
  rate: A scalar `Tensor` with the same type as x. The probability
    that each element is dropped. For example, setting rate=0.1 would drop
    10% of input elements.
  noise_shape: A 1-D integer `Tensor`, representing the
    shape for randomly generated keep/drop flags.
  seed: A Python integer. Used to create random seeds. See
    `tf.random.set_seed` for behavior.
  name: A name for this operation (optional).

Returns:
  A Tensor of the same shape of `x`.

Raises:
  ValueError: If `rate` is not in `[0, 1)` or if `x` is not a floating point
    tensor. `rate=1` is disallowed, because the output would be all zeros,
    which is likely not what was intended.


def _encode_required_field(self, name: str, containing_type_builder: expressions.Builder, builder: expressions.Builder, element_definition: message.Message) -> Optional[validation_pb2.SqlRequirement]:
    element = cast(Any, element_definition)
    if not _is_elem_supported(element):
        return None
    field_name = _last_path_token(builder)
    min_size = element.min.value
    max_size = element.max.value
    element_count = builder.count()
    query_list = []
    if _fhir_path_data_types.is_collection(builder.return_type) and max_size.isdigit():
        query_list.append(element_count <= int(max_size))
    if min_size == 1:
        query_list.append(builder.exists())
    elif min_size > 0:
        query_list.append(element_count >= min_size)
    if not query_list:
        return None
    constraint_key = f'{name}-cardinality-is-valid'
    description = f'The length of {name} must be maximum {max_size} and minimum {min_size}.'
    fhir_path_builder = query_list[0]
    for query in query_list[1:]:
        fhir_path_builder = fhir_path_builder & query
    if constraint_key in self._options.skip_keys:
        return None
    type_codes = _utils.element_type_codes(element)
    if 'Reference' not in type_codes and (not _SKIP_TYPE_CODES.isdisjoint(type_codes)):
        return None
    result = self._encode_fhir_path_builder_constraint(fhir_path_builder, containing_type_builder)
    if result is None:
        return None
    element_definition_path = self._abs_path_invocation(containing_type_builder)
    constraint_key_column_name: str = _key_to_sql_column_name(_path_to_sql_column_name(constraint_key))
    column_name_base: str = _path_to_sql_column_name(element_definition_path)
    column_name = f'{column_name_base}_{constraint_key_column_name}'
    requirement = validation_pb2.SqlRequirement(column_name=column_name, sql_expression=result.sql, fhir_path_sql_expression=result.fhir_path_sql, severity=validation_pb2.ValidationSeverity.SEVERITY_ERROR, type=validation_pb2.ValidationType.VALIDATION_TYPE_CARDINALITY, element_path=element_definition_path, description=description, fhir_path_key=constraint_key, fhir_path_expression=result.builder.fhir_path, fields_referenced_by_expression=[field_name])
    return requirement

Returns `SqlRequirement` for the required field passed.

Args:
  name: name of the constraint key.
  containing_type_builder: The builder of the Structure definition for the
    required field.
  builder: The builder containing the element to encode required field for.
  element_definition: Element definition of the builder.

Returns:
  A `SqlRequirement` representing the requirement generated from
  the element.


def make_pixel_mask(image: np.ndarray, output_size: Tuple[int, int], input_data_format: Optional[Union[str, ChannelDimension]]=None) -> np.ndarray:
    input_height, input_width = get_image_size(image, channel_dim=input_data_format)
    mask = np.zeros(output_size, dtype=np.int64)
    mask[:input_height, :input_width] = 1
    return mask

Make a pixel mask for the image, where 1 indicates a valid pixel and 0 indicates padding.
Args:
    image (`np.ndarray`):
        Image to make the pixel mask for.
    output_size (`Tuple[int, int]`):
        Output size of the mask.


def stft(signals, frame_length, frame_step, fft_length=None, window_fn=window_ops.hann_window, pad_end=False, name=None):
    with ops.name_scope(name, 'stft', [signals, frame_length, frame_step]):
        signals = ops.convert_to_tensor(signals, name='signals')
        signals.shape.with_rank_at_least(1)
        frame_length = ops.convert_to_tensor(frame_length, name='frame_length')
        frame_length.shape.assert_has_rank(0)
        frame_step = ops.convert_to_tensor(frame_step, name='frame_step')
        frame_step.shape.assert_has_rank(0)
        if fft_length is None:
            fft_length = _enclosing_power_of_two(frame_length)
        else:
            fft_length = ops.convert_to_tensor(fft_length, name='fft_length')
        framed_signals = shape_ops.frame(signals, frame_length, frame_step, pad_end=pad_end)
        if window_fn is not None:
            window = window_fn(frame_length, dtype=framed_signals.dtype)
            framed_signals *= window
        return fft_ops.rfft(framed_signals, [fft_length])

Computes the [Short-time Fourier Transform][stft] of `signals`.

Implemented with TPU/GPU-compatible ops and supports gradients.

Args:
  signals: A `[..., samples]` `float32`/`float64` `Tensor` of real-valued
    signals.
  frame_length: An integer scalar `Tensor`. The window length in samples.
  frame_step: An integer scalar `Tensor`. The number of samples to step.
  fft_length: An integer scalar `Tensor`. The size of the FFT to apply.
    If not provided, uses the smallest power of 2 enclosing `frame_length`.
  window_fn: A callable that takes a window length and a `dtype` keyword
    argument and returns a `[window_length]` `Tensor` of samples in the
    provided datatype. If set to `None`, no windowing is used.
  pad_end: Whether to pad the end of `signals` with zeros when the provided
    frame length and step produces a frame that lies partially past its end.
  name: An optional name for the operation.

Returns:
  A `[..., frames, fft_unique_bins]` `Tensor` of `complex64`/`complex128`
  STFT values where `fft_unique_bins` is `fft_length // 2 + 1` (the unique
  components of the FFT).

Raises:
  ValueError: If `signals` is not at least rank 1, `frame_length` is
    not scalar, or `frame_step` is not scalar.

[stft]: https://en.wikipedia.org/wiki/Short-time_Fourier_transform


class JanusVisionConfig(PretrainedConfig):
    model_type = 'janus_vision_model'
    base_config_key = 'vision_config'

    def __init__(self, hidden_size=1024, num_hidden_layers=24, num_attention_heads=16, num_channels=3, patch_size=16, image_size=384, attention_dropout=0.0, layer_norm_eps=1e-06, hidden_act='gelu', mlp_ratio=4.0, attention_bias=True, hidden_dropout_rate=0.0, projection_dim=2048, projection_dropout=0.0, use_qk_norm=False, initializer_range=0.02, depth=2, num_image_tokens=576, **kwargs):
        super().__init__(**kwargs)
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.num_channels = num_channels
        self.patch_size = patch_size
        self.image_size = image_size
        self.attention_dropout = attention_dropout
        self.layer_norm_eps = layer_norm_eps
        self.hidden_act = hidden_act
        self.mlp_ratio = mlp_ratio
        self.attention_bias = attention_bias
        self.hidden_dropout_rate = hidden_dropout_rate
        self.projection_dim = projection_dim
        self.projection_dropout = projection_dropout
        self.use_qk_norm = use_qk_norm
        self.initializer_range = initializer_range
        self.depth = depth
        self.num_image_tokens = num_image_tokens

This is the configuration class to store the configuration of a [`JanusVisionModel`]. It is used to instantiate a
`JanusVisionModel` according to the specified arguments, defining the model architecture.

Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
documentation from [`PretrainedConfig`] for more information.
Args:
    hidden_size (`int`, *optional*, defaults to 1024):
        Dimensionality of the encoder layers and the pooler layer.
    num_hidden_layers (`int`, *optional*, defaults to 24):
        Number of hidden layers in the Transformer encoder.
    num_attention_heads (`int`, *optional*, defaults to 16):
        Number of attention heads for each attention layer in the Transformer encoder.
    num_channels (`int`, *optional*, defaults to 3):
        The number of input channels.
    patch_size (`int`, *optional*, defaults to 16):
        The size (resolution) of each patch.
    image_size (`int`, *optional*, defaults to 384):
        The size (resolution) of each image.
    attention_dropout (`float`, *optional*, defaults to 0.0):
        Dropout probability for attention weights.
    layer_norm_eps (`float`, *optional*, defaults to 1e-06):
        The epsilon used by the layer normalization layers.
    hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
        The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
        `"relu"`, `"selu"`, and `"gelu_new"` are supported.
    mlp_ratio (`float`, *optional*, defaults to 4.0):
        Ratio of MLP hidden dimensionality to embedding dimensionality.
    attention_bias (`bool`, *optional*, defaults to `True`):
        Whether to add a bias to the queries, keys, and values in the attention layers.
    hidden_dropout_rate (`float`, *optional*, defaults to 0.0):
        The dropout probability for fully connected layers in the encoder.
    projection_dim (`int`, *optional*, defaults to 2048):
        Dimensionality of the MLP projection head.
    projection_dropout (`float`, *optional*, defaults to 0.0):
        Dropout probability for the projection layer.
    use_qk_norm (`bool`, *optional*, defaults to `False`):
        Whether to normalize the query and key matrices.
    initializer_range (`float`, *optional*, defaults to 0.02):
        The standard deviation of the truncated normal initializer for initializing all weight matrices.
    depth (`int`, *optional*, defaults to 2):
        Number of hidden layers in the aligner module.
    num_image_tokens (`int`, *optional*, defaults to 576):
        Number of image tokens.


def calc_control_outputs(self, graph):
    control_outputs = {}
    for op in graph.get_operations():
        for control_input in op.control_inputs:
            if control_input not in control_outputs:
                control_outputs[control_input] = set()
            control_outputs[control_input].add(op)
    return control_outputs

Returns the map of control_outputs for a given graph.

Args:
  graph: The graph to parse.

Returns:
  A map of the control outputs.


def bessel_k0e(x, name=None):
    with ops.name_scope(name, 'bessel_k0e', [x]):
        return gen_special_math_ops.bessel_k0e(x)

Computes the Bessel k0e function of `x` element-wise.

Modified Bessel function of order 0.

>>> tf.math.special.bessel_k0e([0.5, 1., 2., 4.]).numpy()
array([1.52410939, 1.14446308, 0.84156822, 0.60929767], dtype=float32)

Args:
  x: A `Tensor` or `SparseTensor`. Must be one of the following types: `half`,
    `float32`, `float64`.
  name: A name for the operation (optional).

Returns:
  A `Tensor` or `SparseTensor`, respectively. Has the same type as `x`.

@compatibility(scipy)
Equivalent to scipy.special.k0e
@end_compatibility


def write_to_pubsub(pcoll, *, topic: str, format: str, schema: Optional[Any]=None, attributes: Optional[Iterable[str]]=None, attributes_map: Optional[str]=None, id_attribute: Optional[str]=None, timestamp_attribute: Optional[str]=None):
    input_schema = schemas.schema_from_element_type(pcoll.element_type)
    extra_fields: list[str] = []
    if isinstance(attributes, str):
        attributes = [attributes]
    if attributes:
        extra_fields.extend(attributes)
    if attributes_map:
        extra_fields.append(attributes_map)

    def attributes_extractor(row):
        if attributes_map:
            attribute_values = dict(getattr(row, attributes_map))
        else:
            attribute_values = {}
        if attributes:
            attribute_values.update({attr: getattr(row, attr) for attr in attributes})
        return attribute_values
    schema_names = set((f.name for f in input_schema.fields))
    missing_attribute_names = set(extra_fields) - schema_names
    if missing_attribute_names:
        raise ValueError(f'Attribute fields {missing_attribute_names} not found in schema fields {schema_names}')
    payload_schema = schema_pb2.Schema(fields=[field for field in input_schema.fields if field.name not in extra_fields])
    formatter = _create_formatter(format, schema, payload_schema)
    return pcoll | beam.Map(lambda row: beam.io.gcp.pubsub.PubsubMessage(formatter(row), attributes_extractor(row))) | beam.io.WriteToPubSub(topic, with_attributes=True, id_label=id_attribute, timestamp_attribute=timestamp_attribute)

Writes messages to Cloud Pub/Sub.

Args:
  topic: Cloud Pub/Sub topic in the form "/topics/<project>/<topic>".
  format: How to format the message payload.  Currently suported
    formats are

      - RAW: Expects a message with a single field (excluding
          attribute-related fields) whose contents are used as the raw bytes
          of the pubsub message.
      - AVRO: Encodes records with a given Avro schema, which may be inferred
          from the input PCollection schema.
      - JSON: Formats records with a given JSON schema, which may be inferred
          from the input PCollection schema.

  schema: Schema specification for the given format.
  attributes: List of attribute keys whose values will be pulled out as
    PubSub message attributes.  For example, if the format is `raw`
    and attributes is `["a", "b"]` then elements of the form
    `Row(any_field=..., a=..., b=...)` will result in PubSub messages whose
    payload has the contents of any_field and whose attribute will be
    populated with the values of `a` and `b`.
  attribute_map: Name of a string-to-string map field in which to pull a set
    of attributes associated with this message.  For example, if the format
    is `raw` and `attribute_map` is set to `"attrs"` then elements of the form
    `Row(any_field=..., attrs=...)` will result in PubSub messages whose
    payload has the contents of any_field and whose attribute will be
    populated with the values from attrs.
    If both `attributes` and `attribute_map` are set, the union of attributes
    from these two sources will be used to populate the PubSub message
    attributes.
  id_attribute: If set, will set an attribute for each Cloud Pub/Sub message
    with the given name and a unique value. This attribute can then be used
    in a ReadFromPubSub PTransform to deduplicate messages.
  timestamp_attribute: If set, will set an attribute for each Cloud Pub/Sub
    message with the given name and the message's publish time as the value.


def add_pop_block_targets(bytecode: list[opcodes.Opcode]) -> None:
    if not bytecode:
        return
    for op in bytecode:
        op.block_target = None
    setup_except_op = (opcodes.SETUP_FINALLY, opcodes.SETUP_EXCEPT_311)
    todo = [(bytecode[0], ())]
    seen = set()
    while todo:
        op, block_stack = todo.pop()
        if op in seen:
            continue
        else:
            seen.add(op)
        if isinstance(op, opcodes.POP_BLOCK):
            assert block_stack, 'POP_BLOCK without block.'
            op.block_target = block_stack[-1].target
            block_stack = block_stack[0:-1]
        elif isinstance(op, opcodes.RAISE_VARARGS):
            for b in reversed(block_stack):
                if isinstance(b, setup_except_op):
                    op.block_target = b.target
                    break
        elif isinstance(op, opcodes.BREAK_LOOP):
            for i in reversed(range(len(block_stack))):
                b = block_stack[i]
                if isinstance(b, opcodes.SETUP_LOOP):
                    op.block_target = b.target
                    assert b.target != op
                    todo.append((op.block_target, block_stack[0:i]))
                    break
        elif isinstance(op, setup_except_op):
            todo.append((op.target, block_stack))
            block_stack += (op,)
        elif op.pushes_block():
            assert op.target, f'{op.name} without target'
            block_stack += (op,)
        elif op.does_jump() and op.target:
            if op.push_exc_block:
                setup_op = op.target
                while not isinstance(setup_op, setup_except_op):
                    setup_op = setup_op.prev
                block_stack += (setup_op,)
            todo.append((op.target, block_stack))
        if not op.no_next():
            assert op.next, f'Bad instruction at end of bytecode: {op!r}.'
            todo.append((op.next, block_stack))

Modifies bytecode so that each POP_BLOCK has a block_target.

This is to achieve better initial ordering of try/except and try/finally code.
try:
  i = 1
  a[i]
except IndexError:
  return i
By connecting a CFG edge from the end of the block (after the "a[i]") to the
except handler, our basic block ordering algorithm knows that the except block
needs to be scheduled last, whereas if there only was an edge before the
"i = 1", it would be able to schedule it too early and thus encounter an
undefined variable. This is only for ordering. The actual analysis of the
code happens later, in vm.py.

Args:
  bytecode: An array of bytecodes.


def reduce_sum(self, x):
    return self.reduce(lambda y: math_ops.reduce_sum(y, axis=0), x)

Performs a sum reduction on `x` across pfor iterations.

Note that this currently may not work inside a control flow construct.
Args:
  x: an unvectorized Tensor.

Returns:
  A Tensor that has same rank as `x`. The value is the sum of the values
  of `x` across the pfor iterations.


def maximum(inputs, **kwargs):
    return Maximum(**kwargs)(inputs)

Functional interface to compute maximum (element-wise) list of `inputs`.

This is equivalent to the `tf.keras.layers.Maximum` layer.

For example:

```python
input1 = tf.keras.layers.Input(shape=(16,))
x1 = tf.keras.layers.Dense(8, activation='relu')(input1) #shape=(None, 8)
input2 = tf.keras.layers.Input(shape=(32,))
x2 = tf.keras.layers.Dense(8, activation='relu')(input2) #shape=(None, 8)
max_inp=tf.keras.layers.maximum([x1,x2]) #shape=(None, 8)
out = tf.keras.layers.Dense(4)(max_inp)
model = tf.keras.models.Model(inputs=[input1, input2], outputs=out)
```

Args:
    inputs: A list of input tensors (at least 2) of same shape.
    **kwargs: Standard layer keyword arguments.

Returns:
    A tensor (of same shape as input tensor) with the element-wise
    maximum of the inputs.

Raises:
    ValueError: If input tensors are of different shape.


def _dict_to_tensor(self, x, k1, k2, k3):
    return array_ops_stack.stack([array_ops_stack.stack([array_ops_stack.stack([x[i, j, k] for k in range(k3)]) for j in range(k2)]) for i in range(k1)])

Convert a dictionary to a tensor.

Args:
  x: A k1 * k2 dictionary.
  k1: First dimension of x.
  k2: Second dimension of x.
  k3: Third dimension of x.

Returns:
  A k1 * k2 * k3 tensor.


def delete_bq_table(project, dataset_id, table_id):
    _LOGGER.info('Clean up a BigQuery table with project: %s, dataset: %s, table: %s.', project, dataset_id, table_id)
    client = bigquery.Client(project=project)
    table_ref = client.dataset(dataset_id).table(table_id)
    try:
        client.delete_table(table_ref)
    except gexc.NotFound:
        raise GcpTestIOError('BigQuery table does not exist: %s' % table_ref)

Delete a BiqQuery table.

Args:
  project: Name of the project.
  dataset_id: Name of the dataset where table is.
  table_id: Name of the table.


class LlavaOnevisionConfig(PretrainedConfig):
    model_type = 'llava_onevision'
    attribute_map = {'image_token_id': 'image_token_index', 'video_token_id': 'video_token_index'}
    sub_configs = {'text_config': AutoConfig, 'vision_config': AutoConfig}

    def __init__(self, vision_config=None, text_config=None, image_token_index=151646, video_token_index=151647, projector_hidden_act='gelu', vision_feature_select_strategy='full', vision_feature_layer=-1, vision_aspect_ratio='anyres_max_9', image_grid_pinpoints=None, tie_word_embeddings=False, multimodal_projector_bias=True, **kwargs):
        self.image_token_index = image_token_index
        self.video_token_index = video_token_index
        self.projector_hidden_act = projector_hidden_act
        self.multimodal_projector_bias = multimodal_projector_bias
        if vision_feature_select_strategy not in ['default', 'full']:
            raise ValueError(f"vision_feature_select_strategy should be one of 'default', 'full'.Got: {vision_feature_select_strategy}")
        self.vision_feature_select_strategy = vision_feature_select_strategy
        self.vision_feature_layer = vision_feature_layer
        self.vision_aspect_ratio = vision_aspect_ratio
        image_grid_pinpoints = image_grid_pinpoints if image_grid_pinpoints is not None else [[384, 384], [384, 768], [384, 1152], [384, 1536], [384, 1920], [384, 2304], [768, 384], [768, 768], [768, 1152], [768, 1536], [768, 1920], [768, 2304], [1152, 384], [1152, 768], [1152, 1152], [1152, 1536], [1152, 1920], [1152, 2304], [1536, 384], [1536, 768], [1536, 1152], [1536, 1536], [1536, 1920], [1536, 2304], [1920, 384], [1920, 768], [1920, 1152], [1920, 1536], [1920, 1920], [1920, 2304], [2304, 384], [2304, 768], [2304, 1152], [2304, 1536], [2304, 1920], [2304, 2304]]
        self.image_grid_pinpoints = image_grid_pinpoints
        if isinstance(vision_config, dict):
            vision_config['model_type'] = vision_config['model_type'] if 'model_type' in vision_config else 'siglip_vision_model'
            vision_config = CONFIG_MAPPING[vision_config['model_type']](**vision_config)
        elif vision_config is None:
            vision_config = CONFIG_MAPPING['siglip_vision_model'](hidden_size=1152, intermediate_size=4304, patch_size=14, image_size=384, num_hidden_layers=26, num_attention_heads=14, vision_use_head=False)
        self.vision_config = vision_config
        if isinstance(text_config, dict):
            text_config['model_type'] = text_config['model_type'] if 'model_type' in text_config else 'qwen2'
            text_config = CONFIG_MAPPING[text_config['model_type']](**text_config)
        elif text_config is None:
            text_config = CONFIG_MAPPING['qwen2']()
        self.text_config = text_config
        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)

This is the configuration class to store the configuration of a [`LlavaOnevisionForConditionalGeneration`]. It is used to instantiate an
Llava-NeXT 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 [llava-hf/llava-onevision-qwen2-7b-ov-hf](https://huggingface.co/llava-hf/llava-onevision-qwen2-7b-ov-hf)
model.

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

Args:
    vision_config (`Union[AutoConfig, dict]`,  *optional*, defaults to `SiglipVisionConfig`):
        The config object or dictionary of the vision backbone.
    text_config (`Union[AutoConfig, dict]`, *optional*, defaults to `Qwen2Config`):
        The config object or dictionary of the text backbone.
    image_token_index (`int`, *optional*, defaults to 151646):
        The image token index to encode the image prompt.
    video_token_index (`int`, *optional*, defaults to 151647):
        The video token index to encode the video prompt.
    projector_hidden_act (`str`, *optional*, defaults to `"gelu"`):
        The activation function used by the multimodal projector.
    vision_feature_select_strategy (`str`, *optional*, defaults to `"full"`):
        The feature selection strategy used to select the vision feature from the vision backbone.
        Can be one of `"default"` or `"full"`. If `"default"`, the CLS token is removed from the vision features.
        If `"full"`, the full vision features are used.
    vision_feature_layer (`Union[int, List[int]]`, *optional*, defaults to -1):
        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.
    vision_aspect_ratio (`str`, *optional*, defaults to `"anyres_max_9"`):
        Aspect ratio used when processong image features. The default value is "anyres_max_9".
    image_grid_pinpoints (`List`, *optional*):
        A list of possible resolutions to use for processing high resolution images. Each item in the list should be a tuple or list
        of the form `(height, width)`.
    tie_word_embeddings (`bool`, *optional*, defaults to `False`):
        Whether the model's input and output word embeddings should be tied.
    multimodal_projector_bias (`bool`, *optional*, defaults to `True`):
        Whether to use bias in the multimodal projector.

Example:

```python
>>> from transformers import LlavaOnevisionForConditionalGeneration, LlavaOnevisionConfig, SiglipVisionConfig, Qwen2Config

>>> # Initializing a CLIP-vision config
>>> vision_config = SiglipVisionConfig()

>>> # Initializing a Llama config
>>> text_config = Qwen2Config()

>>> # Initializing a Llava-Next llava-hf/llava-onevision-qwen2-7b-ov-hf style configuration
>>> configuration = LlavaOnevisionConfig(vision_config, text_config)

>>> # Initializing a model from the llava-hf/llava-onevision-qwen2-7b-ov-hf style configuration
>>> model = LlavaOnevisionForConditionalGeneration(configuration)

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


def __call__(self, request: beam.Row, *args, **kwargs):
    try:
        entity_id = request._asdict()[self.row_key]
    except KeyError:
        raise KeyError('Enrichment requests to Vertex AI Feature Store should contain a field: %s in the input `beam.Row` to join the input with fetched response. This is used as the `FeatureViewDataKey` to fetch feature values corresponding to this key.' % self.row_key)
    try:
        response = self.client.fetch_feature_values(request=aiplatform.gapic.FetchFeatureValuesRequest(data_key=aiplatform.gapic.FeatureViewDataKey(key=entity_id), feature_view=self.feature_view_path, data_format=aiplatform.gapic.FeatureViewDataFormat.PROTO_STRUCT))
    except NotFound:
        if self.exception_level == ExceptionLevel.WARN:
            _LOGGER.warning(_not_found_err_message(self.feature_store_name, self.feature_view_name, entity_id))
            return (request, beam.Row())
        elif self.exception_level == ExceptionLevel.RAISE:
            raise ValueError(_not_found_err_message(self.feature_store_name, self.feature_view_name, entity_id))
    response_dict = dict(response.proto_struct)
    return (request, beam.Row(**response_dict))

Fetches feature value for an entity-id from Vertex AI Feature Store.

Args:
  request: the input `beam.Row` to enrich.


def __init__(self, cell):
    self._cell = cell

Creates a new SamplerCell.

Args:
  cell: A c pointer of TFE_MonitoringSamplerCell.


def has_deprecation_decorator(symbol):
    decorators, symbol = tf_decorator.unwrap(symbol)
    if contains_deprecation_decorator(decorators):
        return True
    if tf_inspect.isfunction(symbol):
        return False
    if not tf_inspect.isclass(symbol):
        return False
    if not hasattr(symbol, '__init__'):
        return False
    init_decorators, _ = tf_decorator.unwrap(symbol.__init__)
    return contains_deprecation_decorator(init_decorators)

Checks if given object has a deprecation decorator.

We check if deprecation decorator is in decorators as well as
whether symbol is a class whose __init__ method has a deprecation
decorator.
Args:
  symbol: Python object.

Returns:
  True if symbol has deprecation decorator.


class SimpleSlidingQuantileTracker(WindowedTracker, QuantileTracker):

    def __init__(self, window_size, q):
        super().__init__(window_mode=WindowMode.SLIDING, window_size=window_size)
        QuantileTracker.__init__(self, q)

    def get(self):
        """Calculates and returns the specified quantile of the current sliding
    window.

    Returns:
      float: The specified quantile of the values in the current sliding window.
             Returns NaN if the window is empty.
    """
        with warnings.catch_warnings(record=False):
            warnings.simplefilter('ignore')
            return np.nanquantile(self._queue, self._q)

Sliding window quantile tracker using NumPy.

This tracker uses NumPy's `nanquantile` function to calculate the specified
quantile of the values currently in the sliding window. It's a simple,
non-incremental approach.

Args:
  window_size: The size of the sliding window.
  q: The quantile to calculate, a float between 0 and 1 (inclusive).


def compute_gradients(self, loss, var_list=None, gate_gradients=GATE_OP, aggregation_method=None, colocate_gradients_with_ops=False, grad_loss=None):
    if callable(loss):
        with backprop.GradientTape() as tape:
            if var_list is not None:
                tape.watch(var_list)
            loss_value = loss()
            loss_value = self._scale_loss(loss_value)
        if var_list is None:
            var_list = tape.watched_variables()
        with ops.control_dependencies([loss_value]):
            grads = tape.gradient(loss_value, var_list, grad_loss)
        return list(zip(grads, var_list))
    if context.executing_eagerly():
        raise RuntimeError('`loss` passed to Optimizer.compute_gradients should be a function when eager execution is enabled.')
    loss = self._scale_loss(loss)
    if gate_gradients not in [Optimizer.GATE_NONE, Optimizer.GATE_OP, Optimizer.GATE_GRAPH]:
        raise ValueError('gate_gradients must be one of: Optimizer.GATE_NONE, Optimizer.GATE_OP, Optimizer.GATE_GRAPH.  Not %s' % gate_gradients)
    self._assert_valid_dtypes([loss])
    if grad_loss is not None:
        self._assert_valid_dtypes([grad_loss])
    if var_list is None:
        var_list = variables.trainable_variables() + ops.get_collection(ops.GraphKeys.TRAINABLE_RESOURCE_VARIABLES)
    else:
        var_list = nest.flatten(var_list)
    var_list += ops.get_collection(ops.GraphKeys._STREAMING_MODEL_PORTS)
    processors = [_get_processor(v) for v in var_list]
    if not var_list:
        raise ValueError('No variables to optimize.')
    var_refs = [p.target() for p in processors]
    grads = gradients.gradients(loss, var_refs, grad_ys=grad_loss, gate_gradients=gate_gradients == Optimizer.GATE_OP, aggregation_method=aggregation_method, colocate_gradients_with_ops=colocate_gradients_with_ops)
    if gate_gradients == Optimizer.GATE_GRAPH:
        grads = control_flow_ops.tuple(grads)
    grads_and_vars = list(zip(grads, var_list))
    self._assert_valid_dtypes([v for g, v in grads_and_vars if g is not None and v.dtype != dtypes.resource])
    return grads_and_vars

Compute gradients of `loss` for the variables in `var_list`.

This is the first part of `minimize()`.  It returns a list
of (gradient, variable) pairs where "gradient" is the gradient
for "variable".  Note that "gradient" can be a `Tensor`, an
`IndexedSlices`, or `None` if there is no gradient for the
given variable.

@compatibility(TF2)
`tf.keras.optimizers.Optimizer` in TF2 does not provide a
`compute_gradients` method, and you should use a `tf.GradientTape` to
obtain the gradients:

```python
@tf.function
def train step(inputs):
  batch_data, labels = inputs
  with tf.GradientTape() as tape:
    predictions = model(batch_data, training=True)
    loss = tf.keras.losses.CategoricalCrossentropy(
        reduction=tf.keras.losses.Reduction.NONE)(labels, predictions)
  gradients = tape.gradient(loss, model.trainable_variables)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))
```

Args:
  loss: A Tensor containing the value to minimize or a callable taking
    no arguments which returns the value to minimize. When eager execution
    is enabled it must be a callable.
  var_list: Optional list or tuple of `tf.Variable` to update to minimize
    `loss`.  Defaults to the list of variables collected in the graph
    under the key `GraphKeys.TRAINABLE_VARIABLES`.
  gate_gradients: How to gate the computation of gradients.  Can be
    `GATE_NONE`, `GATE_OP`, or `GATE_GRAPH`.
  aggregation_method: Specifies the method used to combine gradient terms.
    Valid values are defined in the class `AggregationMethod`.
  colocate_gradients_with_ops: If True, try colocating gradients with
    the corresponding op.
  grad_loss: Optional. A `Tensor` holding the gradient computed for `loss`.

Returns:
  A list of (gradient, variable) pairs. Variable is always present, but
  gradient can be `None`.

Raises:
  TypeError: If `var_list` contains anything else than `Variable` objects.
  ValueError: If some arguments are invalid.
  RuntimeError: If called with eager execution enabled and `loss` is
    not callable.

@compatibility(eager)
When eager execution is enabled, `gate_gradients`, `aggregation_method`,
and `colocate_gradients_with_ops` are ignored.
@end_compatibility


def _sanitize_windows_filename(filename):
    if re.match(WINDOWS_RESERVED_FILENAME_REGEX, filename):
        return WINDOWS_RESERVED_FILENAME_PREFIX + filename
    filename = _truncate_filename(filename, WINDOWS_MAX_FILENAME_LENGTH)
    new_filename_chars = []
    for char in filename:
        if char in WINDOWS_RESERVED_CHARACTERS_REPLACEMENTS:
            new_filename_chars.append(WINDOWS_RESERVED_CHARACTERS_REPLACEMENTS[char])
        else:
            new_filename_chars.append(char)
    filename = ''.join(new_filename_chars)
    if filename.endswith('.') or filename.endswith(' '):
        filename = filename[:-1] + '_'
    return filename

Sanitizes a filename for Windows.

Refer to the following Windows documentation page for the rules:
https://docs.microsoft.com/en-us/windows/win32/fileio/naming-a-file#naming-conventions

If the filename matches one of Window's reserved file namespaces, then the
`WINDOWS_RESERVED_FILENAME_PREFIX` (i.e. "mobly_") prefix will be appended
to the filename to convert it into a valid Windows filename.

Args:
  filename: string, the filename to sanitize for the Windows file system.

Returns:
  A filename that should be safe to use on Windows.


class RaggedTensorDynamicShape:

    def __init__(self, partitioned_dim_sizes, inner_dim_sizes, dim_size_dtype=None):
        """Creates a RaggedTensorDynamicShape.

    Args:
      partitioned_dim_sizes: A `list` of 0-D or 1-D integer `Tensor`, one for
        each partitioned dimension.  If dimension `d` is uniform, then
        `partitioned_dim_sizes[d]` must be an integer scalar, specifying the
        size of all slices across dimension `d`.  If dimension `d` is ragged,
        then `partitioned_dim_sizes[d]` must be an integer vector, specifying
        the size of each slice across dimension `d`.
      inner_dim_sizes: A 1-D integer `Tensor`, whose length is equal to the
        number of inner dimensions.  `inner_dim_sizes[n]` is the size of all
        slices across the `n`th inner dimension (which is the
        `(len(partitioned_dim_sizes)+n)`th dimension in the overall tensor.
      dim_size_dtype: dtype for dimension sizes.  If not specified, then it
        is chosen based on the dtypes of `partitioned_dim_sizes` and
        `inner_dim_sizes`.
    """
        assert isinstance(partitioned_dim_sizes, (list, tuple))
        with ops.name_scope(None, 'RaggedTensorDynamicShape', (partitioned_dim_sizes, inner_dim_sizes)):
            partitioned_dim_sizes = tuple((ops.convert_to_tensor(size, name='partitioned_dimension_size_%d' % i) for i, size in enumerate(partitioned_dim_sizes)))
            inner_dim_sizes = ops.convert_to_tensor(inner_dim_sizes, name='inner_dim_sizes')
            if partitioned_dim_sizes:
                for axis, dimension_size in enumerate(partitioned_dim_sizes):
                    if dimension_size.shape.ndims is None:
                        raise ValueError('rank of partitioned_dim_sizes[%d] is unknown' % axis)
                    dimension_size.shape.with_rank_at_most(1)
                if partitioned_dim_sizes[0].shape.ndims == 1:
                    raise ValueError('outermost partitioned dimension must be uniform')
                if partitioned_dim_sizes[-1].shape.ndims == 0:
                    raise ValueError('innermost partitioned dimension must be ragged')
            inner_dim_sizes.shape.assert_has_rank(1)
            if dim_size_dtype is None:
                dim_size_dtypes = set((p.dtype for p in partitioned_dim_sizes if p.shape.ndims == 1))
                if not dim_size_dtypes:
                    dim_size_dtype = dtypes.int64
                elif len(dim_size_dtypes) == 1:
                    dim_size_dtype = dim_size_dtypes.pop()
                else:
                    if not ragged_config.auto_cast_partition_dtype():
                        raise ValueError('partitioned_dim_sizes must have matching dtypes')
                    dim_size_dtype = dtypes.int64
            partitioned_dim_sizes = tuple((math_ops.cast(p, dim_size_dtype) for p in partitioned_dim_sizes))
            inner_dim_sizes = math_ops.cast(inner_dim_sizes, dim_size_dtype)
            self._partitioned_dim_sizes = partitioned_dim_sizes
            self._inner_dim_sizes = inner_dim_sizes

    def __repr__(self):
        return 'RaggedTensorDynamicShape(partitioned_dim_sizes=%r, inner_dim_sizes=%r)' % (self._partitioned_dim_sizes, self._inner_dim_sizes)

    @staticmethod
    def from_dim_sizes(dim_sizes):
        """Constructs a ragged shape from a list of dimension sizes.

    This list contains a single tensor for each dimension, where the tensor
    is a scalar if the dimension is uniform, or a vector if the dimension is
    ragged.

    Args:
      dim_sizes: List of int32 or int64 scalars or vectors.

    Returns:
      A RaggedTensorDynamicShape.
    """
        with ops.name_scope(None, 'RaggedTensorDynamicShapeFromDimensionSizes', [dim_sizes]):
            dim_sizes = tuple((ops.convert_to_tensor(size, preferred_dtype=dtypes.int64, name='dim_sizes') for size in dim_sizes))
            inner_split = 0
            for dim, dim_size in enumerate(dim_sizes):
                if dim_size.shape.ndims == 1:
                    inner_split = dim + 1
                elif dim_size.shape.ndims != 0:
                    raise ValueError('Each dim_size must be a scalar or a vector')
            return RaggedTensorDynamicShape(dim_sizes[:inner_split], dim_sizes[inner_split:])

    @classmethod
    def from_tensor(cls, rt_input, dim_size_dtype=None):
        """Constructs a ragged shape for a potentially ragged tensor."""
        with ops.name_scope(None, 'RaggedTensorDynamicShapeFromTensor', [rt_input]):
            rt_input = ragged_tensor.convert_to_tensor_or_ragged_tensor(rt_input)
            if not ragged_tensor.is_ragged(rt_input):
                return cls([], array_ops.shape(rt_input), dim_size_dtype=dim_size_dtype)
            else:
                partitioned_dim_sizes = (rt_input.nrows(),) + rt_input.nested_row_lengths()
                return RaggedTensorDynamicShape(partitioned_dim_sizes, array_ops.shape(rt_input.flat_values)[1:], dim_size_dtype=dim_size_dtype)

    def dimension_size(self, axis):
        """Returns the size of slices across the specified dimension."""
        if not isinstance(axis, int):
            raise TypeError('axis must be an integer')
        partitioned_ndims = len(self._partitioned_dim_sizes)
        if axis < partitioned_ndims:
            return self._partitioned_dim_sizes[axis]
        else:
            return self._inner_dim_sizes[axis - partitioned_ndims]

    def is_ragged(self, axis):
        """Returns true if the indicated dimension is ragged."""
        if not isinstance(axis, int):
            raise TypeError('axis must be an integer')
        rank = self.rank
        if axis < 0:
            raise ValueError('Negative axis values are not supported')
        elif rank is not None and axis >= rank:
            raise ValueError('Expected axis=%s < rank=%s' % (axis, rank))
        else:
            return axis > 0 and axis < len(self._partitioned_dim_sizes) and (self._partitioned_dim_sizes[axis].shape.ndims == 1)

    @property
    def rank(self):
        """The number of dimensions in this shape, or None if unknown."""
        inner_ndims = tensor_shape.dimension_value(self._inner_dim_sizes.shape[0])
        if inner_ndims is None:
            return None
        else:
            return len(self._partitioned_dim_sizes) + inner_ndims

    @property
    def partitioned_dim_sizes(self):
        """The partitioned dimension sizes for this shape.

    Returns:
      A `list` of 0-D or 1-D integer `Tensor`.
    """
        return self._partitioned_dim_sizes

    @property
    def inner_dim_sizes(self):
        """The inner dimension sizes for this shape.

    Returns:
      A 1-D integer `Tensor`.
    """
        return self._inner_dim_sizes

    @property
    def num_partitioned_dimensions(self):
        """The number of partitioned dimensions in this shape."""
        return len(self._partitioned_dim_sizes)

    @property
    def num_inner_dimensions(self):
        """The number of inner dimensions, or `None` if not statically known."""
        return tensor_shape.dimension_value(self._inner_dim_sizes.shape[0])

    @property
    def dim_size_dtype(self):
        """DType used by this shape for dimension sizes."""
        return self._inner_dim_sizes.dtype

    def broadcast_to_rank(self, rank):
        """Adds leading size-1 dimensions to broadcast `self` to the given rank.

    E.g., if `shape1` is `[3, (D2), 4]`, then `shape1.broadcast_to_rank(5)`
    is `[1, 1, 3, (D2), 4]`.

    Args:
      rank: The rank for the returned shape.

    Returns:
      A RaggedTensorDynamicShape with `rank` dimensions, whose inner dimensions
      have the same size as `self` and whose outer dimensions have size `1`.

    Raises:
      ValueError: If `self.rank` is unknown or greater than `rank`.
    """
        if self.rank is None:
            raise ValueError('Unable to broadcast: self.rank is unknown')
        dims_to_add = rank - self.rank
        if dims_to_add < 0:
            raise ValueError('Unable to broadcast: rank=%d must be greater than self.rank=%d.' % (rank, self.rank))
        elif dims_to_add == 0:
            return self
        elif self._partitioned_dim_sizes:
            partitioned_dims = (1,) * dims_to_add + self._partitioned_dim_sizes
            return RaggedTensorDynamicShape(partitioned_dims, self.inner_dim_sizes, self.dim_size_dtype)
        else:
            inner_dims = array_ops.concat([array_ops.ones([dims_to_add], self.dim_size_dtype), self.inner_dim_sizes], axis=0)
            return RaggedTensorDynamicShape([], inner_dims, self.dim_size_dtype)

    def broadcast_dimension(self, axis, lengths):
        """Returns a shape that is broadcast-compatible with self & lengths.

    * If dimension[axis] is uniform and lengths is a scalar, the check
      that either lengths==1 or axis==1 or lengths==axis, and tile
      dimension[axis] with tf.where(lengths==axis, 1, axis) repeats.

    * If dimension[axis] is uniform and lengths is a vector, then check
      that dimension[axis]==1, and raggedly tile dimension[axis] with
      lengths repeats.  (we can skip tiling if we statically know that
      slice_lengths == 1??)

    * If dimension[axis] is ragged and lengths is a scalar, then check
      that lengths==1.

    * If dimension[axis] is ragged and lengths is a vector, then check
      that self.dimension_size(axis) == lengths.

    Args:
      axis: `int`.  The dimension to broadcast.
      lengths: 0-D or 1-D integer `Tensor`.

    Returns:
      A `RaggedTensorDynamicShape`.
    """
        lengths = ragged_util.convert_to_int_tensor(lengths, name='lengths', dtype=self.dim_size_dtype)
        if lengths.shape.ndims is None:
            raise ValueError('lengths must have a known rank.')
        elif lengths.shape.ndims > 1:
            raise ValueError('lengths must be a scalar or vector')
        else:
            lengths_is_scalar = lengths.shape.ndims == 0
        if self.is_ragged(axis):
            if lengths_is_scalar:
                condition = math_ops.equal(lengths, 1)
            else:
                condition = math_ops.reduce_all(math_ops.equal(lengths, self.dimension_size(axis)))
        else:
            axis_dim_size = self.dimension_size(axis)
            if lengths_is_scalar:
                condition = math_ops.equal(lengths, 1) | math_ops.equal(axis_dim_size, 1) | math_ops.equal(axis_dim_size, lengths)
            else:
                condition = math_ops.equal(axis_dim_size, 1)
        broadcast_err = ['Unable to broadcast: dimension size mismatch in dimension', axis, 'lengths=', lengths, 'dim_size=', self.dimension_size(axis)]
        broadcast_check = control_flow_assert.Assert(condition, data=broadcast_err, summarize=10)
        with ops.control_dependencies([broadcast_check]):
            if axis < self.num_partitioned_dimensions:
                if self.is_ragged(axis):
                    return RaggedTensorDynamicShape(self._partitioned_dim_sizes, array_ops.identity(self.inner_dim_sizes), self.dim_size_dtype)
                else:
                    return self._broadcast_uniform_partitioned_dimension(axis, lengths)
            elif lengths_is_scalar:
                return self._broadcast_inner_dimension_to_uniform(axis, lengths)
            else:
                if axis == 0:
                    raise ValueError('Unable to broadcast: outermost dimension must be uniform.')
                return self._broadcast_inner_dimension_to_ragged(axis, lengths)

    def num_slices_in_dimension(self, axis):
        """Returns the total number of slices across the indicated dimension."""
        if axis < 0:
            return constant_op.constant(1, dtype=self.dim_size_dtype)
        elif self.is_ragged(axis):
            return math_ops.reduce_sum(self._partitioned_dim_sizes[axis])
        else:
            return self.dimension_size(axis) * self.num_slices_in_dimension(axis - 1)

    def _broadcast_uniform_partitioned_dimension(self, axis, lengths):
        """Broadcasts the partitioned dimension `axis` to match `lengths`."""
        axis_dim_size = self.dimension_size(axis)
        partitioned_sizes = list(self._partitioned_dim_sizes[:axis])
        if lengths.shape.ndims == 0:
            lengths = array_ops.where(math_ops.equal(axis_dim_size, 1), lengths, axis_dim_size)
            repeats = array_ops.where(math_ops.equal(axis_dim_size, 1), lengths, 1)
            splits = array_ops_stack.stack([0, self.num_slices_in_dimension(axis)])
        else:
            splits = math_ops.range(array_ops.size(lengths, out_type=self.dim_size_dtype) + 1)
            repeats = lengths
        partitioned_sizes.append(lengths)
        for dim_size in self._partitioned_dim_sizes[axis + 1:]:
            if dim_size.shape.ndims == 0:
                partitioned_sizes.append(dim_size)
                splits *= dim_size
            else:
                partitioned_sizes.append(ragged_util.repeat_ranges(dim_size, splits, repeats))
                splits = array_ops.gather(ragged_util.lengths_to_splits(dim_size), splits)
        inner_sizes = self._inner_dim_sizes
        return RaggedTensorDynamicShape(partitioned_sizes, inner_sizes, self.dim_size_dtype)

    def _broadcast_inner_dimension_to_uniform(self, axis, length):
        """Broadcasts the inner dimension `axis` to match `lengths`."""
        dim_size = self.dimension_size(axis)
        axis_in_inner_dims = axis - self.num_partitioned_dimensions
        partitioned_sizes = self._partitioned_dim_sizes
        inner_sizes = array_ops.concat([self._inner_dim_sizes[:axis_in_inner_dims], [array_ops.where(math_ops.equal(dim_size, 1), length, dim_size)], self._inner_dim_sizes[axis_in_inner_dims + 1:]], axis=0)
        return RaggedTensorDynamicShape(partitioned_sizes, inner_sizes, self.dim_size_dtype)

    def _broadcast_inner_dimension_to_ragged(self, axis, lengths):
        axis_in_inner_dims = axis - self.num_partitioned_dimensions
        partitioned_sizes = self._partitioned_dim_sizes + tuple([self._inner_dim_sizes[i] for i in range(axis_in_inner_dims)]) + (lengths,)
        inner_sizes = self._inner_dim_sizes[axis_in_inner_dims + 1:]
        return RaggedTensorDynamicShape(partitioned_sizes, inner_sizes)

    def with_dim_size_dtype(self, dtype):
        if dtype not in (dtypes.int32, dtypes.int64):
            raise ValueError('dtype must be int32 or int64')
        if self.dim_size_dtype == dtype:
            return self
        return RaggedTensorDynamicShape([math_ops.cast(p, dtype) for p in self._partitioned_dim_sizes], math_ops.cast(self._inner_dim_sizes, dtype))

A collection of tensors encoding the shape of a potentially ragged tensor.

Each `RaggedTensorDynamicShape` consists of an ordered list of dimension
sizes.  There are two dimension types:

  * "Uniform dimensions" are dimensions where all slices have the same
    length.  `RaggedTensorDynamicShape` records the size of each uniform
    dimension using a single scalar integer.

  * "Ragged dimensions" are dimensions whose slices may have different
    lengths.  `RaggedTensorDynamicShape` records the size of each ragged
    dimension using an integer vector containing the slice lengths for all
    the slices across that dimension.

Furthermore, there are two ways a dimension might be encoded:

  * "Partitioned dimensions" are dimensions that are encoded using a
    `RaggedTensor`'s `nested_row_splits`.  The outermostmost partitioned
    dimension must be uniform, and the innermost partitioned dimension must
    be ragged.

  * "Inner dimensions" are dimensions that are encoded using a
    `RaggedTensor`'s `flat_values`.  Inner dimensions are always uniform.

The sizes of partitioned dimensions are recorded using `partitioned_dim_sizes`
and `inner_dim_sizes`:

  * `partitioned_dim_sizes` is a list of tensors (one for each partitioned
    dimension).

    * For uniform dimensions, the tensor is an integer scalar specifying the
      size of all slices across that dimension.
    * For ragged dimensions, the tensor is an integer vector specifying the
      size of each slice across that dimension.

  * `inner_dim_sizes` is a single integer vector, where each element
    specifies the size of a single inner dimension.

Examples:

Tensor                         | Ragged | Partitioned Dim Sizes  | Inner Dim
                               : Rank   :                        : Sizes
------------------------------ | ------ | ---------------------- | ----------
`[[1, 2, 3], [4, 5, 6]]`       |      0 |                        | `2, 3`
`[[1, 2], [], [3, 4, 5]]`      |      1 | `3, (2, 0, 3)`         |
`[[[1, 2], [3, 4]], [[5, 6]]]` |      1 | `2, (2, 1)`            | 2
`[[[1, 2], [3]], [[4, 5]]]`    |      2 | `2, (2, 1), (2, 1, 2)` |


def increment_max_models(self, increment: int):
    if self._max_models is None:
        self._max_models = 0
    self._max_models += increment

Increments the number of models that this instance of a _ModelManager is
able to hold. If it is never called, no limit is imposed.
Args:
  increment: the amount by which we are incrementing the number of models.


def get_task_states(self, job_configs):
    if self._context_handle:
        job_names, task_nums = zip(*job_configs)
        return pywrap_tfe.TFE_GetTaskStates(self._context_handle, job_names, task_nums)
    else:
        raise ValueError('Context is not initialized.')

Get task states from the Coordination Service.

Args:
  job_configs: A list of tuples of job name and task number.

Returns:
  A list of TF_Status.


class FlavaModelOutput(ModelOutput):
    image_embeddings: Optional[torch.FloatTensor] = None
    image_output: Optional[BaseModelOutputWithPooling] = None
    text_embeddings: Optional[torch.FloatTensor] = None
    text_output: Optional[BaseModelOutputWithPooling] = None
    multimodal_embeddings: Optional[torch.FloatTensor] = None
    multimodal_output: Optional[BaseModelOutputWithPooling] = None

    def to_tuple(self) -> Tuple[Any]:
        return tuple((self[k] if k not in ['text_output', 'image_output', 'multimodal_output'] else getattr(self, k).to_tuple() for k in self.keys()))

Output from FlavaModel containing embeddings and outputs from individual encoders.

Note that `image_embeddings` and `text_embeddigns` returned are similar to pooled output returned from a
transformer. If you want embeddings for contrastive loss or retrieval use a FLAVA model's `image_projection` and
`text_projection` layers on `image_embeddings` and `text_embeddings` respectively.

Args:
    image_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `pixel_values` are present):
        The image embeddings which are basically the pooled output of [`FlavaImageModel`].
    image_output (`BaseModelOutputWithPooling`, *optional*, returned when `pixel_values` are present):
        The output of the [`FlavaImageModel`].
    text_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `input_ids` are present):
        The text embeddings which are basically the pooled output of [`FlavaTextModel`].
    text_output (`BaseModelOutputWithPooling`, *optional*, returned when `input_ids` are present):
        The output of the [`FlavaTextModel`].
    multimodal_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `input_ids` and `pixel_values` are present and `skip_multimodal_encoder` is `None` or `False`):
        The multimodal embeddings which are basically the pooled output of [`FlavaTextModel`].
    multimodal_output (`BaseModelOutputWithPooling`, returned when `input_ids` and `pixel_values` are present and `skip_multimodal_encoder` is `None` or `False`):
        The output of the [`FlavaMultimodalModel`].


def __init__(self, initial_value=None, trainable=None, collections=None, validate_shape=True, caching_device=None, name=None, variable_def=None, dtype=None, expected_shape=None, import_scope=None, constraint=None, use_resource=None, synchronization=variables.VariableSynchronization.AUTO, aggregation=variables.VariableAggregation.NONE, shape=None):

Creates a new variable with value `initial_value`.

The new variable is added to the graph collections listed in `collections`,
which defaults to `[GraphKeys.GLOBAL_VARIABLES]`.

If `trainable` is `True` the variable is also added to the graph collection
`GraphKeys.TRAINABLE_VARIABLES`.

This constructor creates both a `variable` Op and an `assign` Op to set the
variable to its initial value.

Args:
  initial_value: A `Tensor`, or Python object convertible to a `Tensor`,
    which is the initial value for the Variable. The initial value must have
    a shape specified unless `validate_shape` is set to False. Can also be a
    callable with no argument that returns the initial value when called. In
    that case, `dtype` must be specified. (Note that initializer functions
    from init_ops.py must first be bound to a shape before being used here.)
  trainable: If `True`, also adds the variable to the graph collection
    `GraphKeys.TRAINABLE_VARIABLES`. This collection is used as the default
    list of variables to use by the `Optimizer` classes. Defaults to `True`,
    unless `synchronization` is set to `ON_READ`, in which case it defaults
    to `False`.
  collections: List of graph collections keys. The new variable is added to
    these collections. Defaults to `[GraphKeys.GLOBAL_VARIABLES]`.
  validate_shape: If `False`, allows the variable to be initialized with a
    value of unknown shape. If `True`, the default, the shape of
    `initial_value` must be known.
  caching_device: Optional device string describing where the Variable
    should be cached for reading.  Defaults to the Variable's device. If not
    `None`, caches on another device.  Typical use is to cache on the device
    where the Ops using the Variable reside, to deduplicate copying through
    `Switch` and other conditional statements.
  name: Optional name for the variable. Defaults to `'Variable'` and gets
    uniquified automatically.
  variable_def: `VariableDef` protocol buffer. If not `None`, recreates the
    Variable object with its contents, referencing the variable's nodes in
    the graph, which must already exist. The graph is not changed.
    `variable_def` and the other arguments are mutually exclusive.
  dtype: If set, initial_value will be converted to the given type. If
    `None`, either the datatype will be kept (if `initial_value` is a
    Tensor), or `convert_to_tensor` will decide.
  expected_shape: A TensorShape. If set, initial_value is expected to have
    this shape.
  import_scope: Optional `string`. Name scope to add to the `Variable.` Only
    used when initializing from protocol buffer.
  constraint: An optional projection function to be applied to the variable
    after being updated by an `Optimizer` (e.g. used to implement norm
    constraints or value constraints for layer weights). The function must
    take as input the unprojected Tensor representing the value of the
    variable and return the Tensor for the projected value (which must have
    the same shape). Constraints are not safe to use when doing asynchronous
    distributed training.
  use_resource: whether to use resource variables.
  synchronization: Indicates when a distributed a variable will be
    aggregated. Accepted values are constants defined in the class
    `tf.VariableSynchronization`. By default the synchronization is set to
    `AUTO` and the current `DistributionStrategy` chooses when to
    synchronize.
  aggregation: Indicates how a distributed variable will be aggregated.
    Accepted values are constants defined in the class
    `tf.VariableAggregation`.
  shape: (optional) The shape of this variable. If None, the shape of
    `initial_value` will be used. When setting this argument to
    `tf.TensorShape(None)` (representing an unspecified shape), the variable
    can be assigned with values of different shapes.

Raises:
  ValueError: If both `variable_def` and initial_value are specified.
  ValueError: If the initial value is not specified, or does not have a
    shape and `validate_shape` is `True`.
  RuntimeError: If eager execution is enabled.


def _retry_on_appropriate_openai_error(exception):
    return isinstance(exception, (RateLimitError, APIError))

Retry filter that returns True for rate limit (429) or server (5xx) errors.

Args:
  exception: the returned exception encountered during the request/response
    loop.

Returns:
  boolean indication whether or not the exception is a Server Error (5xx) or
    a RateLimitError (429) error.


class MaxPooling1D(keras_layers.MaxPooling1D, base.Layer):

    def __init__(self, pool_size, strides, padding='valid', data_format='channels_last', name=None, **kwargs):
        if strides is None:
            raise ValueError('Argument `strides` must not be None.')
        super(MaxPooling1D, self).__init__(pool_size=pool_size, strides=strides, padding=padding, data_format=data_format, name=name, **kwargs)

Max Pooling layer for 1D inputs.

Args:
  pool_size: An integer or tuple/list of a single integer,
    representing the size of the pooling window.
  strides: An integer or tuple/list of a single integer, specifying the
    strides of the pooling operation.
  padding: A string. The padding method, either 'valid' or 'same'.
    Case-insensitive.
  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, length, channels)` while `channels_first` corresponds to
    inputs with shape `(batch, channels, length)`.
  name: A string, the name of the layer.


def prepare_background_data(self):
    self.background_data = []
    background_dir = os.path.join(self.data_dir, BACKGROUND_NOISE_DIR_NAME)
    if not gfile.Exists(background_dir):
        return self.background_data
    with tf.compat.v1.Session(graph=tf.Graph()) as sess:
        wav_filename_placeholder = tf.compat.v1.placeholder(tf.string, [])
        wav_loader = io_ops.read_file(wav_filename_placeholder)
        wav_decoder = tf.audio.decode_wav(wav_loader, desired_channels=1)
        search_path = os.path.join(self.data_dir, BACKGROUND_NOISE_DIR_NAME, '*.wav')
        for wav_path in gfile.Glob(search_path):
            wav_data = sess.run(wav_decoder, feed_dict={wav_filename_placeholder: wav_path}).audio.flatten()
            self.background_data.append(wav_data)
        if not self.background_data:
            raise Exception('No background wav files were found in ' + search_path)

Searches a folder for background noise audio, and loads it into memory.

It's expected that the background audio samples will be in a subdirectory
named '_background_noise_' inside the 'data_dir' folder, as .wavs that match
the sample rate of the training data, but can be much longer in duration.

If the '_background_noise_' folder doesn't exist at all, this isn't an
error, it's just taken to mean that no background noise augmentation should
be used. If the folder does exist, but it's empty, that's treated as an
error.

Returns:
  List of raw PCM-encoded audio samples of background noise.

Raises:
  Exception: If files aren't found in the folder.


class VitsConfig(PretrainedConfig):
    model_type = 'vits'

    def __init__(self, vocab_size=38, hidden_size=192, num_hidden_layers=6, num_attention_heads=2, window_size=4, use_bias=True, ffn_dim=768, layerdrop=0.1, ffn_kernel_size=3, flow_size=192, spectrogram_bins=513, hidden_act='relu', hidden_dropout=0.1, attention_dropout=0.1, activation_dropout=0.1, initializer_range=0.02, layer_norm_eps=1e-05, use_stochastic_duration_prediction=True, num_speakers=1, speaker_embedding_size=0, upsample_initial_channel=512, upsample_rates=[8, 8, 2, 2], upsample_kernel_sizes=[16, 16, 4, 4], resblock_kernel_sizes=[3, 7, 11], resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]], leaky_relu_slope=0.1, depth_separable_channels=2, depth_separable_num_layers=3, duration_predictor_flow_bins=10, duration_predictor_tail_bound=5.0, duration_predictor_kernel_size=3, duration_predictor_dropout=0.5, duration_predictor_num_flows=4, duration_predictor_filter_channels=256, prior_encoder_num_flows=4, prior_encoder_num_wavenet_layers=4, posterior_encoder_num_wavenet_layers=16, wavenet_kernel_size=5, wavenet_dilation_rate=1, wavenet_dropout=0.0, speaking_rate=1.0, noise_scale=0.667, noise_scale_duration=0.8, sampling_rate=16000, **kwargs):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.window_size = window_size
        self.use_bias = use_bias
        self.ffn_dim = ffn_dim
        self.layerdrop = layerdrop
        self.ffn_kernel_size = ffn_kernel_size
        self.flow_size = flow_size
        self.spectrogram_bins = spectrogram_bins
        self.hidden_act = hidden_act
        self.hidden_dropout = hidden_dropout
        self.attention_dropout = attention_dropout
        self.activation_dropout = activation_dropout
        self.initializer_range = initializer_range
        self.layer_norm_eps = layer_norm_eps
        self.use_stochastic_duration_prediction = use_stochastic_duration_prediction
        self.num_speakers = num_speakers
        self.speaker_embedding_size = speaker_embedding_size
        self.upsample_initial_channel = upsample_initial_channel
        self.upsample_rates = upsample_rates
        self.upsample_kernel_sizes = upsample_kernel_sizes
        self.resblock_kernel_sizes = resblock_kernel_sizes
        self.resblock_dilation_sizes = resblock_dilation_sizes
        self.leaky_relu_slope = leaky_relu_slope
        self.depth_separable_channels = depth_separable_channels
        self.depth_separable_num_layers = depth_separable_num_layers
        self.duration_predictor_flow_bins = duration_predictor_flow_bins
        self.duration_predictor_tail_bound = duration_predictor_tail_bound
        self.duration_predictor_kernel_size = duration_predictor_kernel_size
        self.duration_predictor_dropout = duration_predictor_dropout
        self.duration_predictor_num_flows = duration_predictor_num_flows
        self.duration_predictor_filter_channels = duration_predictor_filter_channels
        self.prior_encoder_num_flows = prior_encoder_num_flows
        self.prior_encoder_num_wavenet_layers = prior_encoder_num_wavenet_layers
        self.posterior_encoder_num_wavenet_layers = posterior_encoder_num_wavenet_layers
        self.wavenet_kernel_size = wavenet_kernel_size
        self.wavenet_dilation_rate = wavenet_dilation_rate
        self.wavenet_dropout = wavenet_dropout
        self.speaking_rate = speaking_rate
        self.noise_scale = noise_scale
        self.noise_scale_duration = noise_scale_duration
        self.sampling_rate = sampling_rate
        if len(upsample_kernel_sizes) != len(upsample_rates):
            raise ValueError(f'The length of `upsample_kernel_sizes` ({len(upsample_kernel_sizes)}) must match the length of `upsample_rates` ({len(upsample_rates)})')
        super().__init__(**kwargs)

This is the configuration class to store the configuration of a [`VitsModel`]. It is used to instantiate a VITS
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 VITS
[facebook/mms-tts-eng](https://huggingface.co/facebook/mms-tts-eng) architecture.

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

Args:
    vocab_size (`int`, *optional*, defaults to 38):
        Vocabulary size of the VITS model. Defines the number of different tokens that can be represented by the
        `inputs_ids` passed to the forward method of [`VitsModel`].
    hidden_size (`int`, *optional*, defaults to 192):
        Dimensionality of the text encoder layers.
    num_hidden_layers (`int`, *optional*, defaults to 6):
        Number of hidden layers in the Transformer encoder.
    num_attention_heads (`int`, *optional*, defaults to 2):
        Number of attention heads for each attention layer in the Transformer encoder.
    window_size (`int`, *optional*, defaults to 4):
        Window size for the relative positional embeddings in the attention layers of the Transformer encoder.
    use_bias (`bool`, *optional*, defaults to `True`):
        Whether to use bias in the key, query, value projection layers in the Transformer encoder.
    ffn_dim (`int`, *optional*, defaults to 768):
        Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
    layerdrop (`float`, *optional*, defaults to 0.1):
        The LayerDrop probability for the encoder. See the [LayerDrop paper](see https://huggingface.co/papers/1909.11556)
        for more details.
    ffn_kernel_size (`int`, *optional*, defaults to 3):
        Kernel size of the 1D convolution layers used by the feed-forward network in the Transformer encoder.
    flow_size (`int`, *optional*, defaults to 192):
        Dimensionality of the flow layers.
    spectrogram_bins (`int`, *optional*, defaults to 513):
        Number of frequency bins in the target spectrogram.
    hidden_act (`str` or `function`, *optional*, defaults to `"relu"`):
        The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
        `"relu"`, `"selu"` and `"gelu_new"` are supported.
    hidden_dropout (`float`, *optional*, defaults to 0.1):
        The dropout probability for all fully connected layers in the embeddings and encoder.
    attention_dropout (`float`, *optional*, defaults to 0.1):
        The dropout ratio for the attention probabilities.
    activation_dropout (`float`, *optional*, defaults to 0.1):
        The dropout ratio for activations inside the fully connected layer.
    initializer_range (`float`, *optional*, defaults to 0.02):
        The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
    layer_norm_eps (`float`, *optional*, defaults to 1e-05):
        The epsilon used by the layer normalization layers.
    use_stochastic_duration_prediction (`bool`, *optional*, defaults to `True`):
        Whether to use the stochastic duration prediction module or the regular duration predictor.
    num_speakers (`int`, *optional*, defaults to 1):
        Number of speakers if this is a multi-speaker model.
    speaker_embedding_size (`int`, *optional*, defaults to 0):
        Number of channels used by the speaker embeddings. Is zero for single-speaker models.
    upsample_initial_channel (`int`, *optional*, defaults to 512):
        The number of input channels into the HiFi-GAN upsampling network.
    upsample_rates (`Tuple[int]` or `List[int]`, *optional*, defaults to `[8, 8, 2, 2]`):
        A tuple of integers defining the stride of each 1D convolutional layer in the HiFi-GAN upsampling network.
        The length of `upsample_rates` defines the number of convolutional layers and has to match the length of
        `upsample_kernel_sizes`.
    upsample_kernel_sizes (`Tuple[int]` or `List[int]`, *optional*, defaults to `[16, 16, 4, 4]`):
        A tuple of integers defining the kernel size of each 1D convolutional layer in the HiFi-GAN upsampling
        network. The length of `upsample_kernel_sizes` defines the number of convolutional layers and has to match
        the length of `upsample_rates`.
    resblock_kernel_sizes (`Tuple[int]` or `List[int]`, *optional*, defaults to `[3, 7, 11]`):
        A tuple of integers defining the kernel sizes of the 1D convolutional layers in the HiFi-GAN
        multi-receptive field fusion (MRF) module.
    resblock_dilation_sizes (`Tuple[Tuple[int]]` or `List[List[int]]`, *optional*, defaults to `[[1, 3, 5], [1, 3, 5], [1, 3, 5]]`):
        A nested tuple of integers defining the dilation rates of the dilated 1D convolutional layers in the
        HiFi-GAN multi-receptive field fusion (MRF) module.
    leaky_relu_slope (`float`, *optional*, defaults to 0.1):
        The angle of the negative slope used by the leaky ReLU activation.
    depth_separable_channels (`int`, *optional*, defaults to 2):
        Number of channels to use in each depth-separable block.
    depth_separable_num_layers (`int`, *optional*, defaults to 3):
        Number of convolutional layers to use in each depth-separable block.
    duration_predictor_flow_bins (`int`, *optional*, defaults to 10):
        Number of channels to map using the unonstrained rational spline in the duration predictor model.
    duration_predictor_tail_bound (`float`, *optional*, defaults to 5.0):
        Value of the tail bin boundary when computing the unconstrained rational spline in the duration predictor
        model.
    duration_predictor_kernel_size (`int`, *optional*, defaults to 3):
        Kernel size of the 1D convolution layers used in the duration predictor model.
    duration_predictor_dropout (`float`, *optional*, defaults to 0.5):
        The dropout ratio for the duration predictor model.
    duration_predictor_num_flows (`int`, *optional*, defaults to 4):
        Number of flow stages used by the duration predictor model.
    duration_predictor_filter_channels (`int`, *optional*, defaults to 256):
        Number of channels for the convolution layers used in the duration predictor model.
    prior_encoder_num_flows (`int`, *optional*, defaults to 4):
        Number of flow stages used by the prior encoder flow model.
    prior_encoder_num_wavenet_layers (`int`, *optional*, defaults to 4):
        Number of WaveNet layers used by the prior encoder flow model.
    posterior_encoder_num_wavenet_layers (`int`, *optional*, defaults to 16):
        Number of WaveNet layers used by the posterior encoder model.
    wavenet_kernel_size (`int`, *optional*, defaults to 5):
        Kernel size of the 1D convolution layers used in the WaveNet model.
    wavenet_dilation_rate (`int`, *optional*, defaults to 1):
        Dilation rates of the dilated 1D convolutional layers used in the WaveNet model.
    wavenet_dropout (`float`, *optional*, defaults to 0.0):
        The dropout ratio for the WaveNet layers.
    speaking_rate (`float`, *optional*, defaults to 1.0):
        Speaking rate. Larger values give faster synthesised speech.
    noise_scale (`float`, *optional*, defaults to 0.667):
        How random the speech prediction is. Larger values create more variation in the predicted speech.
    noise_scale_duration (`float`, *optional*, defaults to 0.8):
        How random the duration prediction is. Larger values create more variation in the predicted durations.
    sampling_rate (`int`, *optional*, defaults to 16000):
        The sampling rate at which the output audio waveform is digitalized expressed in hertz (Hz).

Example:

```python
>>> from transformers import VitsModel, VitsConfig

>>> # Initializing a "facebook/mms-tts-eng" style configuration
>>> configuration = VitsConfig()

>>> # Initializing a model (with random weights) from the "facebook/mms-tts-eng" style configuration
>>> model = VitsModel(configuration)

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