def _assert_tensorlike_all_close(self, sess: session.Session, tensorlike_value_1: core.TensorLike, tensorlike_value_2: core.TensorLike) -> None:
    if isinstance(tensorlike_value_1, core.Tensor):
        tensorlike_value_1 = tensorlike_value_1.eval(session=sess)
    if isinstance(tensorlike_value_2, core.Tensor):
        tensorlike_value_2 = tensorlike_value_2.eval(session=sess)
    self.assertAllClose(tensorlike_value_1, tensorlike_value_2)

Asserts that two different TensorLike values are "all close".

Args:
  sess: Session instance used to evaluate any tf.Tensors.
  tensorlike_value_1: A TensorLike value.
  tensorlike_value_2: A TensorLike value.


def get_size_with_aspect_ratio(image_size: Tuple[int, int], size: int, max_size: Optional[int]=None, mod_size: int=16) -> Tuple[int, int]:
    height, width = image_size
    raw_size = None
    if max_size is not None:
        min_original_size = float(min((height, width)))
        max_original_size = float(max((height, width)))
        if max_original_size / min_original_size * size > max_size:
            raw_size = max_size * min_original_size / max_original_size
            size = int(round(raw_size))
    if width < height:
        ow = size
        if max_size is not None and raw_size is not None:
            oh = int(raw_size * height / width)
        else:
            oh = int(size * height / width)
    elif height <= width and height == size or (width <= height and width == size):
        oh, ow = (height, width)
    else:
        oh = size
        if max_size is not None and raw_size is not None:
            ow = int(raw_size * width / height)
        else:
            ow = int(size * width / height)
    if mod_size is not None:
        ow_mod = np.mod(ow, mod_size)
        oh_mod = np.mod(oh, mod_size)
        ow = ow - ow_mod
        oh = oh - oh_mod
    return (oh, ow)

Computes the output image size given the input image size and the desired output size with multiple of divisible_size.

Args:
    image_size (`Tuple[int, int]`):
        The input image size.
    size (`int`):
        The desired output size.
    max_size (`int`, *optional*):
        The maximum allowed output size.
    mod_size (`int`, *optional*):
        The size to make multiple of mod_size.


def forward(self, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.FloatTensor]=None, encoder_outputs: Optional[Tuple[Tuple[torch.Tensor]]]=None, decoder_input_ids: Optional[torch.LongTensor]=None, decoder_attention_mask: Optional[torch.BoolTensor]=None, past_key_values: Optional[Tuple[Tuple[torch.Tensor]]]=None, context_input_ids: Optional[torch.LongTensor]=None, context_attention_mask: Optional[torch.LongTensor]=None, doc_scores: Optional[torch.FloatTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, output_retrieved: Optional[bool]=None, do_marginalize: Optional[bool]=None, reduce_loss: Optional[bool]=None, labels: Optional[torch.LongTensor]=None, n_docs: Optional[int]=None, **kwargs) -> RetrievAugLMMarginOutput:
    n_docs = n_docs if n_docs is not None else self.config.n_docs
    do_marginalize = do_marginalize if do_marginalize is not None else self.config.do_marginalize
    reduce_loss = reduce_loss if reduce_loss is not None else self.config.reduce_loss
    if labels is not None:
        if decoder_input_ids is None:
            decoder_input_ids = labels
        use_cache = False
    outputs = self.rag(input_ids=input_ids, attention_mask=attention_mask, encoder_outputs=encoder_outputs, decoder_input_ids=decoder_input_ids, decoder_attention_mask=decoder_attention_mask, context_input_ids=context_input_ids, context_attention_mask=context_attention_mask, doc_scores=doc_scores, past_key_values=past_key_values, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, output_retrieved=output_retrieved, n_docs=n_docs)
    loss = None
    logits = outputs.logits
    if labels is not None:
        assert decoder_input_ids is not None
        loss = self.get_nll(outputs.logits, outputs.doc_scores, labels, reduce_loss=reduce_loss, epsilon=self.config.label_smoothing, n_docs=n_docs)
    if do_marginalize:
        logits = self.marginalize(logits, outputs.doc_scores, n_docs)
    return RetrievAugLMMarginOutput(loss=loss, logits=logits, doc_scores=outputs.doc_scores, past_key_values=outputs.past_key_values, context_input_ids=outputs.context_input_ids, context_attention_mask=outputs.context_attention_mask, retrieved_doc_embeds=outputs.retrieved_doc_embeds, retrieved_doc_ids=outputs.retrieved_doc_ids, question_encoder_last_hidden_state=outputs.question_encoder_last_hidden_state, question_enc_hidden_states=outputs.question_enc_hidden_states, question_enc_attentions=outputs.question_enc_attentions, generator_enc_last_hidden_state=outputs.generator_enc_last_hidden_state, generator_enc_hidden_states=outputs.generator_enc_hidden_states, generator_enc_attentions=outputs.generator_enc_attentions, generator_dec_hidden_states=outputs.generator_dec_hidden_states, generator_dec_attentions=outputs.generator_dec_attentions, generator_cross_attentions=outputs.generator_cross_attentions)

input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
    Indices of input sequence tokens in the vocabulary. [`RagConfig`], used to initialize the model, specifies
    which generator to use, it also specifies a compatible generator tokenizer. Use that tokenizer class to
    obtain the indices.

    [What are input IDs?](../glossary#input-ids)
encoder_outputs (`tuple(tuple(torch.FloatTensor)`, *optional*)
    Tuple consists of (`generator_enc_last_hidden_state`, *optional*: `generator_enc_hidden_states`,
    *optional*: `generator_enc_attentions`). `generator_enc_last_hidden_state` of shape `(batch_size, n_docs *
    sequence_length, hidden_size)` is a sequence of hidden-states at the output of the last layer of the
    generator's encoder.

    Used by the ([`RagModel`]) model during decoding.
decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
    Provide for generation tasks. `None` by default, construct as per instructions for the generator model
    you're using with your RAG instance.
decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size,  target_sequence_length)`, *optional*):
    Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
    be used by default.
context_input_ids (`torch.LongTensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`, *optional*, returned when *output_retrieved=True*):
    Input IDs post-processed from the retrieved documents and the question encoder `input_ids` by the
    retriever. If the model was not initialized with a `retriever` ``context_input_ids` has to be provided to
    the forward pass. `context_input_ids` are returned by [`~RagRetriever.__call__`].
context_attention_mask (`torch.LongTensor` of shape `(batch_size * config.n_docs, config.max_combined_length)`,*optional*, returned when *output_retrieved=True*):
    Attention mask post-processed from the retrieved documents and the question encoder `input_ids` by the
    retriever. If the model has is not initialized with a `retriever` `context_attention_mask` has to be
    provided to the forward pass. `context_attention_mask` are returned by [`~RagRetriever.__call__`].
doc_scores (`torch.FloatTensor` of shape `(batch_size, config.n_docs)`):
    Score between each retrieved document embeddings (see `retrieved_doc_embeds`) and
    `question_encoder_last_hidden_state`. If the model has is not initialized with a `retriever` `doc_scores`
    has to be provided to the forward pass. `doc_scores` can be computed via
    `question_encoder_last_hidden_state` and `retrieved_doc_embeds`, see examples for more information.
output_retrieved (`bool`, *optional*):
    Whether or not to return the `retrieved_doc_embeds`, `retrieved_doc_ids`, `context_input_ids` and
    `context_attention_mask`. See returned tensors for more detail.
do_marginalize (`bool`, *optional*):
    If `True`, the logits are marginalized over all documents by making use of
    `torch.nn.functional.log_softmax`.
reduce_loss (`bool`, *optional*):
    Only relevant if `labels` is passed. If `True`, the NLL loss is reduced using the `torch.Tensor.sum`
    operation.
n_docs (`int`, *optional*):
    The number of documents to retrieve.

Example:

```python
>>> from transformers import AutoTokenizer, RagRetriever, RagTokenForGeneration
>>> import torch

>>> tokenizer = AutoTokenizer.from_pretrained("facebook/rag-token-nq")
>>> retriever = RagRetriever.from_pretrained(
...     "facebook/rag-token-nq", index_name="exact", use_dummy_dataset=True
... )
>>> # initialize with RagRetriever to do everything in one forward call
>>> model = RagTokenForGeneration.from_pretrained("facebook/rag-token-nq", retriever=retriever)

>>> inputs = tokenizer("How many people live in Paris?", return_tensors="pt")
>>> targets = tokenizer(text_target="In Paris, there are 10 million people.", return_tensors="pt")
>>> input_ids = inputs["input_ids"]
>>> labels = targets["input_ids"]
>>> outputs = model(input_ids=input_ids, labels=labels)

>>> # or use retriever separately
>>> model = RagTokenForGeneration.from_pretrained("facebook/rag-token-nq", use_dummy_dataset=True)
>>> # 1. Encode
>>> question_hidden_states = model.question_encoder(input_ids)[0]
>>> # 2. Retrieve
>>> docs_dict = retriever(input_ids.numpy(), question_hidden_states.detach().numpy(), return_tensors="pt")
>>> doc_scores = torch.bmm(
...     question_hidden_states.unsqueeze(1), docs_dict["retrieved_doc_embeds"].float().transpose(1, 2)
... ).squeeze(1)
>>> # 3. Forward to generator
>>> outputs = model(
...     context_input_ids=docs_dict["context_input_ids"],
...     context_attention_mask=docs_dict["context_attention_mask"],
...     doc_scores=doc_scores,
...     decoder_input_ids=labels,
... )

>>> # or directly generate
>>> generated = model.generate(
...     context_input_ids=docs_dict["context_input_ids"],
...     context_attention_mask=docs_dict["context_attention_mask"],
...     doc_scores=doc_scores,
... )
>>> generated_string = tokenizer.batch_decode(generated, skip_special_tokens=True)
```


def write_examples(fp, examples):

    def write_tensor(fp, name, x):
        """Write tensor in file format supported by TFLITE example."""
        fp.write('name,%s\n' % name)
        fp.write('dtype,%s\n' % x.dtype)
        fp.write('shape,' + ','.join(map(str, x.shape)) + '\n')
        fp.write('values,' + format_result(x) + '\n')
    fp.write('test_cases,%d\n' % len(examples))
    for example in examples:
        fp.write('inputs,%d\n' % len(example['inputs']))
        for name, value in example['inputs'].items():
            if value is not None:
                write_tensor(fp, name, value)
        fp.write('outputs,%d\n' % len(example['outputs']))
        for name, value in example['outputs'].items():
            write_tensor(fp, name, value)

Given a list `examples`, write a text format representation.

The file format is csv like with a simple repeated pattern. We would ike
to use proto here, but we can't yet due to interfacing with the Android
team using this format.

Args:
  fp: File-like object to write to.
  examples: Example dictionary consisting of keys "inputs" and "outputs"


def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor]=None, encoder_hidden_states: Optional[torch.Tensor]=None, encoder_attention_mask: Optional[torch.Tensor]=None, layer_head_mask: Optional[torch.Tensor]=None, cross_attn_layer_head_mask: Optional[torch.Tensor]=None, past_key_value: Optional[Cache]=None, output_attentions: Optional[bool]=False, use_cache: Optional[bool]=True, cache_position: Optional[torch.Tensor]=None) -> torch.Tensor:
    residual = hidden_states
    hidden_states = self.self_attn_layer_norm(hidden_states)
    hidden_states, self_attn_weights, past_key_value = self.self_attn(hidden_states=hidden_states, past_key_value=past_key_value, attention_mask=attention_mask, layer_head_mask=layer_head_mask, output_attentions=output_attentions, cache_position=cache_position)
    hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
    hidden_states = residual + hidden_states
    cross_attn_weights = None
    if encoder_hidden_states is not None:
        residual = hidden_states
        hidden_states = self.encoder_attn_layer_norm(hidden_states)
        hidden_states, cross_attn_weights, past_key_value = self.encoder_attn(hidden_states=hidden_states, key_value_states=encoder_hidden_states, attention_mask=encoder_attention_mask, layer_head_mask=cross_attn_layer_head_mask, past_key_value=past_key_value, output_attentions=output_attentions)
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states
    residual = hidden_states
    hidden_states = self.final_layer_norm(hidden_states)
    hidden_states = self.activation_fn(self.fc1(hidden_states))
    hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
    hidden_states = self.fc2(hidden_states)
    hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
    hidden_states = residual + hidden_states
    outputs = (hidden_states,)
    if output_attentions:
        outputs += (self_attn_weights, cross_attn_weights)
    if use_cache:
        outputs += (past_key_value,)
    return outputs

Args:
    hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
    attention_mask (`torch.FloatTensor`): attention mask of size
        `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
    encoder_hidden_states (`torch.FloatTensor`):
        cross attention input to the layer of shape `(batch, seq_len, embed_dim)`
    encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size
        `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
    layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size
        `(encoder_attention_heads,)`.
    cross_attn_layer_head_mask (`torch.FloatTensor`): mask for cross-attention heads in a given layer of
        size `(decoder_attention_heads,)`.
    past_key_value (`Tuple(torch.FloatTensor)`): cached past key and value projection states
    output_attentions (`bool`, *optional*):
        Whether or not to return the attentions tensors of all attention layers. See `attentions` under
        returned tensors for more detail.
    cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
        Indices depicting the position of the input sequence tokens in the sequence. It is used to update the
        cache in the correct position and to infer the complete sequence length.


def retrieve_from_golden():
    out_dict = dict()
    with open(CUDA_CC_GOLDEN_DIR) as g_file:
        for line in g_file:
            line_items = line.split(',')
            val_list = []
            for item in line_items[1:]:
                val_list.append(item.strip('\n'))
            out_dict[line_items[0]] = val_list
    return out_dict

Retrieves list of all CUDA compute capability from a golden file.

The following file is set as default:
  `./golden/compute_capability_golden.csv`

Returns:
  Dictionary that lists of all CUDA compute capability in the following
  format:
    {'<GPU name>': ['<version major>.<version minor>', ...], ...}

  If there are multiple versions available for a given GPU, then it
  appends all supported versions in the value list (in the key-value
  pair.)


def _get_main_op_tensor(meta_graph_def_to_load, init_op_key=constants.MAIN_OP_KEY):
    collection_def = meta_graph_def_to_load.collection_def
    init_op = None
    if init_op_key in collection_def:
        init_op_list = collection_def[init_op_key].node_list.value
        if len(init_op_list) != 1:
            raise RuntimeError(f'Expected exactly one SavedModel init op. Found {len(init_op_list)}: {init_op_list}.')
        init_op = ops.get_collection(init_op_key)[0]
    return init_op

Gets the main op tensor, if one exists.

Args:
  meta_graph_def_to_load: The meta graph def from the SavedModel to be loaded.
  init_op_key: name of the collection to check; should be one of MAIN_OP_KEY
    or the deprecated LEGACY_INIT_OP_KEY

Returns:
  The main op tensor, if it exists and `None` otherwise.

Raises:
  RuntimeError: If the collection def corresponding to the main op key has
      other than exactly one tensor.


def set_tensor_shapes(tensors, shapes):
    if shapes:
        tensor_names_to_tensor = {get_tensor_name(tensor): tensor for tensor in tensors}
        for name, shape in shapes.items():
            if name not in tensor_names_to_tensor:
                raise ValueError("Invalid tensor '{}' found in tensor shapes map.".format(name))
            if shape is not None:
                tensor = tensor_names_to_tensor[name]
                try:
                    tensor.set_shape(shape)
                except ValueError as error:
                    message = "The shape of tensor '{0}' cannot be changed from {1} to {2}. {3}".format(name, tensor.shape, shape, str(error))
                    raise ValueError(message)

Sets Tensor shape for each tensor if the shape is defined.

Args:
  tensors: TensorFlow tensor.Tensor.
  shapes: Dict of strings representing input tensor names to list of
    integers representing input shapes (e.g., {"foo": : [1, 16, 16, 3]}).

Raises:
  ValueError:
    `shapes` contains an invalid tensor.
    `shapes` contains an invalid shape for a valid tensor.


def add(inputs, **kwargs):
    return Add(**kwargs)(inputs)

Functional interface to the `keras.layers.Add` layer.

Args:
    inputs: A list of input tensors with the same shape.
    **kwargs: Standard layer keyword arguments.

Returns:
    A tensor as the sum of the inputs. It has the same shape as the inputs.

Examples:

>>> input_shape = (2, 3, 4)
>>> x1 = np.random.rand(*input_shape)
>>> x2 = np.random.rand(*input_shape)
>>> y = keras.layers.add([x1, x2])

Usage in a Keras model:

>>> input1 = keras.layers.Input(shape=(16,))
>>> x1 = keras.layers.Dense(8, activation='relu')(input1)
>>> input2 = keras.layers.Input(shape=(32,))
>>> x2 = keras.layers.Dense(8, activation='relu')(input2)
>>> added = keras.layers.add([x1, x2])
>>> out = keras.layers.Dense(4)(added)
>>> model = keras.models.Model(inputs=[input1, input2], outputs=out)


def get_chief_queue_runner(self):
    if self._gradients_applied is False:
        raise ValueError('Should be called after apply_gradients().')
    return self._chief_queue_runner

Returns the QueueRunner for the chief to execute.

This includes the operations to synchronize replicas: aggregate gradients,
apply to variables, increment global step, insert tokens to token queue.

Note that this can only be called after calling apply_gradients() which
actually generates this queuerunner.

Returns:
  A `QueueRunner` for chief to execute.

Raises:
  ValueError: If this is called before apply_gradients().


def numpy_function(func=None, inp=None, Tout=None, stateful=True, name=None):
    decorator = _check_args_and_maybe_make_decorator(numpy_function, 'tf.numpy_function', func=func, inp=inp, Tout=Tout, stateful=stateful, name=name)
    if decorator is not None:
        return decorator
    return py_func_common(func, inp, Tout, stateful=stateful, name=name)

Wraps a python function and uses it as a TensorFlow op.

Given a python function `func` wrap this function as an operation in a
`tf.function`. `func` must take numpy arrays as its arguments and
return numpy arrays as its outputs.

There are two ways to use `tf.numpy_function`.

### As a decorator

When using `tf.numpy_function` as a decorator:

* you must set `Tout`
* you may set `name`
* you must not set `func` or `inp`

>>> @tf.numpy_function(Tout=tf.float32)
... def my_numpy_func(x):
...   # x will be a numpy array with the contents of the input to the
...   # tf.function
...   print(f'executing eagerly, {x=}')
...   return np.sinh(x)

The function runs eagerly:

>>> my_numpy_func(1.0).numpy()
executing eagerly, x=1.0
1.17520

The behavior doesn't change inside a `tf.function`:

>>> @tf.function(input_signature=[tf.TensorSpec(None, tf.float32)])
... def tf_function(input):
...   y = tf.numpy_function(my_numpy_func, [input], tf.float32)
...   return y
>>> tf_function(tf.constant(1.)).numpy()
executing eagerly, x=array(1.)
1.17520

### Inplace

This form can be useful if you don't control the function's source,
but it is harder to read.

Here is the same function with no decorator:

>>> def my_func(x):
...   # x will be a numpy array with the contents of the input to the
...   # tf.function
...   print(f'executing eagerly, {x=}')
...   return np.sinh(x)

To run `tf.numpy_function` in-place, pass the function, its inputs, and the
output type in a single call to `tf.numpy_function`:

>>> tf.numpy_function(my_func, [tf.constant(1.0)], tf.float32)
executing eagerly, x=array(1.)
1.17520

### More info

Comparison to `tf.py_function`:
`tf.py_function` and `tf.numpy_function` are very similar, except that
`tf.numpy_function` takes numpy arrays, and not `tf.Tensor`s. If you want the
function to contain `tf.Tensors`, and have any TensorFlow operations executed
in the function be differentiable, please use `tf.py_function`.

Note: We recommend to avoid using `tf.numpy_function` outside of
prototyping and experimentation due to the following known limitations:

* Calling `tf.numpy_function` will acquire the Python Global Interpreter Lock
  (GIL) that allows only one thread to run at any point in time. This will
  preclude efficient parallelization and distribution of the execution of the
  program. Therefore, you are discouraged to use `tf.numpy_function` outside
  of prototyping and experimentation.

* The body of the function (i.e. `func`) will not be serialized in a
  `tf.SavedModel`. Therefore, you should not use this function if you need to
  serialize your model and restore it in a different environment.

* The operation must run in the same address space as the Python program
  that calls `tf.numpy_function()`. If you are using distributed
  TensorFlow, you must run a `tf.distribute.Server` in the same process as the
  program that calls `tf.numpy_function`  you must pin the created
  operation to a device in that server (e.g. using `with tf.device():`).

* Currently `tf.numpy_function` is not compatible with XLA. Calling
  `tf.numpy_function` inside `tf.function(jit_compile=True)` will raise an
  error.

* Since the function takes numpy arrays, you cannot take gradients
  through a numpy_function. If you require something that is differentiable,
  please consider using tf.py_function.

Args:
  func: A Python function, which accepts `numpy.ndarray` objects as arguments
    and returns a list of `numpy.ndarray` objects (or a single
    `numpy.ndarray`). This function must accept as many arguments as there are
    tensors in `inp`, and these argument types will match the corresponding
    `tf.Tensor` objects in `inp`. The returns `numpy.ndarray`s must match the
    number and types defined `Tout`. Important Note: Input and output
    `numpy.ndarray`s of `func` are not guaranteed to be copies. In some cases
    their underlying memory will be shared with the corresponding TensorFlow
    tensors. In-place modification or storing `func` input or return values in
    python datastructures without explicit (np.)copy can have
    non-deterministic consequences.
  inp: A list of `tf.Tensor` objects.
  Tout: A list or tuple of tensorflow data types or a single tensorflow data
    type if there is only one, indicating what `func` returns.
  stateful: (Boolean.) Setting this argument to False tells the runtime to
    treat the function as stateless, which enables certain optimizations. A
    function is stateless when given the same input it will return the same
    output and have no side effects; its only purpose is to have a return
    value. The behavior for a stateful function with the `stateful` argument
    False is undefined. In particular, caution should be taken when mutating
    the input arguments as this is a stateful operation.
  name: (Optional) A name for the operation.

Returns:
  * If `func` is `None` this returns a decorator that will ensure the
    decorated function will always run with eager execution even if called
    from a `tf.function`/`tf.Graph`.
  * If used `func` is not `None` this executes `func` with eager execution
    and returns the result: A single or list of `tf.Tensor` which `func`
    computes.


def det(x):
    if any_symbolic_tensors((x,)):
        return Det().symbolic_call(x)
    return _det(x)

Computes the determinant of a square tensor.

Args:
    x: Input tensor of shape `(..., M, M)`.

Returns:
    A tensor of shape `(...,)` representing the determinant of `x`.


Args:
  pipeline_options: Instance of ``PipelineOptions`` or dict of options and
    values (like ``RuntimeValueProvider.runtime_options``).


def mel_to_hertz(mels: Union[float, np.ndarray], mel_scale: str='htk') -> Union[float, np.ndarray]:
    if mel_scale not in ['slaney', 'htk', 'kaldi']:
        raise ValueError('mel_scale should be one of "htk", "slaney" or "kaldi".')
    if mel_scale == 'htk':
        return 700.0 * (np.power(10, mels / 2595.0) - 1.0)
    elif mel_scale == 'kaldi':
        return 700.0 * (np.exp(mels / 1127.0) - 1.0)
    min_log_hertz = 1000.0
    min_log_mel = 15.0
    logstep = np.log(6.4) / 27.0
    freq = 200.0 * mels / 3.0
    if isinstance(mels, np.ndarray):
        log_region = mels >= min_log_mel
        freq[log_region] = min_log_hertz * np.exp(logstep * (mels[log_region] - min_log_mel))
    elif mels >= min_log_mel:
        freq = min_log_hertz * np.exp(logstep * (mels - min_log_mel))
    return freq

Convert frequency from mels to hertz.

Args:
    mels (`float` or `np.ndarray`):
        The frequency, or multiple frequencies, in mels.
    mel_scale (`str`, *optional*, `"htk"`):
        The mel frequency scale to use, `"htk"`, `"kaldi"` or `"slaney"`.

Returns:
    `float` or `np.ndarray`: The frequencies in hertz.


def _handle_per_output_metrics(self, metrics_dict, y_true, y_pred, mask, weights=None):
    metric_results = []
    for metric_name, metric_fn in metrics_dict.items():
        with backend.name_scope(metric_name):
            metric_result = training_utils_v1.call_metric_function(metric_fn, y_true, y_pred, weights=weights, mask=mask)
            metric_results.append(metric_result)
    return metric_results

Calls metric functions for a single output.

Args:
  metrics_dict: A dict with metric names as keys and metric fns as values.
  y_true: Target output.
  y_pred: Predicted output.
  mask: Computed mask value for the current output.
  weights: Weights to be applied on the current output.

Returns:
  A list of metric result tensors.


class Logcat(base_service.BaseService):
    OUTPUT_FILE_TYPE = 'logcat'

    def __init__(self, android_device, configs=None):
        super().__init__(android_device, configs)
        self._ad = android_device
        self._adb_logcat_process = None
        self._adb_logcat_file_obj = None
        self.adb_logcat_file_path = None
        self._config = configs if configs else Config()

    def _enable_logpersist(self):
        """Attempts to enable logpersist daemon to persist logs."""
        if not self._ad.is_rootable:
            return
        logpersist_warning = '%s encountered an error enabling persistent logs, logs may not get saved.'
        if not self._ad.adb.has_shell_command('logpersist.start'):
            logging.warning(logpersist_warning, self)
            return
        try:
            self._ad.adb.shell('logpersist.stop --clear')
            self._ad.adb.shell('logpersist.start')
        except adb.AdbError:
            logging.warning(logpersist_warning, self)

    def _is_timestamp_in_range(self, target, begin_time, end_time):
        low = mobly_logger.logline_timestamp_comparator(begin_time, target) <= 0
        high = mobly_logger.logline_timestamp_comparator(end_time, target) >= 0
        return low and high

    def create_output_excerpts(self, test_info):
        """Convenient method for creating excerpts of adb logcat.

    This copies logcat lines from self.adb_logcat_file_path to an excerpt
    file, starting from the location where the previous excerpt ended.

    Call this method at the end of: `setup_class`, `teardown_test`, and
    `teardown_class`.

    Args:
      test_info: `self.current_test_info` in a Mobly test.

    Returns:
      List of strings, the absolute paths to excerpt files.
    """
        dest_path = test_info.output_path
        utils.create_dir(dest_path)
        filename = self._ad.generate_filename(self.OUTPUT_FILE_TYPE, test_info, 'txt')
        excerpt_file_path = os.path.join(dest_path, filename)
        with open(excerpt_file_path, 'w', encoding='utf-8', errors='replace', newline='') as out:
            while self._adb_logcat_file_obj:
                line = self._adb_logcat_file_obj.readline()
                if not line:
                    break
                out.write(line)
        self._ad.log.debug('logcat excerpt created at: %s', excerpt_file_path)
        return [excerpt_file_path]

    @property
    def is_alive(self):
        return True if self._adb_logcat_process else False

    def clear_adb_log(self):
        """Clears cached adb content."""
        try:
            self._ad.adb.logcat('-c')
        except adb.AdbError as e:
            if b'failed to clear' in e.stderr:
                self._ad.log.warning('Encountered known Android error to clear logcat.')
            else:
                raise

    def _assert_not_running(self):
        """Asserts the logcat service is not running.

    Raises:
      Error, if the logcat service is running.
    """
        if self.is_alive:
            raise Error(self._ad, 'Logcat thread is already running, cannot start another one.')

    def update_config(self, new_config):
        """Updates the configuration for the service.

    The service needs to be stopped before updating, and explicitly started
    after the update.

    This will reset the service. Previous output files may be orphaned if
    output path is changed.

    Args:
      new_config: Config, the new config to use.
    """
        self._assert_not_running()
        self._ad.log.info('[LogcatService] Changing config from %s to %s', self._config, new_config)
        self._config = new_config

    def _open_logcat_file(self):
        """Create a file object that points to the beginning of the logcat file.
    Wait for the logcat file to be created by the subprocess if it doesn't
    exist.
    """
        if not self._adb_logcat_file_obj:
            deadline = time.perf_counter() + CREATE_LOGCAT_FILE_TIMEOUT_SEC
            while not os.path.exists(self.adb_logcat_file_path):
                if time.perf_counter() > deadline:
                    raise Error(self._ad, 'Timeout while waiting for logcat file to be created.')
                time.sleep(1)
            self._adb_logcat_file_obj = open(self.adb_logcat_file_path, 'r', encoding='utf-8', errors='replace', newline='')
            self._adb_logcat_file_obj.seek(0, os.SEEK_END)

    def _close_logcat_file(self):
        """Closes and resets the logcat file object, if it exists."""
        if self._adb_logcat_file_obj:
            self._adb_logcat_file_obj.close()
            self._adb_logcat_file_obj = None

    def start(self):
        """Starts a standing adb logcat collection.

    The collection runs in a separate subprocess and saves logs in a file.
    """
        if self._ad.is_bootloader:
            self._ad.log.warning('Skip starting logcat because the device is in fastboot mode.')
            return
        self._assert_not_running()
        if self._config.clear_log:
            self.clear_adb_log()
        self._start()
        self._open_logcat_file()

    def _start(self):
        """The actual logic of starting logcat."""
        self._enable_logpersist()
        if self._config.output_file_path:
            self._close_logcat_file()
            self.adb_logcat_file_path = self._config.output_file_path
        if not self.adb_logcat_file_path:
            f_name = self._ad.generate_filename(self.OUTPUT_FILE_TYPE, extension_name='txt')
            logcat_file_path = os.path.join(self._ad.log_path, f_name)
            self.adb_logcat_file_path = logcat_file_path
        utils.create_dir(os.path.dirname(self.adb_logcat_file_path))
        cmd = ' "%s" -s %s logcat -v threadtime -T 1 %s >> "%s" ' % (adb.ADB, self._ad.serial, self._config.logcat_params, self.adb_logcat_file_path)
        process = utils.start_standing_subprocess(cmd, shell=True)
        self._adb_logcat_process = process

    def stop(self):
        """Stops the adb logcat service."""
        self._close_logcat_file()
        self._stop()

    def _stop(self):
        """Stops the background process for logcat."""
        if not self._adb_logcat_process:
            return
        try:
            utils.stop_standing_subprocess(self._adb_logcat_process)
        except Exception:
            self._ad.log.exception('Failed to stop adb logcat.')
        self._adb_logcat_process = None

    def pause(self):
        """Pauses logcat.

    Note: the service is unable to collect the logs when paused, if more
    logs are generated on the device than the device's log buffer can hold,
    some logs would be lost.
    """
        self._stop()

    def resume(self):
        """Resumes a paused logcat service."""
        self._assert_not_running()
        self._start()

Android logcat service for Mobly's AndroidDevice controller.

Attributes:
  adb_logcat_file_path: string, path to the file that the service writes
    adb logcat to by default.


def _to_numpy(a):
    if isinstance(a, ops.EagerTensor):
        return a.numpy()
    if isinstance(a, tensor.Tensor):
        sess = ops.get_default_session()
        return sess.run(a)
    if isinstance(a, indexed_slices.IndexedSlicesValue):
        arr = np.zeros(a.dense_shape)
        assert len(a.values) == len(a.indices), 'IndexedSlicesValue has %s value slices but %s indices\n%s' % (a.values, a.indices, a)
        for values_slice, index in zip(a.values, a.indices):
            assert 0 <= index < len(arr), 'IndexedSlicesValue has invalid index %s\n%s' % (index, a)
            arr[index] += values_slice
        return arr
    return a

Converts Tensors, EagerTensors, and IndexedSlicesValue to numpy arrays.

Args:
  a: any value.

Returns:
  If a is EagerTensor or Tensor, returns the evaluation of a by calling
  numpy() or run(). If a is IndexedSlicesValue, constructs the corresponding
  dense numpy array. Otherwise returns a unchanged.


def conv1d(inputs, filters, kernel_size, strides=1, padding='valid', data_format='channels_last', dilation_rate=1, activation=None, use_bias=True, kernel_initializer=None, bias_initializer=init_ops.zeros_initializer(), kernel_regularizer=None, bias_regularizer=None, activity_regularizer=None, kernel_constraint=None, bias_constraint=None, trainable=True, name=None, reuse=None):
    warnings.warn('`tf.layers.conv1d` is deprecated and will be removed in a future version. Please Use `tf.keras.layers.Conv1D` instead.')
    layer = Conv1D(filters=filters, kernel_size=kernel_size, strides=strides, padding=padding, data_format=data_format, dilation_rate=dilation_rate, activation=activation, use_bias=use_bias, kernel_initializer=kernel_initializer, bias_initializer=bias_initializer, kernel_regularizer=kernel_regularizer, bias_regularizer=bias_regularizer, activity_regularizer=activity_regularizer, kernel_constraint=kernel_constraint, bias_constraint=bias_constraint, trainable=trainable, name=name, _reuse=reuse, _scope=name)
    return layer.apply(inputs)

Functional interface for 1D convolution layer (e.g. temporal convolution).

This layer creates a convolution kernel that is convolved
(actually cross-correlated) with the layer input to produce a tensor of
outputs. If `use_bias` is True (and a `bias_initializer` is provided),
a bias vector is created and added to the outputs. Finally, if
`activation` is not `None`, it is applied to the outputs as well.

Args:
  inputs: Tensor input.
  filters: Integer, the dimensionality of the output space (i.e. the number
    of filters in the convolution).
  kernel_size: An integer or tuple/list of a single integer, specifying the
    length of the 1D convolution window.
  strides: An integer or tuple/list of a single integer,
    specifying the stride length of the convolution.
    Specifying any stride value != 1 is incompatible with specifying
    any `dilation_rate` value != 1.
  padding: One of `"valid"` or `"same"` (case-insensitive).
    `"valid"` means no padding. `"same"` results in padding evenly to
    the left/right or up/down of the input such that output has the same
    height/width dimension as the input.
  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)`.
  dilation_rate: An integer or tuple/list of a single integer, specifying
    the dilation rate to use for dilated convolution.
    Currently, specifying any `dilation_rate` value != 1 is
    incompatible with specifying any `strides` value != 1.
  activation: Activation function. Set it to None to maintain a
    linear activation.
  use_bias: Boolean, whether the layer uses a bias.
  kernel_initializer: An initializer for the convolution kernel.
  bias_initializer: An initializer for the bias vector. If None, the default
    initializer will be used.
  kernel_regularizer: Optional regularizer for the convolution kernel.
  bias_regularizer: Optional regularizer for the bias vector.
  activity_regularizer: Optional regularizer function for the output.
  kernel_constraint: Optional projection function to be applied to the
      kernel 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 variable and must return the
      projected variable (which must have the same shape). Constraints are
      not safe to use when doing asynchronous distributed training.
  bias_constraint: Optional projection function to be applied to the
      bias after being updated by an `Optimizer`.
  trainable: Boolean, if `True` also add variables to the graph collection
    `GraphKeys.TRAINABLE_VARIABLES` (see `tf.Variable`).
  name: A string, the name of the layer.
  reuse: Boolean, whether to reuse the weights of a previous layer
    by the same name.

Returns:
  Output tensor.

Raises:
  ValueError: if eager execution is enabled.


def _global_batch_size(self):
    return True

`make_dataset_iterator` and `make_numpy_iterator` use global batch size.

`make_input_fn_iterator` assumes per-replica batching.

Returns:
  Boolean.


class MoeModelOutputWithPast(ModelOutput):
    last_hidden_state: Optional[torch.FloatTensor] = None
    past_key_values: Optional[Cache] = None
    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
    router_logits: Optional[Tuple[torch.FloatTensor]] = None

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

Args:
    last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
        Sequence of hidden-states at the output of the last layer of the model.
    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
        It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).

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

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

        Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
        heads.
    router_logits (`tuple(torch.FloatTensor)`, *optional*, returned when `output_router_probs=True` and `config.add_router_probs=True` is passed or when `config.output_router_probs=True`):
        Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, sequence_length, num_experts)`.

        Raw router logtis (post-softmax) that are computed by MoE routers, these terms are used to compute the auxiliary
        loss for Mixture of Experts models.


def weighted_average(input_tensor: torch.Tensor, weights: Optional[torch.Tensor]=None, dim=None) -> torch.Tensor:
    if weights is not None:
        weighted_tensor = torch.where(weights != 0, input_tensor * weights, torch.zeros_like(input_tensor))
        sum_weights = torch.clamp(weights.sum(dim=dim) if dim else weights.sum(), min=1.0)
        return (weighted_tensor.sum(dim=dim) if dim else weighted_tensor.sum()) / sum_weights
    else:
        return input_tensor.mean(dim=dim)

Computes the weighted average of a given tensor across a given `dim`, masking values associated with weight zero,
meaning instead of `nan * 0 = nan` you will get `0 * 0 = 0`.

Args:
    input_tensor (`torch.FloatTensor`):
        Input tensor, of which the average must be computed.
    weights (`torch.FloatTensor`, *optional*):
        Weights tensor, of the same shape as `input_tensor`.
    dim (`int`, *optional*):
        The dim along which to average `input_tensor`.

Returns:
    `torch.FloatTensor`: The tensor with values averaged along the specified `dim`.


def _compute_mask_indices(shape: Tuple[int, int], mask_prob: float, mask_length: int, attention_mask: Optional[torch.LongTensor]=None, min_masks: int=0) -> np.ndarray:
    batch_size, sequence_length = shape
    if mask_length < 1:
        raise ValueError('`mask_length` has to be bigger than 0.')
    if mask_length > sequence_length:
        raise ValueError(f'`mask_length` has to be smaller than `sequence_length`, but got `mask_length`: {mask_length} and `sequence_length`: {sequence_length}`')
    epsilon = np.random.rand(1).item()

    def compute_num_masked_span(input_length):
        """Given input length, compute how many spans should be masked"""
        num_masked_span = int(mask_prob * input_length / mask_length + epsilon)
        num_masked_span = max(num_masked_span, min_masks)
        if num_masked_span * mask_length > sequence_length:
            num_masked_span = sequence_length // mask_length
        if input_length - (mask_length - 1) < num_masked_span:
            num_masked_span = max(input_length - (mask_length - 1), 0)
        return num_masked_span
    input_lengths = attention_mask.detach().sum(-1).tolist() if attention_mask is not None else [sequence_length for _ in range(batch_size)]
    spec_aug_mask = np.zeros((batch_size, sequence_length), dtype=bool)
    spec_aug_mask_idxs = []
    max_num_masked_span = compute_num_masked_span(sequence_length)
    if max_num_masked_span == 0:
        return spec_aug_mask
    for input_length in input_lengths:
        num_masked_span = compute_num_masked_span(input_length)
        spec_aug_mask_idx = np.random.choice(np.arange(input_length - (mask_length - 1)), num_masked_span, replace=False)
        if len(spec_aug_mask_idx) == 0:
            dummy_mask_idx = sequence_length - 1
        else:
            dummy_mask_idx = spec_aug_mask_idx[0]
        spec_aug_mask_idx = np.concatenate([spec_aug_mask_idx, np.ones(max_num_masked_span - num_masked_span, dtype=np.int32) * dummy_mask_idx])
        spec_aug_mask_idxs.append(spec_aug_mask_idx)
    spec_aug_mask_idxs = np.array(spec_aug_mask_idxs)
    spec_aug_mask_idxs = np.broadcast_to(spec_aug_mask_idxs[:, :, None], (batch_size, max_num_masked_span, mask_length))
    spec_aug_mask_idxs = spec_aug_mask_idxs.reshape(batch_size, max_num_masked_span * mask_length)
    offsets = np.arange(mask_length)[None, None, :]
    offsets = np.broadcast_to(offsets, (batch_size, max_num_masked_span, mask_length)).reshape(batch_size, max_num_masked_span * mask_length)
    spec_aug_mask_idxs = spec_aug_mask_idxs + offsets
    if spec_aug_mask_idxs.max() > sequence_length - 1:
        spec_aug_mask_idxs[spec_aug_mask_idxs > sequence_length - 1] = sequence_length - 1
    np.put_along_axis(spec_aug_mask, spec_aug_mask_idxs, 1, -1)
    return spec_aug_mask

Computes random mask spans for a given shape. Used to implement [SpecAugment: A Simple Data Augmentation Method for
ASR](https://huggingface.co/papers/1904.08779). Note that this method is not optimized to run on TPU and should be run on
CPU as part of the preprocessing during training.

Args:
    shape: The shape for which to compute masks. This should be of a tuple of size 2 where
           the first element is the batch size and the second element is the length of the axis to span.
    mask_prob:  The percentage of the whole axis (between 0 and 1) which will be masked. The number of
                independently generated mask spans of length `mask_length` is computed by
                `mask_prob*shape[1]/mask_length`. Note that due to overlaps, `mask_prob` is an upper bound and the
                actual percentage will be smaller.
    mask_length: size of the mask
    min_masks: minimum number of masked spans
    attention_mask: A (right-padded) attention mask which independently shortens the feature axis of
                    each batch dimension.


def broadcast_recv_v2(shape, dtype, group_size, group_key, instance_key, communication_hint='auto', timeout=0):
    return gen_collective_ops.collective_bcast_recv_v2(T=dtype, group_size=group_size, group_key=group_key, instance_key=instance_key, shape=shape, communication_hint=communication_hint.lower(), timeout_seconds=timeout)

Receives a broadcasts tensor, across devices.

Args:
  shape: an int tensor.  Shape of the tensor to be received.
  dtype: Type of the tensor to be received.
  group_size: an int32 tensor.  One plus the number of receiving tensors, i.e.
      the total number of devices participating.  Each tensor must reside on a
      different device.
  group_key: an int32 tensor identifying the group of devices.
  instance_key: an int32 tensor identifying the participating group of Ops.
  communication_hint: preferred collective communication.  The implementation
    may fall back to another mechanism.  Options include `auto`, `ring`, and
    `nccl`.
  timeout: If set to a non zero, set a completion timeout to detect staleness.
    If the timer goes off, a DeadlineExceededError is raised.
    The timeout value in seconds. This feature is experimental.

Returns:
  An Op implementing the broadcast receive.


def crossed_column(keys, hash_bucket_size, hash_key=None):
    if not hash_bucket_size or hash_bucket_size < 1:
        raise ValueError('hash_bucket_size must be > 1. hash_bucket_size: {}'.format(hash_bucket_size))
    if not keys or len(keys) < 2:
        raise ValueError('keys must be a list with length > 1. Given: {}'.format(keys))
    for key in keys:
        if not isinstance(key, six.string_types) and (not isinstance(key, (CategoricalColumn, fc_old._CategoricalColumn))):
            raise ValueError('Unsupported key type. All keys must be either string, or categorical column except HashedCategoricalColumn. Given: {}'.format(key))
        if isinstance(key, (HashedCategoricalColumn, fc_old._HashedCategoricalColumn)):
            raise ValueError('categorical_column_with_hash_bucket is not supported for crossing. Hashing before crossing will increase probability of collision. Instead, use the feature name as a string. Given: {}'.format(key))
    return CrossedColumn(keys=tuple(keys), hash_bucket_size=hash_bucket_size, hash_key=hash_key)

Returns a column for performing crosses of categorical features.

Crossed features will be hashed according to `hash_bucket_size`. Conceptually,
the transformation can be thought of as:
  Hash(cartesian product of features) % `hash_bucket_size`

For example, if the input features are:

* SparseTensor referred by first key:

  ```python
  shape = [2, 2]
  {
      [0, 0]: "a"
      [1, 0]: "b"
      [1, 1]: "c"
  }
  ```

* SparseTensor referred by second key:

  ```python
  shape = [2, 1]
  {
      [0, 0]: "d"
      [1, 0]: "e"
  }
  ```

then crossed feature will look like:

```python
 shape = [2, 2]
{
    [0, 0]: Hash64("d", Hash64("a")) % hash_bucket_size
    [1, 0]: Hash64("e", Hash64("b")) % hash_bucket_size
    [1, 1]: Hash64("e", Hash64("c")) % hash_bucket_size
}
```

Here is an example to create a linear model with crosses of string features:

```python
keywords_x_doc_terms = crossed_column(['keywords', 'doc_terms'], 50K)
columns = [keywords_x_doc_terms, ...]
features = tf.io.parse_example(..., features=make_parse_example_spec(columns))
linear_prediction = linear_model(features, columns)
```

You could also use vocabulary lookup before crossing:

```python
keywords = categorical_column_with_vocabulary_file(
    'keywords', '/path/to/vocabulary/file', vocabulary_size=1K)
keywords_x_doc_terms = crossed_column([keywords, 'doc_terms'], 50K)
columns = [keywords_x_doc_terms, ...]
features = tf.io.parse_example(..., features=make_parse_example_spec(columns))
linear_prediction = linear_model(features, columns)
```

If an input feature is of numeric type, you can use
`categorical_column_with_identity`, or `bucketized_column`, as in the example:

```python
# vertical_id is an integer categorical feature.
vertical_id = categorical_column_with_identity('vertical_id', 10K)
price = numeric_column('price')
# bucketized_column converts numerical feature to a categorical one.
bucketized_price = bucketized_column(price, boundaries=[...])
vertical_id_x_price = crossed_column([vertical_id, bucketized_price], 50K)
columns = [vertical_id_x_price, ...]
features = tf.io.parse_example(..., features=make_parse_example_spec(columns))
linear_prediction = linear_model(features, columns)
```

To use crossed column in DNN model, you need to add it in an embedding column
as in this example:

```python
vertical_id_x_price = crossed_column([vertical_id, bucketized_price], 50K)
vertical_id_x_price_embedded = embedding_column(vertical_id_x_price, 10)
dense_tensor = input_layer(features, [vertical_id_x_price_embedded, ...])
```

Args:
  keys: An iterable identifying the features to be crossed. Each element can
    be either:
    * string: Will use the corresponding feature which must be of string type.
    * `CategoricalColumn`: Will use the transformed tensor produced by this
      column. Does not support hashed categorical column.
  hash_bucket_size: An int > 1. The number of buckets.
  hash_key: Specify the hash_key that will be used by the `FingerprintCat64`
    function to combine the crosses fingerprints on SparseCrossOp (optional).

Returns:
  A `CrossedColumn`.

Raises:
  ValueError: If `len(keys) < 2`.
  ValueError: If any of the keys is neither a string nor `CategoricalColumn`.
  ValueError: If any of the keys is `HashedCategoricalColumn`.
  ValueError: If `hash_bucket_size < 1`.


def HasCalls(self, *calls, **kwargs):
    if len(calls) == 1 and _IsIterable(calls[0]) and (not isinstance(calls[0], mock._Call)):
        calls = calls[0]
    contains_all = AssertThat(self._actual.mock_calls).ContainsAllIn(calls)
    any_order = kwargs.get('any_order')
    if any_order or any_order is None:
        return contains_all
    return contains_all.InOrder()

Assert that the mocked function was called with all the given calls.

Args:
  *calls: iterable of mock.call objects. Developers may also pass a single
      iterable of mock.call objects, for compatibility with mock's
      assert_has_calls() method, although this form is not preferred.
  **kwargs: optional parameters. The only recognized parameter is any_order:
      If any_order=True, the assertion succeeds if the mocked function was
          called with all the given calls, regardless of the call order.
      If any_order=False, the assertion succeeds if the mocked function was
          called with all of the given calls in the given order.
      If any_order is omitted, it behaves like any_order=True. This is the
          preferred way of calling HasCalls(). Developers who wish to
          enforce an order should call InOrder() on the returned predicate.
          If the order is unimportant, simply omit the InOrder() call.
          This is an intentional divergence from the mock library's syntax.

Returns:
  If any_order is True or omitted, and the mocked function was called all of
      with the expected calls, returns an _Ordered predicate on which
      .InOrder() can be subsequently called.
  If any_order=False, invokes the InOrder() predicate and returns its value.

Raises:
  TruthAssertionError: the mocked function is missing any of the expected
      calls.


def constant_to_value(self, pyval, subst=None, node=None):
    node = node or self.ctx.root_node
    if pyval.__class__ is tuple:
        type_key = tuple((type(v) for v in pyval))
    else:
        type_key = type(pyval)
    key = ('constant', pyval, type_key)
    if key in self._convert_cache:
        if self._convert_cache[key] is None:
            self._convert_cache[key] = self.unsolvable
            if not self.ctx.recursion_allowed:
                name = getattr(pyval, 'name', None) or pyval.__class__.__name__
                self.ctx.errorlog.recursion_error(self.ctx.vm.frames, name)
        return self._convert_cache[key]
    else:
        self._convert_cache[key] = None
        need_node = [False]

        def get_node():
            need_node[0] = True
            return node
        recursive = isinstance(pyval, pytd.LateType) and pyval.recursive
        if recursive:
            context = self.ctx.allow_recursive_convert()
        else:
            context = contextlib.nullcontext()
        with context:
            try:
                value = self._constant_to_value(pyval, subst, get_node)
            except NotImplementedError:
                del self._convert_cache[key]
                raise
        if not need_node[0] or node is self.ctx.root_node:
            if recursive:
                annot = abstract.LateAnnotation(pyval.name, self.ctx.vm.frames, self.ctx)
                annot.set_type(value)
                value = annot
            self._convert_cache[key] = value
        return value

Like constant_to_var, but convert to an abstract.BaseValue.

This also memoizes the results.  We don't memoize on name, as builtin types
like str or list might be reinitialized under different names (e.g. "param
1"), but we want the canonical name and type.  We *do* memoize on the type
as well, to make sure that e.g. "1.0" and "1" get converted to different
constants.  Memoization is an optimization, but an important one - mapping
constants like "None" to the same AbstractValue greatly simplifies the
cfg structures we're building.

Args:
  pyval: The constant to convert.
  subst: The current type parameters.
  node: The current CFG node. (For instances)

Returns:
  The converted constant. (Instance of BaseValue)


def __init__(self, chunk_id_fn: Optional[ChunkIdFn]=None):
    self.assign_chunk_id_fn = functools.partial(_assign_chunk_id, chunk_id_fn) if chunk_id_fn is not None else None

Base class for chunking transforms in RAG pipelines.

ChunkingTransformProvider defines the interface for splitting documents
into chunks for embedding and retrieval. Implementations should define how
to split content while preserving metadata and managing chunk IDs.

The transform flow:
- Takes input documents with content and metadata
- Splits content into chunks using implementation-specific logic
- Preserves document metadata in resulting chunks
- Optionally assigns unique IDs to chunks (configurable via chunk_id_fn

Example usage:
  >>> class MyChunker(ChunkingTransformProvider):
  ...     def get_splitter_transform(self):
  ...         return beam.ParDo(MySplitterDoFn())
  ...
  >>> chunker = MyChunker(chunk_id_fn=my_id_function)
  >>>
  >>> with beam.Pipeline() as p:
  ...     chunks = (
  ...         p
  ...         | beam.Create([{'text': 'document...', 'source': 'doc.txt'}])
  ...         | MLTransform(...).with_transform(chunker))

Args:
  chunk_id_fn: Optional function to generate chunk IDs. If not provided,
    random UUIDs will be used. Function should take a Chunk and return str.


class DefaultThrottler(PreCallThrottler):

    def __init__(self, window_ms: int=1, bucket_ms: int=1, overload_ratio: float=2, delay_secs: int=5):
        self.throttler = AdaptiveThrottler(window_ms=window_ms, bucket_ms=bucket_ms, overload_ratio=overload_ratio)
        self.delay_secs = delay_secs

Default throttler that uses
:class:`apache_beam.io.components.adaptive_throttler.AdaptiveThrottler`

Args:
  window_ms (int): length of history to consider, in ms, to set throttling.
  bucket_ms (int): granularity of time buckets that we store data in, in ms.
  overload_ratio (float): the target ratio between requests sent and
    successful requests. This is "K" in the formula in
    https://landing.google.com/sre/book/chapters/handling-overload.html.
  delay_secs (int): minimum number of seconds to throttle a request.


def to_json_file(self, json_file_path: Union[str, os.PathLike]):
    with open(json_file_path, 'w', encoding='utf-8') as writer:
        writer.write(self.to_json_string())

Save this instance to a JSON file.

Args:
    json_file_path (`str` or `os.PathLike`):
        Path to the JSON file in which this image_processor instance's parameters will be saved.


def get_check_numerics_error_message(slot, num_outputs, op_type, tensor, inputs, graph=None, traceback=None, stack_height_limit=30, path_length_limit=50):
    eager_vs_graph_qualifier = 'graph' if graph else 'eagerly-executing'
    message = '\n'
    message += '\n!!! Detected Infinity or NaN in output %d of %s op "%s" (# of outputs: %d) !!!\n' % (slot, eager_vs_graph_qualifier, op_type, num_outputs)
    message += '  dtype: %s\n' % tensor.dtype
    message += '  shape: %s\n' % (tensor.shape,)
    if not graph:
        is_inf = np.isinf(tensor)
        num_neg_inf = np.sum(np.logical_and(np.less(tensor, 0.0), is_inf))
        num_pos_inf = np.sum(np.logical_and(np.greater(tensor, 0.0), is_inf))
        num_nan = np.sum(np.isnan(tensor))
        if num_neg_inf > 0:
            message += '  # of -Inf elements: %s\n' % num_neg_inf
        if num_pos_inf > 0:
            message += '  # of +Inf elements: %s\n' % num_pos_inf
        if num_nan:
            message += '  # of +NaN elements: %s\n' % num_nan
    if len(inputs) > 1:
        message += '\n  Input tensors (%d):\n' % len(inputs)
        for slot, input_tensor in enumerate(inputs):
            message += '         %d: %s\n' % (slot, _maybe_lookup_original_input_tensor(graph, input_tensor))
    elif len(inputs) == 1:
        message += '\n  Input tensor: %s\n' % _maybe_lookup_original_input_tensor(graph, inputs[0])
    if graph and hasattr(graph, 'name') and graph.name:
        message += '  Graph name: "%s"\n' % graph.name
    if graph and traceback:
        message += '\n  Stack trace of op\'s creation ("->": inferred user code):\n'
        if stack_height_limit is not None and len(traceback) > stack_height_limit:
            num_omitted_frames = len(traceback) - stack_height_limit
            message += '    + ... (Omitted %d frames)\n' % num_omitted_frames
        for filepath, lineno, function_name, source_line in traceback[-stack_height_limit:]:
            user_code_indicator = '    '
            if not source_utils.guess_is_tensorflow_py_library(filepath):
                user_code_indicator = ' -> '
            message += '    + %s (L%d) %s\n' % (limit_string_length(filepath, path_length_limit), lineno, function_name)
            if source_line is not None:
                message += '%s|   %s\n' % (user_code_indicator, source_line)
    message += '\n'
    return message

Create a meaningful and user-friendly error message about offending tensor.

The error message reveals the following info about the op that outputs
NaN/Infinity: dtype, shape (to the extent known at graph-construction time),
input tensors, stack trace for op creation (if is graph mode).

Args:
  slot: (int) slot index of the tensor output.
  num_outputs: (int) total number of outputs of the op.
  op_type: (str) Type of the that generates `tensor`.
  tensor: (Tensor) the offending tensor, i.e., the tensor that contains
    Infinities or NaNs.
  inputs: (array of Tensor) inputs to the op that generates `tensor`.
  graph: (tf.Graph) the graph object that `tensor` belongs to. Available only
    under graph mode.
  traceback: (list of trace frames) the stack trace of the op's creation.
    Available only under graph model.
  stack_height_limit: (int or None) If int, limit to the height of the stack
    trace printed in the error message. If None, no limit to the height.
  path_length_limit: (int or None) Length limit for file paths included in the
    formatted stack trace.

Returns:
  (str) A formatted error message.


def __call__(self, images: Optional[ImageInput]=None, text: Optional[Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]]]=None, audio=None, videos=None, **kwargs: Unpack[Llama4ProcessorKwargs]) -> BatchFeature:
    if text is None:
        raise ValueError('You have to specify text.')
    output_kwargs = self._merge_kwargs(Llama4ProcessorKwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, **kwargs)
    if not isinstance(text, (list, tuple)):
        text = [text]
    image_inputs = {}
    if images is not None:
        images = make_flat_list_of_images(images)
        image_inputs = self.image_processor(images=images, **output_kwargs['images_kwargs'])
        image_height, image_width = image_inputs['pixel_values'][0].shape[-2:]
        num_patches_per_chunk = int(image_height // self.patch_size * (image_width // self.patch_size) // self.downsample_ratio)
        aspect_ratios = image_inputs.pop('aspect_ratios')
        total_placeholders = sum((prompt.count(self.fake_image_token) for prompt in text))
        if total_placeholders != len(images):
            raise ValueError(f'Found {total_placeholders} placeholders across the batch, but have {len(images)} flattened images.')
        image_index = 0
        processed_text = []
        for prompt in text:
            placeholder_count = prompt.count(self.fake_image_token)
            if placeholder_count == 0:
                processed_text.append(prompt)
                continue
            prompt_splits = prompt.split(self.fake_image_token)
            new_prompt = []
            for local_image_index, split_part in enumerate(prompt_splits):
                new_prompt.append(split_part)
                if local_image_index < placeholder_count:
                    tokens_for_this_image = self._prompt_split_image(aspect_ratios[image_index], num_patches_per_chunk)
                    image_index += 1
                    new_prompt.append(tokens_for_this_image)
            processed_text.append(''.join(new_prompt))
        if image_index != len(images):
            raise ValueError('Number of image placeholders in the prompt does not match the number of images.')
        text = processed_text
    return_tensors = output_kwargs['text_kwargs'].pop('return_tensors', None)
    text_inputs = self.tokenizer(text, **output_kwargs['text_kwargs'])
    self._check_special_mm_tokens(text, text_inputs, modalities=['image'])
    return BatchFeature(data={**text_inputs, **image_inputs}, tensor_type=return_tensors)

Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
and `kwargs` arguments to PreTrainedTokenizerFast's [`~PreTrainedTokenizerFast.__call__`] to encode the text.
To prepare the vision inputs, this method forwards the `images` and `kwargs` arguments to
Llama4ImageProcessor's [`~Llama4ImageProcessor.__call__`] if `images` is not `None`.

Args:
    images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`):
        The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
        tensor. Both channels-first and channels-last formats are supported.
    text (`str`, `List[str]`, `List[List[str]]`):
        The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
        (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
        `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
    return_tensors (`str` or [`~utils.TensorType`], *optional*):
        If set, will return tensors of a particular framework. Acceptable values are:
        - `'tf'`: Return TensorFlow `tf.constant` objects.
        - `'pt'`: Return PyTorch `torch.Tensor` objects.
        - `'np'`: Return NumPy `np.ndarray` objects.
        - `'jax'`: Return JAX `jnp.ndarray` objects.

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

    - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
    - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
      `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
      `None`).
    - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.


def start_services(self, service_alises):
    for name in service_alises:
        if name not in self._service_objects:
            raise Error(self._device, 'No service is registered under the name "%s", cannot start.' % name)
        service = self._service_objects[name]
        if not service.is_alive:
            service.start()

Starts the specified services.

Services will be started in the order specified by the input list.
No-op for services that are already running.

Args:
  service_alises: list of strings, the aliases of services to start.


def create(cls, num_quantiles, epsilon=None, max_num_elements=None, key=None, reverse=False, weighted=False, input_batched=False):
    max_num_elements = max_num_elements or cls._MAX_NUM_ELEMENTS
    if not epsilon:
        epsilon = min(0.01, 1.0 / num_quantiles) if weighted else 1.0 / num_quantiles
    b = 2
    while (b - 2) * (1 << b - 2) < epsilon * max_num_elements:
        b = b + 1
    b = b - 1
    k = max(2, int(math.ceil(max_num_elements / float(1 << b - 1))))
    return cls(num_quantiles=num_quantiles, buffer_size=k, num_buffers=b, key=key, reverse=reverse, weighted=weighted, input_batched=input_batched)

Creates an approximate quantiles combiner with the given key and desired
number of quantiles.

Args:
  num_quantiles: Number of quantiles to produce. It is the size of the
    final output list, including the mininum and maximum value items.
  epsilon: (optional) The default error bound is `epsilon`, which holds as
    long as the number of elements is less than `_MAX_NUM_ELEMENTS`.
    Specifically, if one considers the input as a sorted list x_1, ...,
    x_N, then the distance between each exact quantile x_c and its
    approximation x_c' is bounded by `|c - c'| < epsilon * N`. Note that
    these errors are worst-case scenarios. In practice the accuracy tends
    to be much better.
  max_num_elements: (optional) The cost (in time and space) to compute
    quantiles to a given accuracy is a function of the total number of
    elements in the data set.
  key: (optional) Key is a mapping of elements to a comparable key, similar
    to the key argument of Python's sorting methods.
  reverse: (optional) whether to order things smallest to largest, rather
    than largest to smallest.
  weighted: (optional) if set to True, the combiner produces weighted
    quantiles. The input elements are then expected to be tuples of values
    with the corresponding weight.
  input_batched: (optional) if set to True, inputs are expected to be
    batches of elements.


class ProtoAssertions(object):

    def assertProtoEqual(self, *args, **kwargs):
        return assertProtoEqual(self, *args, **kwargs)

Mix this into a googletest.TestCase class to get proto2 assertions.

Usage:

class SomeTestCase(compare.ProtoAssertions, googletest.TestCase):
  ...
  def testSomething(self):
    ...
    self.assertProtoEqual(a, b)

See module-level definitions for method documentation.


def _has_own(self, node, cls, method):
    assert method in ('__new__', '__init__')
    if not isinstance(cls, abstract.Class):
        return False
    self.load_lazy_attribute(method)
    obj_method = self.members[method]
    _, cls_method = self.ctx.attribute_handler.get_attribute(node, cls, method)
    return obj_method.data != cls_method.data

Whether a class has its own implementation of a particular method.

Args:
  node: The current node.
  cls: An abstract.Class.
  method: The method name. So that we don't have to handle the cases when
    the method doesn't exist, we only support "__new__" and "__init__".

Returns:
  True if the class's definition of the method is different from the
  definition in builtins.object, False otherwise.


def _generate_enqueue_op(self, flat_inputs: List[internal_types.NativeObject], flat_weights: List[Optional[internal_types.NativeObject]], flat_features: List[tpu_embedding_v2_utils.FeatureConfig], device_ordinal: int, mode_override: Text) -> ops.Operation:
    combiners = [table.combiner for table in self._table_config]
    indices_or_row_splits = []
    values = []
    weights = []
    int_zeros = array_ops.zeros((0,), dtype=dtypes.int32)
    float_zeros = array_ops.zeros((0,), dtype=dtypes.float32)
    for inp, weight, (path, feature) in zip(flat_inputs, flat_weights, flat_features):
        if isinstance(inp, tensor_lib.Tensor):
            self._add_data_for_tensor(inp, weight, indices_or_row_splits, values, weights, int_zeros, float_zeros, path)
        elif isinstance(inp, sparse_tensor.SparseTensor):
            self._add_data_for_sparse_tensor(inp, weight, indices_or_row_splits, values, weights, int_zeros, float_zeros, path, feature)
        elif isinstance(inp, ragged_tensor.RaggedTensor):
            self._add_data_for_ragged_tensor(inp, weight, indices_or_row_splits, values, weights, int_zeros, float_zeros, path, feature)
        else:
            raise ValueError('Input {} is of unknown type {}. Please only pass Tensor, SparseTensor or RaggedTensor as input to enqueue.'.format(path, type(inp)))
    return tpu_ops.enqueue_tpu_embedding_arbitrary_tensor_batch(sample_indices_or_row_splits=indices_or_row_splits, embedding_indices=values, aggregation_weights=weights, mode_override=mode_override, device_ordinal=device_ordinal, combiners=combiners)

Outputs a the enqueue op given the inputs and weights.

Args:
  flat_inputs: A list of input tensors.
  flat_weights: A list of input weights (or None) of the same length as
    flat_inputs.
  flat_features: A list of FeatureConfigs of the same length as flat_inputs.
  device_ordinal: The device to create the enqueue op for.
  mode_override: A tensor containing the string "train" or "inference".

Returns:
  The enqueue op.


def __mod__(self, other):
    other = as_dimension(other)
    if self._value is None or other.value is None:
        return Dimension(None)
    else:
        return Dimension(self._value % other.value)

Returns `self` modulo `other`.

Dimension modulo are computed as follows:

```python
tf.compat.v1.Dimension(m)    % tf.compat.v1.Dimension(n)     ==
tf.compat.v1.Dimension(m % n)
tf.compat.v1.Dimension(m)    % tf.compat.v1.Dimension(None)  # equiv. to
tf.compat.v1.Dimension(None)
tf.compat.v1.Dimension(None) % tf.compat.v1.Dimension(n)     # equiv. to
tf.compat.v1.Dimension(None)
tf.compat.v1.Dimension(None) % tf.compat.v1.Dimension(None)  # equiv. to
tf.compat.v1.Dimension(None)
```

Args:
  other: Another Dimension, or a value accepted by `as_dimension`.

Returns:
  A Dimension whose value is `self` modulo `other`.


class custom_gradient:

    def __init__(self, fun):
        warnings.warn('`custom_gradient` for the numpy backend acts as a pass-through to support the forward pass. No gradient computation or modification takes place.')
        self.fun = fun

    def __call__(self, *args, **kwargs):
        outputs, _ = self.fun(*args, **kwargs)
        return outputs

Decorator for custom gradients.

Args:
    fun: Forward pass function.


def coerce(lhs: FhirPathDataType, rhs: FhirPathDataType) -> FhirPathDataType:
    if not is_coercible(lhs, rhs):
        raise TypeError(f'Unsupported Standard SQL coercion between {lhs} and {rhs}.')
    if isinstance(rhs, _Any) or isinstance(lhs, _Any):
        return _Any
    if rhs in lhs.supported_coercion:
        return rhs
    else:
        return lhs

Performs implicit type coercion between two datatypes.

See more at: https://hl7.org/fhirpath/#conversion.

Args:
  lhs: The left operand.
  rhs: The right operand.

Returns:
  The resulting coerced datatype, if successful.

Raises:
  TypeError: In the event that coercion is not supported.
  ValueError: In the event that a coercion cycle is detected.


class Idefics2PerceiverConfig(PretrainedConfig):
    model_type = 'idefics2_perceiver'

    def __init__(self, hidden_act='silu', hidden_size=4096, rms_norm_eps=1e-06, resampler_n_latents=64, resampler_depth=3, resampler_n_heads=16, resampler_head_dim=96, num_key_value_heads=4, attention_dropout=0.0, initializer_range=0.02, **kwargs):
        self.hidden_act = hidden_act
        self.hidden_size = hidden_size
        self.rms_norm_eps = rms_norm_eps
        self.resampler_n_latents = resampler_n_latents
        self.resampler_depth = resampler_depth
        self.resampler_n_heads = resampler_n_heads
        self.num_key_value_heads = num_key_value_heads
        self.resampler_head_dim = resampler_head_dim
        self.attention_dropout = attention_dropout
        self.initializer_range = initializer_range
        if self.num_key_value_heads > self.resampler_n_heads:
            raise ValueError(f'num_key_value_heads={self.num_key_value_heads} must be less than or equal to resampler_n_heads={self.resampler_n_heads}')
        super().__init__(**kwargs)

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

Args:
    hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
        The non-linear activation function (function or string) in the perceiver block.
    hidden_size (`int`, *optional*, defaults to 4096):
        Dimension of the hidden representations.
    rms_norm_eps (`float`, *optional*, defaults to 1e-06):
        The epsilon used by the rms normalization layers.
    resampler_n_latents (`int`, *optional*, defaults to 64):
        Number of latent embeddings to resample ("compress") the input sequence to (usually < 128).
    resampler_depth (`int`, *optional*, defaults to 3):
        Depth of the Perceiver Resampler (Transformer w/ cross attention). Should be shallow (<= 3).
    resampler_n_heads (`int`, *optional*, defaults to 16):
        Number of heads in each Transformer block (for multi-headed self-attention).
    resampler_head_dim (`int`, *optional*, defaults to 96):
        Dimensionality of each head projection in the Transformer block.
    num_key_value_heads (`int`, *optional*, defaults to 4):
        Number of key-value heads in the perceiver attention block.
    attention_dropout (`float`, *optional*, defaults to 0.0):
        The dropout ratio for the attention probabilities.
    initializer_range (`float`, *optional*, defaults to 0.02):
        The standard deviation for initializing all weight matrices in the model.


def _StopOps(from_ops: list[ops.Operation], stop_gradient_ops: list[ops.Operation], pending_count, xs_set):
    stop_ops = set()
    for op in from_ops:
        is_stop_op = True
        for inp in _NonEagerInputs(op, xs_set):
            if pending_count[inp.op] > 0:
                is_stop_op = False
                break
        if is_stop_op:
            stop_ops.add(op)
    stop_ops.update((op for op in stop_gradient_ops))
    return stop_ops

The set of ops that terminate the gradient computation.

This computes the frontier of the forward graph *before* which backprop
should stop. Operations in the returned set will not be differentiated.
This set is defined as the subset of `from_ops` containing ops that have
no predecessor in `from_ops`. `pending_count` is the result of
`_PendingCount(xs, from_ops)`. An 'op' has predecessors in `from_ops`
iff pending_count[op] > 0.

In addition, none of `stop_gradient_ops` will be differentiated.

Args:
  from_ops: list of Operations.
  stop_gradient_ops: list of Operations never to backprop through.
  pending_count: mapping from operation to number of backprop inputs.
  xs_set: ObjectIdentitySet of Tensors.

Returns:
  The set of operations.


def _assert_sparse_compatible(sparse_tensors):
    checks = []
    first = sparse_tensors[0]
    for t in sparse_tensors[1:]:
        checks.append(check_ops.assert_equal(first.dense_shape, t.dense_shape, message='Mismatched shapes!'))
        checks.append(check_ops.assert_equal(first.indices, t.indices, message='Mismatched indices!'))
    return checks

Check that all of `sparse_tensors` have same `indices` and `dense_shape`.

Args:
  sparse_tensors: A list of sparse tensors.

Returns:
  An op to be used as a control dependency.


def gelu(features, approximate=False, name=None):
    with ops.name_scope(name, 'Gelu', [features]):
        features = ops.convert_to_tensor(features, name='features')
        if not features.dtype.is_floating:
            raise ValueError(f'`features.dtype` must be a floating point tensor.Received:features.dtype={features.dtype}')
        if approximate:
            coeff = math_ops.cast(0.044715, features.dtype)
            return 0.5 * features * (1.0 + math_ops.tanh(0.7978845608028654 * (features + coeff * math_ops.pow(features, 3))))
        else:
            return 0.5 * features * math_ops.erfc(-features * math_ops.cast(0.7071067811865476, features.dtype))

Compute the Gaussian Error Linear Unit (GELU) activation function.

Gaussian error linear unit (GELU) computes
`x * P(X <= x)`, where `P(X) ~ N(0, 1)`.
The (GELU) nonlinearity weights inputs by their value, rather than gates
inputs by their sign as in ReLU.

For example:

>>> x = tf.constant([-3.0, -1.0, 0.0, 1.0, 3.0], dtype=tf.float32)
>>> y = tf.nn.gelu(x)
>>> y.numpy()
array([-0.00404951, -0.15865529,  0.        ,  0.8413447 ,  2.9959507 ],
    dtype=float32)
>>> y = tf.nn.gelu(x, approximate=True)
>>> y.numpy()
array([-0.00363752, -0.15880796,  0.        ,  0.841192  ,  2.9963627 ],
    dtype=float32)

Args:
  features: A `float Tensor` representing preactivation values.
  approximate: An optional `bool`. Defaults to `False`. Whether to enable
    approximation.
  name: A name for the operation (optional).

Returns:
  A `Tensor` with the same type as `features`.

Raises:
  ValueError: if `features` is not a floating point `Tensor`.

References:
  [Gaussian Error Linear Units (GELUs)](https://arxiv.org/abs/1606.08415).


def set_cpu0(device_string):
    if context.is_custom_device(device_string):
        return device_string
    parsed_device = pydev.DeviceSpec.from_string(device_string)
    parsed_device = parsed_device.replace(device_type='CPU', device_index=0)
    return parsed_device.to_string()

Creates a new device string based on `device_string` but using /CPU:0.

If the device is already on /CPU:0 or it is a custom device, this is a no-op.

Args:
  device_string: A device string.

Returns:
  A device string.


def handle_http_error(error: HTTPException) -> ResponseReturnValue:
    code = error.code or 500
    return (DQMResponse(name=error.name, description=error.description, code=code), code)

DQM HTTP Error Response.

Args:
    * error: HTTP error

Returns:
    * DQMResponse for the error with the relevant status code


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

Creates a new IntGaugeCell.

Args:
  cell: A c pointer of TFE_MonitoringIntGaugeCell.


def traverse(root, visit):
    _traverse_internal(root, visit, [], '')

Recursively enumerate all members of `root`.

Similar to the Python library function `os.path.walk`.

Traverses the tree of Python objects starting with `root`, depth first.
Parent-child relationships in the tree are defined by membership in modules or
classes. The function `visit` is called with arguments
`(path, parent, children)` for each module or class `parent` found in the tree
of python objects starting with `root`. `path` is a string containing the name
with which `parent` is reachable from the current context. For example, if
`root` is a local class called `X` which contains a class `Y`, `visit` will be
called with `('Y', X.Y, children)`).

If `root` is not a module or class, `visit` is never called. `traverse`
never descends into built-in modules.

`children`, a list of `(name, object)` pairs are determined by
`tf_inspect.getmembers`. To avoid visiting parts of the tree, `children` can
be modified in place, using `del` or slice assignment.

Cycles (determined by reference equality, `is`) stop the traversal. A stack of
objects is kept to find cycles. Objects forming cycles may appear in
`children`, but `visit` will not be called with any object as `parent` which
is already in the stack.

Traversing system modules can take a long time, it is advisable to pass a
`visit` callable which denylists such modules.

Args:
  root: A python object with which to start the traversal.
  visit: A function taking arguments `(path, parent, children)`. Will be
    called for each object found in the traversal.


def Visit(self, visitor, *args, **kwargs):
    return _Visit(self, visitor, *args, **kwargs)

Visitor interface for transforming a tree of nodes to a new tree.

You can pass a visitor, and callback functions on that visitor will be
called for all nodes in the tree. Note that nodes are also allowed to
be stored in lists and as the values of dictionaries, as long as these
lists/dictionaries are stored in the named fields of the Node class.
It's possible to overload the Visit function on Nodes, to do your own
processing.

Arguments:
  visitor: An instance of a visitor for this tree. For every node type you
    want to transform, this visitor implements a "Visit<Classname>"
    function named after the class of the node this function should
    target. Note that <Classname> is the *actual* class of the node, so
    if you subclass a Node class, visitors for the superclasses will *not*
    be triggered anymore. Also, visitor callbacks are only triggered
    for subclasses of Node.
  *args: Passed to the visitor callback.
  **kwargs: Passed to the visitor callback.

Returns:
  Transformed version of this node.


def _compute_router_probabilities(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    self.input_dtype = hidden_states.dtype
    hidden_states = hidden_states.to(self.dtype)
    if self.training and self.jitter_noise > 0:
        hidden_states *= torch.empty_like(hidden_states).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise)
    self._cast_classifier()
    router_logits = self.classifier(hidden_states)
    router_probabilities = nn.functional.softmax(router_logits, dim=-1, dtype=self.dtype).to(self.input_dtype)
    return (router_probabilities, router_logits)

Computes router probabilities from input hidden states.

Args:
    hidden_states (`torch.Tensor`):
        (batch_size, sequence_length, hidden_dim) from which router probabilities are computed.
Returns:
    router_probabilities (`torch.Tensor`):
        Tensor of shape (batch_size, sequence_length, num_experts) corresponding to the probabilities for each
        token and expert. Used for routing tokens to experts.
    router_logits (`torch.Tensor`):
        Logits tensor of shape (batch_size, sequence_length, num_experts) corresponding to raw router logits.
        This is used later for computing router z-loss.


def pack(self, tensors: Sequence[Any], layout: layout_lib.Layout) -> Any:
    if not context.executing_eagerly():
        raise RuntimeError('`pack` must be called eagerly.')
    self._register_mesh(layout.mesh)
    with ops.device(self.name):
        if all((isinstance(t, sparse_tensor.SparseTensor) for t in tensors)):
            if not all((t.shape == tensors[0].shape for t in tensors)):
                raise TypeError('All input SparseTensors to Pack must be same shape.')
            is_sparse = True
            tensors = [t.indices for t in tensors] + [t.values for t in tensors] + [ops.convert_to_tensor(t.shape, dtype=dtypes.int64) for t in tensors]
        elif any((isinstance(t, sparse_tensor.SparseTensor) for t in tensors)):
            raise TypeError('Cannot Pack SparseTensors with Tensors.')
        else:
            is_sparse = False
        try:
            return _pywrap_dtensor_device.Pack(context.context()._handle, tensors, layout.to_string(), self._device_info, is_sparse)
        except core._NotOkStatusException as e:
            raise core._status_to_exception(e) from None

Packs tensors into a DTensor handle on this DTensor device.

Packing and unpacking are inverse operations:

```
* unpack(pack(tensors)) == tensors
* pack(unpack(dtensor)) == dtensor
```

Refer to `dtensor.pack` for more information.

Args:
  tensors: The list of tensors to pack into a DTensor.
  layout: The layout of the DTensor to be created.

Returns:
  A DTensor created from the individual component tensors.

Raises:
  RuntimeError: When not called eagerly.


def to_json_string(self, use_diff: bool=True, ignore_metadata: bool=False) -> str:
    if use_diff is True:
        config_dict = self.to_diff_dict()
    else:
        config_dict = self.to_dict()
    if ignore_metadata:
        for metadata_field in METADATA_FIELDS:
            config_dict.pop(metadata_field, None)

    def convert_keys_to_string(obj):
        if isinstance(obj, dict):
            return {str(key): convert_keys_to_string(value) for key, value in obj.items()}
        elif isinstance(obj, list):
            return [convert_keys_to_string(item) for item in obj]
        else:
            return obj

    def convert_dataclass_to_dict(obj):
        if isinstance(obj, dict):
            return {key: convert_dataclass_to_dict(value) for key, value in obj.items()}
        elif is_dataclass(obj):
            return obj.to_dict()
        else:
            return obj
    config_dict = convert_keys_to_string(config_dict)
    config_dict = convert_dataclass_to_dict(config_dict)
    return json.dumps(config_dict, indent=2, sort_keys=True) + '\n'

Serializes this instance to a JSON string.

Args:
    use_diff (`bool`, *optional*, defaults to `True`):
        If set to `True`, only the difference between the config instance and the default `GenerationConfig()`
        is serialized to JSON string.
    ignore_metadata (`bool`, *optional*, defaults to `False`):
        Whether to ignore the metadata fields present in the instance

Returns:
    `str`: String containing all the attributes that make up this configuration instance in JSON format.


def WriteModifyTimestamp(self, timestamp):
    if timestamp is None:
        return True
    self.modify_time = None
    return self._WriteTimestamp(timestamp, self.modify_file)

Convenience method for writing the last modify timestamp.

Args:
  timestamp:  An int with the number of seconds since epoch.
    If timestamp is None, performs no action.

Returns:
  A boolean indicating success of the write.


class SquadResult:

    def __init__(self, unique_id, start_logits, end_logits, start_top_index=None, end_top_index=None, cls_logits=None):
        self.start_logits = start_logits
        self.end_logits = end_logits
        self.unique_id = unique_id
        if start_top_index:
            self.start_top_index = start_top_index
            self.end_top_index = end_top_index
            self.cls_logits = cls_logits

Constructs a SquadResult which can be used to evaluate a model's output on the SQuAD dataset.

Args:
    unique_id: The unique identifier corresponding to that example.
    start_logits: The logits corresponding to the start of the answer
    end_logits: The logits corresponding to the end of the answer


class GlobalAveragePooling1D(BaseGlobalPooling):

    def __init__(self, data_format=None, keepdims=False, **kwargs):
        super().__init__(pool_dimensions=1, data_format=data_format, keepdims=keepdims, **kwargs)
        self.supports_masking = True

    def call(self, inputs, mask=None):
        steps_axis = 1 if self.data_format == 'channels_last' else 2
        if mask is not None:
            mask = backend.cast(mask, inputs[0].dtype)
            mask = ops.expand_dims(mask, 2 if self.data_format == 'channels_last' else 1)
            inputs *= mask
            return ops.sum(inputs, axis=steps_axis, keepdims=self.keepdims) / ops.sum(mask, axis=steps_axis, keepdims=self.keepdims)
        else:
            return ops.mean(inputs, axis=steps_axis, keepdims=self.keepdims)

    def compute_mask(self, inputs, mask=None):
        return None

Global average pooling operation for temporal data.

Args:
    data_format: string, either `"channels_last"` or `"channels_first"`.
        The ordering of the dimensions in the inputs. `"channels_last"`
        corresponds to inputs with shape `(batch, steps, features)`
        while `"channels_first"` corresponds to inputs with shape
        `(batch, features, steps)`. It defaults to the `image_data_format`
        value found in your Keras config file at `~/.keras/keras.json`.
        If you never set it, then it will be `"channels_last"`.
    keepdims: A boolean, whether to keep the temporal dimension or not.
        If `keepdims` is `False` (default), the rank of the tensor is
        reduced for spatial dimensions. If `keepdims` is `True`, the
        temporal dimension are retained with length 1.
        The behavior is the same as for `tf.reduce_mean` or `np.mean`.

Call arguments:
    inputs: A 3D tensor.
    mask: Binary tensor of shape `(batch_size, steps)` indicating whether
        a given step should be masked (excluded from the average).

Input shape:

- If `data_format='channels_last'`:
    3D tensor with shape:
    `(batch_size, steps, features)`
- If `data_format='channels_first'`:
    3D tensor with shape:
    `(batch_size, features, steps)`

Output shape:

- If `keepdims=False`:
    2D tensor with shape `(batch_size, features)`.
- If `keepdims=True`:
    - If `data_format="channels_last"`:
        3D tensor with shape `(batch_size, 1, features)`
    - If `data_format="channels_first"`:
        3D tensor with shape `(batch_size, features, 1)`

Example:

>>> x = np.random.rand(2, 3, 4)
>>> y = keras.layers.GlobalAveragePooling1D()(x)
>>> y.shape
(2, 4)


def data_gen_sig1() -> repr_dataset.RepresentativeDataset:
    for _ in range(4):
        yield {'matmul_input': random_ops.random_uniform(shape=(1, 4))}

Generates tuple-style samples for signature 'sig1'.

The first element of the tuple identifies the signature key the input data
is for.

Yields:
  Representative sample for 'sig1'.