def Print(input_, data, message=None, first_n=None, summarize=None, name=None):
    return gen_logging_ops._print(input_, data, message, first_n, summarize, name)

Prints a list of tensors.

    This is an identity op (behaves like `tf.identity`) with the side effect
    of printing `data` when evaluating.

    Note: This op prints to the standard error. It is not currently compatible
        with jupyter notebook (printing to the notebook *server's* output, not into
        the notebook).

    @compatibility(TF2)
    This API is deprecated. Use `tf.print` instead. `tf.print` does not need the
    `input_` argument.

    `tf.print` works in TF2 when executing eagerly and inside a `tf.function`.

    In TF1-styled sessions, an explicit control dependency declaration is needed
    to execute the `tf.print` operation. Refer to the documentation of
    `tf.print` for more details.
    @end_compatibility

Args:
        input_: A tensor passed through this op.
        data: A list of tensors to print out when op is evaluated.
        message: A string, prefix of the error message.
        first_n: Only log `first_n` number of times. Negative numbers log always;
        this is the default.
        summarize: Only print this many entries of each tensor. If None, then a
        maximum of 3 elements are printed per input tensor.
        name: A name for the operation (optional).

Returns:
        A `Tensor`. Has the same type and contents as `input_`.

        ```python
        sess = tf.compat.v1.Session()
        with sess.as_default():
        tensor = tf.range(10)
        print_op = tf.print(tensor)
        with tf.control_dependencies([print_op]):
            out = tf.add(tensor, tensor)
        sess.run(out)
        ```


def ee_initialize(use_personal_account: bool=False, enforce_high_volume: bool=False, service_account: t.Optional[str]=None, private_key: t.Optional[str]=None, project_id: t.Optional[str]=None) -> None:
    creds = get_creds(use_personal_account, service_account, private_key)
    on_compute_engine = is_compute_engine()
    if on_compute_engine:
        if project_id is None and use_personal_account:
            raise RuntimeError('Project_name should not be None!')
        params = {'credentials': creds, 'opt_url': 'https://earthengine-highvolume.googleapis.com'}
        if project_id:
            params['project'] = project_id
        ee.Initialize(**params)
    elif enforce_high_volume and (not on_compute_engine):
        raise RuntimeError('Must run on a compute engine VM to use the high volume earth engine api.')
    else:
        ee.Initialize(creds)

Initializes earth engine with the high volume API when using a compute engine VM.

Args:
        use_personal_account: A flag to use personal account for ee authentication. Default: False.
        enforce_high_volume: A flag to use the high volume API when using a compute engine VM. Default: False.
        service_account: Service account address when using a private key for earth engine authentication.
        private_key: A private key path to authenticate earth engine using private key. Default: None.
        Project ID: An identifier that represents the name of a project present in Earth Engine.

Raises:
        RuntimeError: Earth Engine did not initialize.


def SendCapture(self, request, global_params=None):
    config = self.GetMethodConfig('SendCapture')
    return self._RunMethod(config, request, global_params=global_params)

Send encoded debug capture data for component.

Args:
        request: (DataflowProjectsLocationsJobsDebugSendCaptureRequest) input message
        global_params: (StandardQueryParameters, default: None) global arguments

Returns:
        (SendDebugCaptureResponse) The response message.


def __init__(self, name='mean', dtype=None):
    super().__init__(name=name, dtype=dtype)
    self.total = self.add_variable(shape=(), initializer=initializers.Zeros(), dtype=self.dtype, name='total')
    self.count = self.add_variable(shape=(), initializer=initializers.Zeros(), dtype=self.dtype, name='count')

Compute the (weighted) mean of the given values.

    For example, if values is `[1, 3, 5, 7]` then the mean is 4.
    If `sample_weight` was specified as `[1, 1, 0, 0]` then the mean would be 2.

    This metric creates two variables, `total` and `count`.
    The mean value returned is simply `total` divided by `count`.

Args:
        name: (Optional) string name of the metric instance.
        dtype: (Optional) data type of the metric result.

Example:
    >>> m = Mean()
    >>> m.update_state([1, 3, 5, 7])
    >>> m.result()
    4.0

    >>> m.reset_state()
    >>> m.update_state([1, 3, 5, 7], sample_weight=[1, 1, 0, 0])
    >>> m.result()
    2.0


def __init__(self, target_shape, **kwargs):
    super().__init__(**kwargs)
    self.target_shape = tuple(target_shape)

Layer that reshapes inputs into the given shape.

Args:
        target_shape: Target shape. Tuple of integers, does not include the
            samples dimension (batch size).

    Input shape:
        Arbitrary, although all dimensions in the input shape must be
        known/fixed. Use the keyword argument `input_shape` (tuple of integers,
        does not include the samples/batch size axis) when using this layer as
        the first layer in a model.

    Output shape:
        `(batch_size, *target_shape)`

Example:
    >>> x = keras.Input(shape=(12,))
    >>> y = keras.layers.Reshape((3, 4))(x)
    >>> y.shape
    (None, 3, 4)

    >>> # also supports shape inference using `-1` as dimension
    >>> y = keras.layers.Reshape((-1, 2, 2))(x)
    >>> y.shape
    (None, 3, 2, 2)


def outer_graph(self):
    current = self._weak_outer_graph()
    if current is None:
        return self._fallback_outer_graph
    return current

The Graph this FuncGraph is nested in.

    Functions may capture Tensors from graphs they are nested in (transitive).

Returns:
        A Graph object. Initially set to the current default graph when the
        FuncGraph was created. If the previous `outer_graph` was deleted because
        the function that owns it was deleted, `outer_graph` is reset to the
        outermost default graph active when the FuncGraph was created. This
        FuncGraph won't have captured anything from the new `outer_graph` (and
        likely not from the previous setting, since that would have created a
        strong reference), but it is returned so that FuncGraphs always have a
        parent.


def forward(self, input_ids: Optional[torch.Tensor]=None, bbox: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None, bbox_first_token_mask: Optional[torch.Tensor]=None, token_type_ids: Optional[torch.Tensor]=None, position_ids: Optional[torch.Tensor]=None, head_mask: Optional[torch.Tensor]=None, inputs_embeds: Optional[torch.Tensor]=None, initial_token_labels: Optional[torch.Tensor]=None, subsequent_token_labels: Optional[torch.Tensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[Tuple[torch.Tensor], BrosSpadeOutput]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    outputs = self.bros(input_ids=input_ids, bbox=bbox, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    last_hidden_states = outputs[0]
    last_hidden_states = last_hidden_states.transpose(0, 1).contiguous()
    initial_token_logits = self.initial_token_classifier(last_hidden_states).transpose(0, 1).contiguous()
    subsequent_token_logits = self.subsequent_token_classifier(last_hidden_states, last_hidden_states).squeeze(0)
    inv_attention_mask = 1 - attention_mask
    batch_size, max_seq_length = inv_attention_mask.shape
    device = inv_attention_mask.device
    invalid_token_mask = torch.cat([inv_attention_mask, torch.zeros([batch_size, 1]).to(device)], axis=1).bool()
    subsequent_token_logits = subsequent_token_logits.masked_fill(invalid_token_mask[:, None, :], torch.finfo(subsequent_token_logits.dtype).min)
    self_token_mask = torch.eye(max_seq_length, max_seq_length + 1).to(device=device, dtype=torch.bool)
    subsequent_token_logits = subsequent_token_logits.masked_fill(self_token_mask[None, :, :], torch.finfo(subsequent_token_logits.dtype).min)
    subsequent_token_mask = attention_mask.view(-1).bool()
    loss = None
    if initial_token_labels is not None and subsequent_token_labels is not None:
        loss_fct = CrossEntropyLoss()
        initial_token_labels = initial_token_labels.view(-1)
        if bbox_first_token_mask is not None:
            bbox_first_token_mask = bbox_first_token_mask.view(-1)
            initial_token_loss = loss_fct(initial_token_logits.view(-1, self.num_labels)[bbox_first_token_mask], initial_token_labels[bbox_first_token_mask])
        else:
            initial_token_loss = loss_fct(initial_token_logits.view(-1, self.num_labels), initial_token_labels)
        subsequent_token_labels = subsequent_token_labels.view(-1)
        subsequent_token_loss = loss_fct(subsequent_token_logits.view(-1, max_seq_length + 1)[subsequent_token_mask], subsequent_token_labels[subsequent_token_mask])
        loss = initial_token_loss + subsequent_token_loss
    if not return_dict:
        output = (initial_token_logits, subsequent_token_logits) + outputs[2:]
        return (loss,) + output if loss is not None else output
    return BrosSpadeOutput(loss=loss, initial_token_logits=initial_token_logits, subsequent_token_logits=subsequent_token_logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions)

bbox ('torch.FloatTensor' of shape '(batch_size, num_boxes, 4)'):
    Bounding box coordinates for each token in the input sequence. Each bounding box is a list of four values
    (x1, y1, x2, y2), where (x1, y1) is the top left corner, and (x2, y2) is the bottom right corner of the
    bounding box.
    bbox_first_token_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
    Mask to indicate the first token of each bounding box. Mask values selected in `[0, 1]`:

    - 1 for tokens that are **not masked**,
    - 0 for tokens that are **masked**.
    initial_token_labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
    Labels for the initial token classification.
    subsequent_token_labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
    Labels for the subsequent token classification.

Example:
    ```python
    >>> import torch
    >>> from transformers import BrosProcessor, BrosSpadeEEForTokenClassification

    >>> processor = BrosProcessor.from_pretrained("jinho8345/bros-base-uncased")

    >>> model = BrosSpadeEEForTokenClassification.from_pretrained("jinho8345/bros-base-uncased")

    >>> encoding = processor("Hello, my dog is cute", add_special_tokens=False, return_tensors="pt")
    >>> bbox = torch.tensor([[[0, 0, 1, 1]]]).repeat(1, encoding["input_ids"].shape[-1], 1)
    >>> encoding["bbox"] = bbox

    >>> outputs = model(**encoding)
    ```


def in_top_k(predictions, targets, k):
    return nn.in_top_k(predictions, targets, k)

Returns whether the `targets` are in the top `k` `predictions`.

Args:
        predictions: A tensor of shape `(batch_size, classes)` and type `float32`.
        targets: A 1D tensor of length `batch_size` and type `int32` or `int64`.
        k: An `int`, number of top elements to consider.

Returns:
        A 1D tensor of length `batch_size` and type `bool`.
        `output[i]` is `True` if `predictions[i, targets[i]]` is within top-`k`
        values of `predictions[i]`.


def load_graph_from_args(pipeline_name: str, framework: str, model: str, tokenizer: Optional[str]=None, **models_kwargs) -> Pipeline:
    if tokenizer is None:
        tokenizer = model
    if framework == 'pt' and (not is_torch_available()):
        raise Exception('Cannot convert because PyTorch is not installed. Please install torch first.')
    if framework == 'tf' and (not is_tf_available()):
        raise Exception('Cannot convert because TF is not installed. Please install tensorflow first.')
    print(f'Loading pipeline (model: {model}, tokenizer: {tokenizer})')
    return pipeline(pipeline_name, model=model, tokenizer=tokenizer, framework=framework, model_kwargs=models_kwargs)

Convert the set of arguments provided through the CLI to an actual pipeline reference (tokenizer + model

Args:
        pipeline_name: The kind of pipeline to use (ner, question-answering, etc.)
        framework: The actual model to convert the pipeline from ("pt" or "tf")
        model: The model name which will be loaded by the pipeline
        tokenizer: The tokenizer name which will be loaded by the pipeline, default to the model's value

    Returns: Pipeline object


def _prop(self):
    return 'prop_with_doc'


def get(self):
    raise NotImplementedError

Creates a generator to extract data from the queue.

    Skip the data if it is `None`.

    This method is called from the main thread.

Yields:
        The next element in the queue, i.e. a tuple
        `(inputs, targets)` or
        `(inputs, targets, sample_weights)`.


def get_dense_tensor(self, transformation_cache, state_manager):
    if isinstance(self.categorical_column, SequenceCategoricalColumn):
        raise ValueError('In indicator_column: {}. categorical_column must not be of type SequenceCategoricalColumn. Suggested fix A: If you wish to use DenseFeatures, use a non-sequence categorical_column_with_*. Suggested fix B: If you wish to create sequence input, use SequenceFeatures instead of DenseFeatures. Given (type {}): {}'.format(self.name, type(self.categorical_column), self.categorical_column))
    return transformation_cache.get(self, state_manager)

Returns dense `Tensor` representing feature.

Args:
        transformation_cache: A `FeatureTransformationCache` object to access
        features.
        state_manager: A `StateManager` to create / access resources such as
        lookup tables.

Returns:
        Dense `Tensor` created within `transform_feature`.

Raises:
        ValueError: If `categorical_column` is a `SequenceCategoricalColumn`.


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

Return:

Example:
    ```python
    >>> import tensorflow as tf
    >>> from transformers import AutoTokenizer, TFMobileBertForPreTraining

    >>> tokenizer = AutoTokenizer.from_pretrained("google/mobilebert-uncased")
    >>> model = TFMobileBertForPreTraining.from_pretrained("google/mobilebert-uncased")
    >>> input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute"))[None, :]  # Batch size 1
    >>> outputs = model(input_ids)
    >>> prediction_scores, seq_relationship_scores = outputs[:2]
    ```


def track_and_add_endpoint(self, name, resource, input_signature, **kwargs):
    if name in self._endpoint_names:
        raise ValueError(f"Endpoint name '{name}' is already taken.")
    if not isinstance(resource, layers.Layer):
        raise ValueError(f'Invalid resource type. Expected an instance of a Keras `Layer` or `Model`. Received: resource={resource} (of type {type(resource)})')
    if not resource.built:
        raise ValueError('The layer provided has not yet been built. It must be built before export.')
    if backend.backend() != 'jax':
        if 'jax2tf_kwargs' in kwargs or 'is_static' in kwargs:
            raise ValueError(f"'jax2tf_kwargs' and 'is_static' are only supported with the jax backend. Current backend: {backend.backend()}")
    input_signature = tree.map_structure(make_tf_tensor_spec, input_signature)
    if not hasattr(BackendExportArchive, 'track_and_add_endpoint'):
        self.track(resource)
        return self.add_endpoint(name, resource.__call__, input_signature, **kwargs)
    else:
        decorated_fn = super().track_and_add_endpoint(name, resource, input_signature, **kwargs)
        self._endpoint_signatures[name] = input_signature
        setattr(self._tf_trackable, name, decorated_fn)
        self._endpoint_names.append(name)
        return decorated_fn

Track the variables and register a new serving endpoint.

    This function combines the functionality of `track` and `add_endpoint`.
    It tracks the variables of the `resource` (either a layer or a model)
    and registers a serving endpoint using `resource.__call__`.

Args:
        name: `str`. The name of the endpoint.
        resource: A trackable Keras resource, such as a layer or model.
        input_signature: Optional. Specifies the shape and dtype of `fn`.
            Can be a structure of `keras.InputSpec`, `tf.TensorSpec`,
            `backend.KerasTensor`, or backend tensor (see below for an
            example showing a `Functional` model with 2 input arguments). If
            not provided, `fn` must be a `tf.function` that has been called
            at least once. Defaults to `None`.
        **kwargs: Additional keyword arguments:
            - Specific to the JAX backend:
                - `is_static`: Optional `bool`. Indicates whether `fn` is
                    static. Set to `False` if `fn` involves state updates
                    (e.g., RNG seeds).
                - `jax2tf_kwargs`: Optional `dict`. Arguments for
                    `jax2tf.convert`. See [`jax2tf.convert`](
                        https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md).
                    If `native_serialization` and `polymorphic_shapes` are
                    not provided, they are automatically computed.


def __init__(self, sess, hooks):
    _WrappedSession.__init__(self, sess)
    self._hooks = hooks
    self._should_stop = False

Initializes a _HookedSession object.

Args:
        sess: A `tf.compat.v1.Session` or a `_WrappedSession` object.
        hooks: An iterable of `SessionRunHook' objects.


def explicit_pass(msg, extras=None):
    raise signals.TestPass(msg, extras)

Explicitly pass a test.

    This will pass the test explicitly regardless of any other error happened
    in the test body. E.g. even if errors have been recorded with `expects`,
    the test will still be marked pass if this is called.

    A test without uncaught exception will pass implicitly so this should be
    used scarcely.

Args:
        msg: A string explaining the details of the passed test.
        extras: An optional field for extra information to be included in
        test result.

Raises:
        signals.TestPass: Mark a test as passed.


def forward(self, hidden_states: torch.FloatTensor, p_mask: Optional[torch.FloatTensor]=None) -> torch.FloatTensor:
    x = self.dense(hidden_states).squeeze(-1)
    if p_mask is not None:
        if p_mask.dtype == torch.float16:
            x = x * (1 - p_mask) - 65500 * p_mask
        else:
            x = x * (1 - p_mask) - 1e+30 * p_mask
    return x

Args:
    hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`):
        The final hidden states of the model.
    p_mask (`torch.FloatTensor` of shape `(batch_size, seq_len)`, *optional*):
        Mask for tokens at invalid position, such as query and special symbols (PAD, SEP, CLS). 1.0 means token
        should be masked.

Returns:
    `torch.FloatTensor`: The start logits for SQuAD.


def _convert_tf1_model(flags):
    if flags.custom_opdefs:
        register_custom_opdefs(_parse_array(flags.custom_opdefs))
    converter = _get_tflite_converter(flags)
    if flags.inference_type:
        converter.inference_type = _parse_inference_type(flags.inference_type, 'inference_type')
    if flags.inference_input_type:
        converter.inference_input_type = _parse_inference_type(flags.inference_input_type, 'inference_input_type')
    if flags.output_format:
        converter.output_format = _toco_flags_pb2.FileFormat.Value(flags.output_format)
    if flags.mean_values and flags.std_dev_values:
        input_arrays = converter.get_input_arrays()
        std_dev_values = _parse_array(flags.std_dev_values, type_fn=float)
        if converter.inference_type == dtypes.float32:
            mean_values = _parse_array(flags.mean_values, type_fn=float)
        else:
            mean_values = _parse_array(flags.mean_values, type_fn=int)
        quant_stats = list(zip(mean_values, std_dev_values))
        if not flags.input_arrays and len(input_arrays) > 1 or len(input_arrays) != len(quant_stats):
            raise ValueError("Mismatching --input_arrays, --std_dev_values, and --mean_values. The flags must have the same number of items. The current input arrays are '{0}'. --input_arrays must be present when specifying --std_dev_values and --mean_values with multiple input tensors in order to map between names and values.".format(','.join(input_arrays)))
        converter.quantized_input_stats = dict(list(zip(input_arrays, quant_stats)))
    if flags.default_ranges_min is not None and flags.default_ranges_max is not None:
        converter.default_ranges_stats = (flags.default_ranges_min, flags.default_ranges_max)
    if flags.drop_control_dependency:
        converter.drop_control_dependency = flags.drop_control_dependency
    if flags.reorder_across_fake_quant:
        converter.reorder_across_fake_quant = flags.reorder_across_fake_quant
    if flags.change_concat_input_ranges:
        converter.change_concat_input_ranges = flags.change_concat_input_ranges == 'TRUE'
    if flags.allow_custom_ops:
        converter.allow_custom_ops = flags.allow_custom_ops
    if flags.target_ops:
        ops_set_options = lite.OpsSet.get_options()
        converter.target_spec.supported_ops = set()
        for option in flags.target_ops.split(','):
            if option not in ops_set_options:
                raise ValueError('Invalid value for --target_ops. Options: {0}'.format(','.join(ops_set_options)))
            converter.target_spec.supported_ops.add(lite.OpsSet(option))
    if flags.experimental_select_user_tf_ops:
        if lite.OpsSet.SELECT_TF_OPS not in converter.target_spec.supported_ops:
            raise ValueError('--experimental_select_user_tf_ops can only be set if --target_ops contains SELECT_TF_OPS.')
        user_op_set = set()
        for op_name in flags.experimental_select_user_tf_ops.split(','):
            user_op_set.add(op_name)
        converter.target_spec.experimental_select_user_tf_ops = list(user_op_set)
    if flags.post_training_quantize:
        converter.optimizations = [lite.Optimize.DEFAULT]
        if converter.inference_type != dtypes.float32:
            print('--post_training_quantize quantizes a graph of inference_type FLOAT. Overriding inference_type to FLOAT.')
            converter.inference_type = dtypes.float32
    if flags.quantize_to_float16:
        converter.target_spec.supported_types = [dtypes.float16]
        if not flags.post_training_quantize:
            print('--quantize_to_float16 will only take effect with the --post_training_quantize flag enabled.')
    if flags.dump_graphviz_dir:
        converter.dump_graphviz_dir = flags.dump_graphviz_dir
    if flags.dump_graphviz_video:
        converter.dump_graphviz_vode = flags.dump_graphviz_video
    if flags.conversion_summary_dir:
        converter.conversion_summary_dir = flags.conversion_summary_dir
    converter.experimental_new_converter = flags.experimental_new_converter
    if flags.experimental_new_quantizer is not None:
        converter.experimental_new_quantizer = flags.experimental_new_quantizer
    output_data = converter.convert()
    with gfile.GFile(flags.output_file, 'wb') as f:
        f.write(output_data)

Calls function to convert the TensorFlow 1.X model into a TFLite model.

Args:
        flags: argparse.Namespace object.

Raises:
        ValueError: Invalid flags.


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

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

    Note that this currently may not work inside a control flow construct.

Args:
        x: an unvectorized Tensor.

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


def _autopacking_helper(list_or_tuple, dtype, name):
    if context.executing_eagerly():
        if all((isinstance(elem, core.Tensor) for elem in list_or_tuple)):
            return gen_array_ops.pack(list_or_tuple, name=name)
    must_pack = False
    converted_elems = []
    with ops.name_scope(name) as scope:
        for i, elem in enumerate(list_or_tuple):
            if isinstance(elem, core.Tensor):
                if dtype is not None and elem.dtype.base_dtype != dtype:
                    raise TypeError(f'Cannot convert a list containing a tensor of dtype {elem.dtype} to {dtype} (Tensor is: {elem!r})')
                converted_elems.append(elem)
                must_pack = True
            elif isinstance(elem, (list, tuple)):
                converted_elem = _autopacking_helper(elem, dtype, str(i))
                if isinstance(converted_elem, core.Tensor):
                    must_pack = True
                converted_elems.append(converted_elem)
            else:
                converted_elems.append(elem)
        if must_pack:
            elems_as_tensors = []
            for i, elem in enumerate(converted_elems):
                if isinstance(elem, core.Tensor):
                    elems_as_tensors.append(elem)
                else:
                    elems_as_tensors.append(constant_op.constant(elem, dtype=dtype, name=str(i)))
            return gen_array_ops.pack(elems_as_tensors, name=scope)
        else:
            return converted_elems

Converts the given list or tuple to a tensor by packing.

Args:
        list_or_tuple: A (possibly nested) list or tuple containing a tensor.
        dtype: The element type of the returned tensor.
        name: A name for the returned tensor.

Returns:
        A `tf.Tensor` with value equivalent to `list_or_tuple`.


def _bigbird_block_rand_mask_with_head(self, from_seq_length, to_seq_length, from_block_size, to_block_size, num_heads, plan_from_length, plan_num_rand_blocks, indices_prng_key: Optional[jax.random.PRNGKey]=None, deterministic: Optional[bool]=True, window_block_left=1, window_block_right=1, global_block_top=1, global_block_bottom=1, global_block_left=1, global_block_right=1):
    if from_seq_length // from_block_size != to_seq_length // to_block_size:
        raise ValueError('Error the number of blocks needs to be same!')
    if from_seq_length not in plan_from_length:
        raise ValueError('Error from sequence length not in plan!')
    num_blocks = from_seq_length // from_block_size
    plan_block_length = jnp.array(plan_from_length) // from_block_size
    max_plan_idx = plan_from_length.index(from_seq_length)
    rand_attn = [jnp.zeros((num_blocks, sum(plan_num_rand_blocks[:max_plan_idx + 1])), dtype=jnp.int32) for i in range(num_heads)]
    if deterministic:
        for nh in range(num_heads):
            rand_attn[nh] = rand_attn[nh][global_block_top:num_blocks - global_block_bottom, :]
        return rand_attn
    for plan_idx in range(max_plan_idx + 1):
        rnd_r_cnt = 0
        if plan_idx > 0:
            if plan_num_rand_blocks[plan_idx] > 0:
                rnd_r_cnt = int(sum(plan_num_rand_blocks[:plan_idx]))
                curr_r_cnt = int(sum(plan_num_rand_blocks[:plan_idx + 1]))
                for blk_rw_idx in range(global_block_top, plan_block_length[plan_idx - 1]):
                    for h in range(num_heads):
                        single_block_row_attention = self._get_single_block_row_attention(block_id=blk_rw_idx, to_start_block_id=plan_block_length[plan_idx - 1], to_end_block_id=plan_block_length[plan_idx], num_rand_blocks=plan_num_rand_blocks[plan_idx], window_block_left=window_block_left, window_block_right=window_block_right, global_block_left=global_block_left, global_block_right=global_block_right, indices_prng_key=indices_prng_key)
                        rand_attn[h] = rand_attn[h].at[blk_rw_idx, rnd_r_cnt:curr_r_cnt].set(single_block_row_attention)
            for pl_id in range(plan_idx):
                if plan_num_rand_blocks[pl_id] == 0:
                    continue
                for blk_rw_idx in range(plan_block_length[plan_idx - 1], plan_block_length[plan_idx]):
                    rnd_r_cnt = 0
                    to_start_block_id = 0
                    if pl_id > 0:
                        rnd_r_cnt = int(sum(plan_num_rand_blocks[:pl_id]))
                        to_start_block_id = plan_block_length[pl_id - 1]
                    curr_r_cnt = int(sum(plan_num_rand_blocks[:pl_id + 1]))
                    for h in range(num_heads):
                        single_block_row_attention = self._get_single_block_row_attention(block_id=blk_rw_idx, to_start_block_id=to_start_block_id, to_end_block_id=plan_block_length[pl_id], num_rand_blocks=plan_num_rand_blocks[pl_id], window_block_left=window_block_left, window_block_right=window_block_right, global_block_left=global_block_left, global_block_right=global_block_right, indices_prng_key=indices_prng_key)
                        rand_attn[h] = rand_attn[h].at[blk_rw_idx, rnd_r_cnt:curr_r_cnt].set(single_block_row_attention)
        if plan_num_rand_blocks[plan_idx] == 0:
            continue
        curr_r_cnt = int(sum(plan_num_rand_blocks[:plan_idx + 1]))
        from_start_block_id = global_block_top
        to_start_block_id = 0
        if plan_idx > 0:
            rnd_r_cnt = int(sum(plan_num_rand_blocks[:plan_idx]))
            from_start_block_id = plan_block_length[plan_idx - 1]
            to_start_block_id = plan_block_length[plan_idx - 1]
        for blk_rw_idx in range(from_start_block_id, plan_block_length[plan_idx]):
            for h in range(num_heads):
                single_block_row_attention = self._get_single_block_row_attention(block_id=blk_rw_idx, to_start_block_id=to_start_block_id, to_end_block_id=plan_block_length[plan_idx], num_rand_blocks=plan_num_rand_blocks[plan_idx], window_block_left=window_block_left, window_block_right=window_block_right, global_block_left=global_block_left, global_block_right=global_block_right, indices_prng_key=indices_prng_key)
                rand_attn[h] = rand_attn[h].at[blk_rw_idx, rnd_r_cnt:curr_r_cnt].set(single_block_row_attention)
    for nh in range(num_heads):
        rand_attn[nh] = rand_attn[nh][global_block_top:num_blocks - global_block_bottom, :]
    return rand_attn

Create adjacency list of random attention.

Args:
        from_seq_length: int. length of from sequence.
        to_seq_length: int. length of to sequence.
        from_block_size: int. size of block in from sequence.
        to_block_size: int. size of block in to sequence.
        num_heads: int. total number of heads.
        plan_from_length: list. plan from length where num_random_blocks are chosen from.
        plan_num_rand_blocks: list. number of rand blocks within the plan.
        indices_prng_key: jax.random.PRNGKey. PRNG key that is used to perform random jax operations.
        deterministic: bool. When False random attention will be used.
        window_block_left: int. number of blocks of window to left of a block.
        window_block_right: int. number of blocks of window to right of a block.
        global_block_top: int. number of blocks at the top.
        global_block_bottom: int. number of blocks at the bottom.
        global_block_left: int. Number of blocks globally used to the left.
        global_block_right: int. Number of blocks globally used to the right.

Returns:
        adjacency list of size num_head where each element is of size from_seq_length//from_block_size-2 by
        num_rand_blocks


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

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

    Uses the new sink in an example pipeline.

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

    This method runs two pipelines.

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

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

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

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

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


def _benchmarkFeed(self, name, target, size, iters):
    feed_val = np.random.rand(size).astype(np.float32)
    times = []
    with ops.Graph().as_default():
        p = array_ops.placeholder(dtypes.float32, shape=[size])
        no_op = array_ops.identity(p).op
        with session.Session(target) as sess:
            sess.run(no_op, feed_dict={p: feed_val})
            for _ in range(iters):
                start_time = time.time()
                sess.run(no_op, feed_dict={p: feed_val})
                end_time = time.time()
                times.append(end_time - start_time)
    print('%s %d %f' % (name, size, np.median(times)))
    self.report_benchmark(iters=1, wall_time=np.median(times), name=name)

Runs a microbenchmark to measure the cost of feeding a tensor.

    Reports the median cost of feeding a tensor of `size` * `sizeof(float)`
    bytes.

Args:
        name: A human-readable name for logging the output.
        target: The session target to use for the benchmark.
        size: The number of floating-point numbers to be feed.
        iters: The number of iterations to perform.


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

Returns the map of control_outputs for a given graph.

Args:
        graph: The graph to parse.

Returns:
        A map of the control outputs.


def erf(x):
    if any_symbolic_tensors((x,)):
        return Erf().symbolic_call(x)
    x = backend.convert_to_tensor(x)
    return backend.math.erf(x)

Computes the error function of `x`, element-wise.

Args:
        x: Input tensor.

Returns:
        A tensor with the same dtype as `x`.

Example:
    >>> x = np.array([-3.0, -2.0, -1.0, 0.0, 1.0])
    >>> keras.ops.erf(x)
    array([-0.99998 , -0.99532, -0.842701,  0.,  0.842701], dtype=float32)


def with_options(self, options, name=None) -> 'DatasetV2':
    return _OptionsDataset(self, options, name=name)

Returns a new `tf.data.Dataset` with the given options set.

    The options are "global" in the sense they apply to the entire dataset.
    If options are set multiple times, they are merged as long as different
    options do not use different non-default values.

    >>> ds = tf.data.Dataset.range(5)
    >>> ds = ds.interleave(lambda x: tf.data.Dataset.range(5),
    ...                    cycle_length=3,
    ...                    num_parallel_calls=3)
    >>> options = tf.data.Options()
    >>> # This will make the interleave order non-deterministic.
    >>> options.deterministic = False
    >>> ds = ds.with_options(options)

Args:
        options: A `tf.data.Options` that identifies the options the use.
        name: (Optional.) A name for the tf.data operation.

Returns:
        A new `Dataset` with the transformation applied as described above.

Raises:
        ValueError: when an option is set more than once to a non-default value


def create_primary_api_files(output_files, packages, packages_to_ignore, root_init_template, output_dir, output_package, api_name, api_version, compat_api_versions, compat_init_templates, lazy_loading=_LAZY_LOADING, use_relative_imports=False):
    module_name_to_file_path = {}
    for output_file in output_files:
        module_name = get_module(os.path.dirname(output_file), output_dir)
        module_name_to_file_path[module_name] = os.path.normpath(output_file)
    for module, file_path in module_name_to_file_path.items():
        if not os.path.isdir(os.path.dirname(file_path)):
            os.makedirs(os.path.dirname(file_path))
        open(file_path, 'a').close()
    module_text_map, deprecation_footer_map, root_module_footer = get_api_init_text(packages, packages_to_ignore, output_package, api_name, api_version, compat_api_versions, lazy_loading, use_relative_imports)
    missing_output_files = []
    root_module = ''
    compat_module_to_template = {_COMPAT_MODULE_TEMPLATE % v: t for v, t in zip(compat_api_versions, compat_init_templates)}
    for v in compat_api_versions:
        compat_module_to_template.update({_SUBCOMPAT_MODULE_TEMPLATE % (v, vs): t for vs, t in zip(compat_api_versions, compat_init_templates)})
    for module, text in module_text_map.items():
        if module not in module_name_to_file_path:
            module_file_path = '"%s/__init__.py"' % module.replace('.', '/')
            missing_output_files.append(module_file_path)
            continue
        contents = ''
        if module == root_module and root_init_template:
            with open(root_init_template, 'r') as root_init_template_file:
                contents = root_init_template_file.read()
                contents = contents.replace('# API IMPORTS PLACEHOLDER', text)
                contents = contents.replace('# __all__ PLACEHOLDER', root_module_footer)
        elif module in compat_module_to_template:
            with open(compat_module_to_template[module], 'r') as init_template_file:
                contents = init_template_file.read()
                contents = contents.replace('# API IMPORTS PLACEHOLDER', text)
        else:
            contents = _GENERATED_FILE_HEADER % get_module_docstring(module, packages[0], api_name) + text + _GENERATED_FILE_FOOTER
        if module in deprecation_footer_map:
            if '# WRAPPER_PLACEHOLDER' in contents:
                contents = contents.replace('# WRAPPER_PLACEHOLDER', deprecation_footer_map[module])
            else:
                contents += deprecation_footer_map[module]
        with open(module_name_to_file_path[module], 'w') as fp:
            fp.write(contents)
    if missing_output_files:
        missing_files = ',\n'.join(sorted(missing_output_files))
        raise ValueError(f'Missing outputs for genrule:\n{missing_files}. Be sure to add these targets to tensorflow/python/tools/api/generator/api_init_files_v1.bzl and tensorflow/python/tools/api/generator/api_init_files.bzl (tensorflow repo), or tf_keras/api/api_init_files.bzl (tf_keras repo)')

Creates __init__.py files for the Python API.

Args:
        output_files: List of __init__.py file paths to create.
        packages: Base python packages containing python with target tf_export
        decorators.
        packages_to_ignore: python packages to be ignored when checking for
        tf_export decorators.
        root_init_template: Template for top-level __init__.py file. "# API IMPORTS
        PLACEHOLDER" comment in the template file will be replaced with imports.
        output_dir: output API root directory.
        output_package: Base output package where generated API will be added.
        api_name: API you want to generate Currently, only `tensorflow`.
        api_version: API version to generate (`v1` or `v2`).
        compat_api_versions: Additional API versions to generate in compat/
        subdirectory.
        compat_init_templates: List of templates for top level compat init files in
        the same order as compat_api_versions.
        lazy_loading: Boolean flag. If True, a lazy loading `__init__.py` file is
        produced and if `False`, static imports are used.
        use_relative_imports: True if we should use relative imports when import
        submodules.

Raises:
        ValueError: if output_files list is missing a required file.


def matmul(self, matmul_input: core.Tensor) -> Mapping[str, core.Tensor]:
    out = math_ops.matmul(matmul_input, self.matmul_filters)
    return {'output': out}

Performs a matrix multiplication.

Args:
        matmul_input: Input tensor to matmul with the filter.

Returns:
        A map of: output key -> output result.


def layer_statistics_dump(self, file: IO[str]) -> None:
    fields = ['op_name', 'tensor_idx'] + list(self._layer_debug_metrics.keys())
    if self._debug_options.layer_direct_compare_metrics is not None:
        fields += list(self._debug_options.layer_direct_compare_metrics.keys())
    fields += ['scale', 'zero_point', 'tensor_name']
    writer = csv.DictWriter(file, fields)
    writer.writeheader()
    if self.layer_statistics:
        for name, metrics in self.layer_statistics.items():
            data = metrics.copy()
            data['tensor_name'], _ = self._get_operand_name_and_index(name)
            data['tensor_idx'] = self._numeric_verify_op_details[name]['inputs'][0]
            data['op_name'] = self._quant_interpreter._get_op_details(self._defining_op[data['tensor_idx']])['op_name']
            details = self._quant_interpreter._get_tensor_details(data['tensor_idx'], subgraph_index=0)
            data['scale'], data['zero_point'] = (details['quantization_parameters']['scales'][0], details['quantization_parameters']['zero_points'][0])
            writer.writerow(data)

Dumps layer statistics into file, in csv format.

Args:
        file: file, or file-like object to write.


def _tensordot_reshape(a, axes, flipped=False):
    if a.get_shape().is_fully_defined() and isinstance(axes, (list, tuple)):
        shape_a = a.get_shape().as_list()
        axes = [i if i >= 0 else i + len(shape_a) for i in axes]
        free = [i for i in builtins.range(len(shape_a)) if i not in axes]
        free_dims = [shape_a[i] for i in free]
        prod_free = int(np.prod([shape_a[i] for i in free]))
        prod_axes = int(np.prod([shape_a[i] for i in axes]))
        perm = list(axes) + free if flipped else free + list(axes)
        new_shape = [prod_axes, prod_free] if flipped else [prod_free, prod_axes]
        if (perm != np.arange(len(shape_a))).any():
            a_trans = array_ops.transpose(a, perm)
        else:
            a_trans = a
        if a_trans.get_shape().as_list() != new_shape:
            reshaped_a = array_ops.reshape(a_trans, new_shape)
        else:
            reshaped_a = a_trans
        return (reshaped_a, free_dims, free_dims)
    else:
        if a.get_shape().ndims is not None and isinstance(axes, (list, tuple)):
            shape_a = a.get_shape().as_list()
            axes = [i if i >= 0 else i + len(shape_a) for i in axes]
            free = [i for i in builtins.range(len(shape_a)) if i not in axes]
            axes_dims = [shape_a[i] for i in axes]
            free_dims = [shape_a[i] for i in free]
            free_dims_static = free_dims
            axes = ops.convert_to_tensor(axes, dtype=dtypes.int32, name='axes')
            free = ops.convert_to_tensor(free, dtype=dtypes.int32, name='free')
            shape_a = array_ops.shape(a)
        else:
            free_dims_static = None
            shape_a = array_ops.shape(a)
            rank_a = array_ops.rank(a)
            axes = ops.convert_to_tensor(axes, dtype=dtypes.int32, name='axes')
            axes = array_ops.where(axes >= 0, axes, axes + rank_a)
            free, _ = gen_array_ops.list_diff(range(rank_a), axes, dtypes.int32)
        free_dims = array_ops.gather(shape_a, free)
        axes_dims = array_ops.gather(shape_a, axes)
        prod_free_dims = reduce_prod(free_dims)
        prod_axes_dims = reduce_prod(axes_dims)
        if flipped:
            perm = array_ops.concat([axes, free], 0)
            new_shape = array_ops_stack.stack([prod_axes_dims, prod_free_dims])
        else:
            perm = array_ops.concat([free, axes], 0)
            new_shape = array_ops_stack.stack([prod_free_dims, prod_axes_dims])
        reshaped_a = array_ops.reshape(array_ops.transpose(a, perm), new_shape)
        return (reshaped_a, free_dims, free_dims_static)

Helper method to perform transpose and reshape for contraction op.

    This method is helpful in reducing `math_ops.tensordot` to `math_ops.matmul`
    using `array_ops.transpose` and `array_ops.reshape`. The method takes a
    tensor and performs the correct transpose and reshape operation for a given
    set of indices. It returns the reshaped tensor as well as a list of indices
    necessary to reshape the tensor again after matrix multiplication.

Args:
        a: `Tensor`.
        axes: List or `int32` `Tensor` of unique indices specifying valid axes of
        `a`.
        flipped: An optional `bool`. Defaults to `False`. If `True`, the method
        assumes that `a` is the second argument in the contraction operation.

Returns:
        A tuple `(reshaped_a, free_dims, free_dims_static)` where `reshaped_a` is
        the tensor `a` reshaped to allow contraction via `matmul`, `free_dims` is
        either a list of integers or an `int32` `Tensor`, depending on whether
        the shape of a is fully specified, and free_dims_static is either a list
        of integers and None values, or None, representing the inferred
        static shape of the free dimensions


def from_vision_text_configs(cls, vision_config: PretrainedConfig, text_config: PretrainedConfig, **kwargs):
    return cls(vision_config=vision_config.to_dict(), text_config=text_config.to_dict(), **kwargs)

Instantiate a [`VisionTextDualEncoderConfig`] (or a derived class) from text model configuration and vision
    model configuration.

Returns:
        [`VisionTextDualEncoderConfig`]: An instance of a configuration object


def _safe_close(self, sess: session.Session):
    try:
        sess.close()
    except Exception:
        pass

Closes a session without raising an exception.

    Just like sess.close() but ignores exceptions.

Args:
        sess: A `Session`.


def _wrap_call_and_conditional_losses(layer):
    layer_call = _get_layer_call_method(layer)

    def call_and_return_conditional_losses(*args, **kwargs):
        """Returns layer (call_output, conditional losses) tuple."""
        call_output = layer_call(*args, **kwargs)
        if version_utils.is_v1_layer_or_model(layer):
            conditional_losses = layer.get_losses_for(_filtered_inputs([args, kwargs]))
        else:
            conditional_losses = [l for l in layer.losses if not hasattr(l, '_unconditional_loss')]
        return (call_output, conditional_losses)
    return _create_call_fn_decorator(layer, call_and_return_conditional_losses)

Wraps call function that returns a tuple of (outputs, losses).

    The losses returned are conditional on the inputs passed to the call function.
    Unconditional losses (e.g. weight regularizeration) are wrapped separately.

Args:
        layer: a Keras layer object

Returns:
        python call function that returns outputs and conditional losses -- excludes
        activity regularizer


def get_impacted_files_from_tiny_model_summary(diff_with_last_commit: bool=False) -> List[str]:
    repo = Repo(PATH_TO_REPO)
    folder = Path(repo.working_dir)
    if not diff_with_last_commit:
        print(f'main is at {repo.refs.main.commit}')
        print(f'Current head is at {repo.head.commit}')
        commits = repo.merge_base(repo.refs.main, repo.head)
        for commit in commits:
            print(f'Branching commit: {commit}')
    else:
        print(f'main is at {repo.head.commit}')
        commits = repo.head.commit.parents
        for commit in commits:
            print(f'Parent commit: {commit}')
    if not os.path.isfile(folder / 'tests/utils/tiny_model_summary.json'):
        return []
    files = set()
    for commit in commits:
        with checkout_commit(repo, commit):
            with open(folder / 'tests/utils/tiny_model_summary.json', 'r', encoding='utf-8') as f:
                old_content = f.read()
        with open(folder / 'tests/utils/tiny_model_summary.json', 'r', encoding='utf-8') as f:
            new_content = f.read()
        old_content = json.loads(old_content)
        new_content = json.loads(new_content)
        old_keys = set(old_content.keys())
        new_keys = set(new_content.keys())
        keys_with_diff = old_keys.symmetric_difference(new_keys)
        common_keys = old_keys.intersection(new_keys)
        for key in common_keys:
            if old_content[key] != new_content[key]:
                keys_with_diff.add(key)
        impacted_model_classes = []
        for key in keys_with_diff:
            if key in new_keys:
                impacted_model_classes.extend(new_content[key]['model_classes'])
        with open(folder / 'src/transformers/__init__.py') as fp:
            lines = fp.readlines()
            new_lines = []
            for line in lines:
                if line == '_import_structure = {\n':
                    new_lines.append(line)
                elif line == '# Direct imports for type-checking\n':
                    break
                elif len(new_lines) > 0:
                    line = re.sub('is_.+_available\\(\\)', 'True', line)
                    line = line.replace('OptionalDependencyNotAvailable', 'Exception')
                    line = line.replace('Exception()', 'Exception')
                    new_lines.append(line)
        with tempfile.TemporaryDirectory() as tmpdirname:
            with open(os.path.join(tmpdirname, 'temp_init.py'), 'w') as fp:
                fp.write(''.join(new_lines))
            spec = importlib.util.spec_from_file_location('temp_init', os.path.join(tmpdirname, 'temp_init.py'))
            module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(module)
            import_structure = module._import_structure
            reversed_structure = {}
            for key, values in import_structure.items():
                for value in values:
                    reversed_structure[value] = key
            for model_class in impacted_model_classes:
                module = reversed_structure[model_class]
                framework = ''
                if model_class.startswith('TF'):
                    framework = 'tf'
                elif model_class.startswith('Flax'):
                    framework = 'flax'
                fn = f'modeling_{module.split('.')[-1]}.py' if framework == '' else f'modeling_{framework}_{module.split('.')[-1]}.py'
                files.add(f'src.transformers.{module}.{fn}'.replace('.', os.path.sep).replace(f'{os.path.sep}py', '.py'))
    return sorted(files)

Return a list of python modeling files that are impacted by the changes of `tiny_model_summary.json` in between:

    - the current head and the main branch if `diff_with_last_commit=False` (default)
    - the current head and its parent commit otherwise.

Returns:
        `List[str]`: The list of Python modeling files that are impacted by the changes of `tiny_model_summary.json`.


def remove_low_and_no_objects(masks, scores, labels, object_mask_threshold, num_labels):
    if not masks.shape[0] == scores.shape[0] == labels.shape[0]:
        raise ValueError('mask, scores and labels must have the same shape!')
    to_keep = labels.ne(num_labels) & (scores > object_mask_threshold)
    return (masks[to_keep], scores[to_keep], labels[to_keep])

Binarize the given masks using `object_mask_threshold`, it returns the associated values of `masks`, `scores` and
    `labels`.

Args:
        masks (`torch.Tensor`):
            A tensor of shape `(num_queries, height, width)`.
        scores (`torch.Tensor`):
            A tensor of shape `(num_queries)`.
        labels (`torch.Tensor`):
            A tensor of shape `(num_queries)`.
        object_mask_threshold (`float`):
            A number between 0 and 1 used to binarize the masks.

Raises:
        `ValueError`: Raised when the first dimension doesn't match in all input tensors.

Returns:
        `Tuple[`torch.Tensor`, `torch.Tensor`, `torch.Tensor`]`: The `masks`, `scores` and `labels` without the region
        < `object_mask_threshold`.


def __call__(self, image: Union[str, List[str], 'Image.Image', List['Image.Image']], candidate_labels: List[str], **kwargs: Any) -> Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]]:
    if 'images' in kwargs:
        image = kwargs.pop('images')
    if image is None:
        raise ValueError('Cannot call the zero-shot-image-classification pipeline without an images argument!')
    return super().__call__(image, candidate_labels=candidate_labels, **kwargs)

Assign labels to the image(s) passed as inputs.

Args:
        image (`str`, `List[str]`, `PIL.Image` or `List[PIL.Image]`):
            The pipeline handles three types of images:

            - A string containing a http link pointing to an image
            - A string containing a local path to an image
            - An image loaded in PIL directly

        candidate_labels (`List[str]`):
            The candidate labels for this image. They will be formatted using *hypothesis_template*.

        hypothesis_template (`str`, *optional*, defaults to `"This is a photo of {}"`):
            The format used in conjunction with *candidate_labels* to attempt the image classification by
            replacing the placeholder with the candidate_labels. Pass "{}" if *candidate_labels* are
            already formatted.

        timeout (`float`, *optional*, defaults to None):
            The maximum time in seconds to wait for fetching images from the web. If None, no timeout is set and
            the call may block forever.

Returns:
        A list of dictionaries containing one entry per proposed label. Each dictionary contains the
        following keys:
        - **label** (`str`) -- One of the suggested *candidate_labels*.
        - **score** (`float`) -- The score attributed by the model to that label. It is a value between
            0 and 1, computed as the `softmax` of `logits_per_image`.


def save(self, file_prefix, options=None):
    if options and options.experimental_enable_async_checkpoint:
        self._checkpoint_options = options
        if checkpoint_context.in_preemption_save_context():
            if self._async_checkpointer_impl is not None:
                self._async_checkpointer_impl.sync()
            logging.warning('Switching to regular sync checkpoint for preemption checkpoint.')
        elif context.executing_eagerly():
            return self._async_checkpointer().save(file_prefix, options)
        else:
            logging.warning('Saving async checkpoint in graph mode is currently not supported; switching to regular sync checkpoint instead.')
    if isinstance(file_prefix, os.PathLike):
        file_prefix = os.fspath(file_prefix)
    options = copy.copy(options) or checkpoint_options.CheckpointOptions()
    graph_building = not context.executing_eagerly()
    if graph_building:
        if ops.inside_function():
            raise NotImplementedError('Calling tf.train.Checkpoint.save() from a function is not supported, as save() modifies saving metadata in ways not supported by TensorFlow Operations. Consider using tf.train.Checkpoint.write(), a lower-level API which does not update metadata. tf.train.latest_checkpoint and related APIs will not see this checkpoint.')
        session = get_session()
        if self._save_counter is None:
            session.run(self.save_counter.initializer)
    if not graph_building or self._save_assign_op is None:
        with ops.colocate_with(self.save_counter):
            assign_op = self.save_counter.assign_add(1, read_value=True)
        if graph_building:
            self._save_assign_op = data_structures.NoDependency(assign_op)
    if graph_building:
        checkpoint_number = session.run(self._save_assign_op)
    else:
        checkpoint_number = assign_op.numpy()
    if options.experimental_write_callbacks is None:
        options.experimental_write_callbacks = [_update_checkpoint_state_internal]
    else:
        options.experimental_write_callbacks.append(_update_checkpoint_state_internal)
    return self._write('%s-%d' % (file_prefix, checkpoint_number), options=options)

Saves a training checkpoint and provides basic checkpoint management.

    The saved checkpoint includes variables created by this object and any
    trackable objects it depends on at the time `Checkpoint.save()` is
    called.

    `save` is a basic convenience wrapper around the `write` method,
    sequentially numbering checkpoints using `save_counter` and updating the
    metadata used by `tf.train.latest_checkpoint`. More advanced checkpoint
    management, for example garbage collection and custom numbering, may be
    provided by other utilities which also wrap `write` and `read`.
    (`tf.train.CheckpointManager` for example).

    ```
    step = tf.Variable(0, name="step")
    checkpoint = tf.train.Checkpoint(step=step)
    checkpoint.save("/tmp/ckpt")

    # Later, read the checkpoint with restore()
    checkpoint.restore("/tmp/ckpt-1")

    # You can also pass options to save() and restore(). For example this
    # runs the IO ops on the localhost:
    options = tf.train.CheckpointOptions(experimental_io_device="/job:localhost")
    checkpoint.save("/tmp/ckpt", options=options)

    # Later, read the checkpoint with restore()
    checkpoint.restore("/tmp/ckpt-1", options=options)
    ```

Args:
        file_prefix: A prefix to use for the checkpoint filenames
        (/path/to/directory/and_a_prefix). Names are generated based on this
        prefix and `Checkpoint.save_counter`.
        options: Optional `tf.train.CheckpointOptions` object.

Returns:
        The full path to the checkpoint.


def _batched_mask_to_box(masks: 'torch.Tensor'):
    if torch.numel(masks) == 0:
        return torch.zeros(*masks.shape[:-2], 4, device=masks.device)
    shape = masks.shape
    height, width = shape[-2:]
    in_height, _ = torch.max(masks, dim=-1)
    in_height_coords = in_height * torch.arange(height, device=in_height.device)[None, :]
    bottom_edges, _ = torch.max(in_height_coords, dim=-1)
    in_height_coords = in_height_coords + height * ~in_height
    top_edges, _ = torch.min(in_height_coords, dim=-1)
    in_width, _ = torch.max(masks, dim=-2)
    in_width_coords = in_width * torch.arange(width, device=in_width.device)[None, :]
    right_edges, _ = torch.max(in_width_coords, dim=-1)
    in_width_coords = in_width_coords + width * ~in_width
    left_edges, _ = torch.min(in_width_coords, dim=-1)
    empty_filter = (right_edges < left_edges) | (bottom_edges < top_edges)
    out = torch.stack([left_edges, top_edges, right_edges, bottom_edges], dim=-1)
    out = out * (~empty_filter).unsqueeze(-1)
    out = out.reshape(*shape[:-2], 4)
    return out

Computes the bounding boxes around the given input masks. The bounding boxes are in the XYXY format which
    corresponds the following required indices:
        - LEFT: left hand side of the bounding box
        - TOP: top of the bounding box
        - RIGHT: right of the bounding box
        - BOTTOM: bottom of the bounding box

    Return [0,0,0,0] for an empty mask. For input shape channel_1 x channel_2 x ... x height x width, the output shape
    is channel_1 x channel_2 x ... x 4.

Args:
        - masks (`torch.Tensor` of shape `(batch, nb_mask, height, width)`)


def get_num_bytes(self, batch: Sequence[scipy.sparse.csr_matrix]) -> int:
    return sum((sys.getsizeof(element) for element in batch))

Returns:
    The number of bytes of data for a batch.


def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor]=None, position_ids: Optional[torch.LongTensor]=None, past_key_value: Optional[Cache]=None, output_attentions: Optional[bool]=False, use_cache: Optional[bool]=False, cache_position: Optional[torch.LongTensor]=None, **kwargs) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
    residual = hidden_states
    hidden_states, self_attn_weights, present_key_value = self.self_attn(hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, **kwargs)
    hidden_states = self.input_layernorm(hidden_states)
    hidden_states = residual + hidden_states
    residual = hidden_states
    hidden_states = self.mlp(hidden_states)
    hidden_states = self.post_attention_layernorm(hidden_states)
    hidden_states = residual + hidden_states
    outputs = (hidden_states,)
    if output_attentions:
        outputs += (self_attn_weights,)
    if use_cache:
        outputs += (present_key_value,)
    return outputs

Args:
    hidden_states (`torch.FloatTensor`):
        input to the layer of shape `(batch, seq_len, embed_dim)`
    attention_mask (`torch.FloatTensor`, *optional*):
        attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,
        query_sequence_length, key_sequence_length)` if default attention is used.
    position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
        Indices of positions of each input sequence tokens in the position embeddings
    past_key_value (`Tuple(torch.FloatTensor)`, *optional*): 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.
    use_cache (`bool`, *optional*):
        If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
        (see `past_key_values`).
    cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
        Indices depicting the position of the input sequence tokens in the sequence.


def _sym_missing(self) -> Dict[str, Any]:
    missing = dict()
    for k, v in self._sym_attributes.items():
        if pg_typing.MISSING_VALUE != v and isinstance(v, base.Symbolic):
            missing_child = v.sym_missing(flatten=False)
            if missing_child:
                missing[k] = missing_child
    return missing

Returns missing values for Functor.

    Semantically unbound arguments are not missing, thus we only return partial
    bound arguments in `sym_missing`. As a result, a functor is partial only
    when any of its bound arguments is partial.

Returns:
        A dict of missing key (or path) to missing value.


def checksum(path):
    filesystem = FileSystems.get_filesystem(path)
    return filesystem.checksum(path)

Fetch checksum metadata of a file on the
    :class:`~apache_beam.io.filesystem.FileSystem`.

    This operation returns checksum metadata as stored in the underlying
    FileSystem. It should not read any file data. Checksum type and format are
    FileSystem dependent and are not compatible between FileSystems.

Args:
        path: string path of a file.

    Returns: string containing checksum

Raises:
        ``BeamIOError``: if path isn't a file or doesn't exist.


def set_weight_collections(self, weight_collections):
    self._weight_collections = weight_collections

Sets the weight collections for the layer.

Args:
        weight_collections: A list of collection names to which the Variable will
        be added.


def extrapolation_scheme(value_grid, t1, t2, equation_params_fn):
    first_half_step = implicit_scheme(value_grid, t1, (t1 + t2) / 2, equation_params_fn)
    two_half_steps = implicit_scheme(first_half_step, (t1 + t2) / 2, t2, equation_params_fn)
    full_step = implicit_scheme(value_grid, t1, t2, equation_params_fn)
    return 2 * two_half_steps - full_step

Constructs extrapolation implicit-explicit scheme.

    Performs two implicit half-steps, one full implicit step, and combines them
    with such coefficients that ensure second-order errors. More computationally
    expensive than Crank-Nicolson scheme, but provides a better approximation for
    high-wavenumber components, which results in absence of oscillations typical
    for Crank-Nicolson scheme in case of non-smooth initial conditions. See [1]
    for details.

    #### References:
    [1]: D. Lawson, J & Ll Morris, J. The Extrapolation of First Order Methods
    for Parabolic Partial Differential Equations. I. 1978
    SIAM Journal on Numerical Analysis. 15. 1212-1224.
    https://epubs.siam.org/doi/abs/10.1137/0715082

Args:
        value_grid: A `Tensor` of real dtype. Grid of solution values at the current
        time.
        t1: Time before the step.
        t2: Time after the step.
        equation_params_fn: A callable that takes a scalar `Tensor` argument
        representing time and constructs the tridiagonal matrix `A`
        (a tuple of three `Tensor`s, main, upper, and lower diagonals)
        and the inhomogeneous term `b`. All of the `Tensor`s are of the same
        `dtype` as `inner_value_grid` and of the shape broadcastable with the
        shape of `inner_value_grid`.

Returns:
        A `Tensor` of the same shape and `dtype` a
        `values_grid` and represents an approximate solution `u(t2)`.


def __init__(self, use_variable):
    super(GatherModel, self).__init__()
    w_val = np.random.randn(128, 32).astype('f4')
    if use_variable:
        self.w = variables.Variable(w_val)
    else:
        self.w = w_val

Initializes a GatherModel.

Args:
        use_variable: If True, creates a variable for weight.


def _write_object_proto(self, proto, options):
    write_object_proto_for_resource_variable(self, proto, options)

Writes additional information of the variable into the SavedObject proto.

    Subclasses of ResourceVariables could choose to override this method to
    customize extra information to provide when saving a SavedModel.

    Ideally, this should contain the logic in
    write_object_proto_for_resource_variable but `DistributedValue` is an
    outlier at the momemnt. Once `DistributedValue` becomes a proper
    ResourceVariable, we should remove the helper method below.

Args:
        proto: `SavedObject` proto to update.
        options: A `SaveOption` instance that configures save behavior.


def stop_standing_subprocess(proc):
    logging.debug('Stopping standing subprocess %d', proc.pid)
    _kill_process_tree(proc)
    if proc.stdout:
        proc.stdout.close()
    if proc.stderr:
        proc.stderr.close()
    proc.wait()
    logging.debug('Stopped standing subprocess %d', proc.pid)

Stops a subprocess started by start_standing_subprocess.

    Before killing the process, we check if the process is running, if it has
    terminated, Error is raised.

    Catches and ignores the PermissionError which only happens on Macs.

Args:
        proc: Subprocess to terminate.

Raises:
        Error: if the subprocess could not be stopped.


def elementwise_division(func):
    sparse_func = jax_sparse.sparsify(func)

    @functools.wraps(func)
    def sparse_wrapper(x1, x2):
        if isinstance(x1, jax_sparse.JAXSparse):
            if isinstance(x2, jax_sparse.JAXSparse):
                x1 = x1.todense()
                x2 = x2.todense()
            elif isinstance(x1, jax_sparse.BCOO):
                if not hasattr(x2, 'shape') or len(x2.shape) == 0:
                    return jax_sparse.BCOO((func(x1.data, x2), x1.indices), shape=x1.shape, indices_sorted=x1.indices_sorted, unique_indices=x1.unique_indices)
                else:
                    if not jax_utils.is_in_jax_tracing_scope(x2):
                        x2_zeros_and_nans = jnp.equal(x2, 0)
                        if not jnp.issubdtype(x2.dtype, jnp.integer):
                            x2_zeros_and_nans = jnp.logical_or(x2_zeros_and_nans, jnp.isnan(x2))
                        x2_zeros_and_nans = jax_sparse.bcoo_fromdense(x2_zeros_and_nans, n_batch=x1.n_batch, n_dense=x1.n_dense, index_dtype=x1.indices.dtype)
                        x1 = bcoo_add_indices(x1, x2_zeros_and_nans, sum_duplicates=True)
                    return sparse_func(x1, x2)
            else:
                raise ValueError(f'Unsupported sparse format: {x1.__class__}')
        elif isinstance(x2, jax_sparse.JAXSparse):
            x2 = x2.todense()
        return func(x1, x2)
    return sparse_wrapper

Decorator to add support for `BCOO` sparse tensors to element-wise binary
    division and related operators.

    This decorator is designed for operations related to the division of two
    two operands (e.g. `divide`). It accepts `BCOO` tensors for both the
    dividend and the divisor, but handles them differently based on whether they
    are the dividend or the divisor.

    - If the divisor is sparse, it is densified and the result is dense because
        the result contains Inf or Nan outside of the indices of the dividend.
    - If the dividend is sparse and the divisor is dense, it finds occurrences
        of zeros and NaNs in the divisor. The result may therefore have more
        indices than there were in the dividend to return correct values where the
        divisor was zero or NaN.
    - If the dividend is sparse and the divisor is a scalar, it does the
        division element-wise. Note that the result is incorrectly sparse if the
        scalar divisor is zero.

Args:
        func: The function to wrap.

Returns:
        Wrapped function that supports `BCOO` sparse tensors.


def remove_objects_from_args(args: Iterable[Any], kwargs: Dict[str, Any], pvalue_class: Union[Type[T], Tuple[Type[T], ...]]) -> Tuple[List[Any], Dict[str, Any], List[T]]:
    pvals = []

    def swapper(value):
        pvals.append(value)
        return ArgumentPlaceholder()
    new_args = [swapper(v) if isinstance(v, pvalue_class) else v for v in args]
    new_kwargs = dict(((k, swapper(v)) if isinstance(v, pvalue_class) else (k, v) for k, v in sorted(kwargs.items())))
    return (new_args, new_kwargs, pvals)

For internal use only; no backwards-compatibility guarantees.

    Replaces all objects of a given type in args/kwargs with a placeholder.

Args:
        args: A list of positional arguments.
        kwargs: A dictionary of keyword arguments.
        pvalue_class: A class object representing the types of arguments that must
        be replaced with a placeholder value (instance of ArgumentPlaceholder).

Returns:
        A 3-tuple containing a modified list of positional arguments, a modified
        dictionary of keyword arguments, and a list of all objects replaced with
        a placeholder value.


def get_registered_object(name, custom_objects=None, module_objects=None):
    if name in _GLOBAL_CUSTOM_OBJECTS:
        return _GLOBAL_CUSTOM_OBJECTS[name]
    elif custom_objects and name in custom_objects:
        return custom_objects[name]
    elif module_objects and name in module_objects:
        return module_objects[name]
    return None

Returns the class associated with `name` if it is registered with Keras.

    This function is part of the Keras serialization and deserialization
    framework. It maps strings to the objects associated with them for
    serialization/deserialization.

Example:
    ```
    def from_config(cls, config, custom_objects=None):
        if 'my_custom_object_name' in config:
        config['hidden_cls'] = tf.keras.utils.get_registered_object(
            config['my_custom_object_name'], custom_objects=custom_objects)
    ```

Args:
        name: The name to look up.
        custom_objects: A dictionary of custom objects to look the name up in.
        Generally, custom_objects is provided by the user.
        module_objects: A dictionary of custom objects to look the name up in.
        Generally, module_objects is provided by midlevel library implementers.

Returns:
        An instantiable class associated with 'name', or None if no such class
        exists.


def InsertNodesBefore(new_nodes, target):
    for node in new_nodes:
        _InsertNodeAt(node, target, after=False)

Insert new_nodes before the given target location in the tree.

Args:
        new_nodes: a sequence of new nodes to insert (the nodes should not be in the
        tree).
        target: the target node before which the new node node will be inserted.

Raises:
        RuntimeError: if the tree is corrupted, or the insertion would corrupt it.


def non_slot_devices(self, var_list):
    raise NotImplementedError('must be implemented in descendants')

Device(s) for non-slot variables.

    DEPRECATED: TF 1.x ONLY.

    This method returns non-slot devices where non-slot variables are placed.
    Users can create non-slot variables on these devices by using a block:

    ```python
    with tf.distribute.StrategyExtended.colocate_vars_with(tf.distribute.StrategyExtended.non_slot_devices(...)):
        ...
    ```

Args:
        var_list: The list of variables being optimized, needed with the
        default `tf.distribute.Strategy`.

Returns:
        A sequence of devices for non-slot variables.


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, past_key_value: Optional[Tuple[torch.Tensor]]=None, output_attentions: Optional[bool]=False, use_cache: Optional[bool]=True) -> torch.Tensor:
    residual = hidden_states
    hidden_states = self.self_attn_layer_norm(hidden_states)
    self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
    hidden_states, self_attn_weights, present_key_value = self.self_attn(hidden_states=hidden_states, past_key_value=self_attn_past_key_value, attention_mask=attention_mask, output_attentions=output_attentions)
    hidden_states = self.attn_dropout(hidden_states)
    hidden_states = residual + hidden_states
    cross_attn_present_key_value = None
    cross_attn_weights = None
    if encoder_hidden_states is not None:
        residual = hidden_states
        hidden_states = self.cross_attention_layer_norm(hidden_states)
        cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
        hidden_states, cross_attn_weights, cross_attn_present_key_value = self.cross_attention(hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, past_key_value=cross_attn_past_key_value, attention_mask=encoder_attention_mask, output_attentions=output_attentions)
        hidden_states = self.attn_dropout(hidden_states)
        hidden_states = residual + hidden_states
        present_key_value += cross_attn_present_key_value
    residual = hidden_states
    hidden_states = self.ffn_layer_norm(hidden_states)
    hidden_states = self.ffn(hidden_states)
    hidden_states = self.ffn_dropout(hidden_states)
    hidden_states = residual + hidden_states
    outputs = (hidden_states, present_key_value)
    if output_attentions:
        outputs += (self_attn_weights, cross_attn_weights)
    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.
    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.


def _run_graph(self, device, input_shape, variable, num_inputs, axis, grad, num_iters):
    graph = ops.Graph()
    with graph.as_default():
        outputs = build_graph(device, input_shape, variable, num_inputs, axis, grad)
    config = config_pb2.ConfigProto(graph_options=config_pb2.GraphOptions(optimizer_options=config_pb2.OptimizerOptions(opt_level=config_pb2.OptimizerOptions.L0)))
    with session_lib.Session(graph=graph, config=config) as session:
        variables.global_variables_initializer().run()
        _ = session.run(outputs)
        start_time = time.time()
        for _ in range(num_iters):
            _ = session.run(outputs)
        duration = time.time() - start_time
        print('%s shape:%d/%d var: %r #inputs:%d axis:%d grad:%r - %f secs - %f GB/sec' % (device, input_shape[0], input_shape[1], variable, num_inputs, axis, grad, duration / num_iters, num_inputs * input_shape[0] * input_shape[1] * 4 * 2 * 100 / (duration / num_iters) / 1000000000.0))
    name_template = 'concat_bench_{device}_input_shape_{shape}_variable_{variable}_num_inputs_{num_inputs}_axis_{axis}_grad_{grad}'
    self.report_benchmark(name=name_template.format(device=device, num_inputs=num_inputs, variable=variable, grad=grad, shape=str(input_shape).replace(' ', ''), axis=str(axis), iters=num_iters))
    return duration

Run the graph and print its execution time.

Args:
        device: string, the device to run on.
        input_shape: shape of the input tensors.
        variable: whether or not the input shape should be fixed
        num_inputs: the number of inputs to concat
        axis: axis to be concat'ed
        grad: if True compute the gradient
        num_iters: number of steps to run.

Returns:
        The duration of the run in seconds.


def write(self, value, *labels):
    raise NotImplementedError

Writes the value to the given cache.

Args:
        value: An encodable (with corresponding PCoder) value
        *labels: List of labels for PCollection instance


def get_min_max_value(self) -> tuple[float, float]:
    average_min_max_statistics = self._statistics.average_min_max_statistics
    num_samples = average_min_max_statistics.num_samples
    if num_samples == 0:
        raise ValueError(f'num_samples must not be 0 when calibration method is AverageMinMax: {self._calib_opts}')
    min_value, max_value = (average_min_max_statistics.min_sum / num_samples, average_min_max_statistics.max_sum / num_samples)
    return (min_value, max_value)

Calculates the average of min and max values.

Returns:
        (min_value, max_value): Min and max calculated using AverageMinMax

Raises:
        ValueError: num_samples is 0.


def all_matches(pcoll, regex):
    regex = Regex._regex_compile(regex)

    def _process(element):
        m = regex.match(element)
        if m:
            yield [m.group(ix) for ix in range(m.lastindex + 1)]
    return pcoll | FlatMap(_process)

Returns all matches (groups) if zero or more characters at the beginning
    of string match the regular expression.

Args:
        regex: the regular expression string or (re.compile) pattern.


def recipe_manual(config, auth_read):
    hello(config, {'auth': auth_read, 'hour': [], 'say': 'Hello Manual', 'sleep': 0})

Used by tests.

Args:
        auth_read (authentication) - Credentials used for reading data.


def matches_alias(self, alias: str) -> bool:
    del self
    del alias
    return False

Indicates whether the expression will be selected as the given alias.

    Intended to be over-ridden by sub-classes which can safely implement it.
    Given an expression and an alias, indicates whether the expression will be
    SELECT'd as the given alias. For example, an expression like `SELECT a.b`
    matches the alias 'b', making it equivalent to the expression
    `SELECT a.b AS b`.

Args:
        alias: The alias to compare the expression against.

Returns:
        True when the expression evaluates to the same name as the alias and False
        otherwise.


def discount_bond_price(self, state: types.RealTensor, times: types.RealTensor, maturities: types.RealTensor, name: str=None) -> types.RealTensor:
    name = name or self._name + '_discount_bond_prices'
    with tf.name_scope(name):
        x_t = tf.convert_to_tensor(state, self._dtype)
        times = tf.convert_to_tensor(times, self._dtype)
        maturities = tf.convert_to_tensor(maturities, self._dtype)
        input_shape_times = tf.shape(times)
        mean_reversion = self._mean_reversion
        y_t = self.state_y(times)
        y_t = tf.reshape(tf.transpose(y_t), tf.concat([input_shape_times, [self._dim, self._dim]], axis=0))
        values = self._bond_reconstitution(times, maturities, mean_reversion, x_t, y_t, 1, tf.shape(times)[0])
        return values[0][0]

Returns zero-coupon bond prices `P(t,T)` conditional on `x(t)`.

Args:
        state: A `Tensor` of real dtype and shape compatible with
        `(num_times, dim)` specifying the state `x(t)`.
        times: A `Tensor` of real dtype and shape `(num_times,)`. The time `t`
        at which discount bond prices are computed.
        maturities: A `Tensor` of real dtype and shape `(num_times,)`. The time
        to maturity of the discount bonds.
        name: Str. The name to give this op.
        Default value: `discount_bond_prices`.

Returns:
        A `Tensor` of real dtype and the same shape as `(num_times,)`
        containing the price of zero-coupon bonds.


def _pad_all_input(inputs: Iterable[core_types.Tensor], padded_shapes: List[Optional[tensor_shape.TensorShape]], padding_spec: PaddingSpec) -> Tuple[List[List[Any]], List[dynamic_padding.PaddingMap]]:
    maximum_static_shapes = []
    need_padding = []
    input_shape_tensors = []
    for core_idx, inputs_per_core in enumerate(inputs):
        for idx, input_tensor in enumerate(inputs_per_core):
            input_shape = input_tensor.get_shape().as_list()
            if core_idx == 0:
                input_shape_tensors.append([])
                maximum_static_shapes.append(input_shape)
                need_padding.append(np.full_like(input_shape, False, dtype=bool))
            else:
                for i, s in enumerate(input_shape):
                    if s is None or s != maximum_static_shapes[idx][i]:
                        need_padding[idx][i] = True
                maximum_static_shapes[idx] = max(input_shape, maximum_static_shapes[idx])
            real_input_shape = array_ops.shape(input_tensor)
            real_input_shape.op._set_attr(_POST_DEVICE_REWRITE_ATTR, attr_value_pb2.AttrValue(b=True))
            input_shape_tensors[idx].append(real_input_shape)
    maximum_shapes = []
    for shapes_per_input in input_shape_tensors:
        maximum_shapes.append(math_ops.reduce_max(array_ops_stack.stack(shapes_per_input), axis=0))
    padded_inputs = []
    real_shapes = []
    padding_maps = []
    for core_idx, inputs_per_core in enumerate(inputs):
        padded_inputs.append([])
        real_shapes.append([])
        real_shape_idx = len(inputs_per_core) - 1
        for idx, input_tensor in enumerate(inputs_per_core):
            input_shape_tensor = input_shape_tensors[idx][core_idx]
            input_shape = input_tensor.get_shape().as_list()
            padded_shape = padded_shapes[idx]
            if any(need_padding[idx]) and padded_shape is not None:
                for i, s in enumerate(input_shape):
                    if need_padding[idx][i]:
                        if core_idx == 0:
                            real_shape_idx += 1
                            padding_map = dynamic_padding.PaddingMap()
                            padding_map.arg_index = idx
                            padding_map.shape_index = i
                            padding_map.padding_arg_index = real_shape_idx
                            padding_maps.append(padding_map)
                        real_shapes[core_idx].append(math_ops.cast(input_shape_tensor[i], dtypes.int32))
                paddings = []
                for i, s in enumerate(padded_shape.dims):
                    if need_padding[idx][i]:
                        minimum_dynamic_dim_size = 2
                        if s.value is not None:
                            max_dim_size = max(s.value, minimum_dynamic_dim_size)
                        else:
                            max_dim_size = math_ops.maximum(maximum_shapes[idx][i], minimum_dynamic_dim_size)
                            if padding_spec == PaddingSpec.POWER_OF_TWO:
                                max_dim_size = _ceil_to_pow_of_n(max_dim_size, 2)
                        padding = [0, max_dim_size - input_shape_tensor[i]]
                    else:
                        padding = [0, 0]
                    paddings.append(padding)
                if input_tensor.get_shape().is_fully_defined():
                    padded_input = cond.cond(array_ops.constant(True), lambda: array_ops.pad(input_tensor, paddings), lambda: input_tensor)
                else:
                    padded_input = array_ops.pad(input_tensor, paddings)
                padded_input.op._set_attr(_POST_DEVICE_REWRITE_ATTR, attr_value_pb2.AttrValue(b=True))
                padded_inputs[core_idx].append(padded_input)
            else:
                padded_inputs[core_idx].append(input_tensor)
    num_replicas = len(padded_inputs)
    for i in range(num_replicas):
        padded_inputs[i].extend(real_shapes[i])
    return (padded_inputs, padding_maps)

Pad all input tensors given padded_shapes.

    The real shape tensors will be concatenated with the padded original inputs.

Args:
        inputs: The original inputs.
        padded_shapes: A list of padded shapes for each input. If an entry is None,
        no padding is performed.
        padding_spec: An enum specified by `tpu.PaddingSpec`. This describes the
        padding policy when the `inputs` to `tf.tpu.replicate` is dynamic.
        One usage is to enable automatic bucketizing on the inputs by setting the
        value to `tpu.PaddingSpec.POWER_OF_TWO`, which can help to reduce the
        recompilation in the XLA side.

Returns:
        The padded inputs and a PaddingMap list which maps the padded input
        dimension to the real shape argument index.


def test_antithetic_sample_paths_mean_2d(self, random_type, seed):
    mu = np.array([0.2, 0.7])
    a = np.array([[0.4, 0.1], [0.3, 0.2]])
    b = np.array([[0.33, -0.03], [0.21, 0.5]])

    def drift_fn(t, x):
        del x
        return mu * tf.sqrt(t)

    def vol_fn(t, x):
        del x
        return (a * t + b) * tf.ones([2, 2], dtype=t.dtype)
    times = np.array([0.1, 0.21, 0.32, 0.43, 0.55])
    num_samples = 5000
    x0 = np.array([0.1, -1.1])
    paths = self.evaluate(euler_sampling.sample(dim=2, drift_fn=drift_fn, volatility_fn=vol_fn, times=times, num_samples=num_samples, initial_state=x0, random_type=random_type, time_step=0.01, seed=seed))
    self.assertAllClose(paths.shape, (num_samples, 5, 2), atol=0)
    means = np.mean(paths, axis=0)
    times = np.reshape(times, [-1, 1])
    expected_means = x0 + 2.0 / 3.0 * mu * np.power(times, 1.5)
    self.assertAllClose(means, expected_means, rtol=0.005, atol=0.005)

Tests path properties for 2-dimentional anthithetic variates method.

    The same test as above but with `PSEUDO_ANTITHETIC` random type.
    We construct the following Ito processes.

    dX_1 = mu_1 sqrt(t) dt + s11 dW_1 + s12 dW_2
    dX_2 = mu_2 sqrt(t) dt + s21 dW_1 + s22 dW_2

    mu_1, mu_2 are constants.
    s_ij = a_ij t + b_ij

    For this process expected value at time t is (x_0)_i + 2/3 * mu_i * t^1.5.

Args:
        random_type: Random number type defined by tff.math.random.RandomType
        enum.
        seed: Random seed.


def _colocation_dict(self) -> dict[str, traceable_stack.TraceableObject]:
    locations_dict = self._colocation_code_locations or {}
    return locations_dict.copy()

Code locations for colocation context managers active at op creation.

    This property will return a dictionary for which the keys are nodes with
    which this Operation is colocated, and for which the values are
    traceable_stack.TraceableObject instances.  The TraceableObject instances
    record the location of the relevant colocation context manager but have the
    "obj" field set to None to prevent leaking private data.

    For example, suppose file_a contained these lines:

        file_a.py:
        14: node_a = tf.constant(3, name='NODE_A')
        15: with tf.compat.v1.colocate_with(node_a):
        16:   node_b = tf.constant(4, name='NODE_B')

    Then a TraceableObject t_obj representing the colocation context manager
    would have these member values:

        t_obj.obj -> None
        t_obj.filename = 'file_a.py'
        t_obj.lineno = 15

    and node_b.op._colocation_dict would return the dictionary

        { 'NODE_A': t_obj }

Returns:
        {str: traceable_stack.TraceableObject} as per this method's description,
        above.


def __init__(self, gain=1.0, dtype=dtypes.float32):
    self.gain = gain
    self.dtype = _assert_float_dtype(dtypes.as_dtype(dtype))

Initializer that generates the identity matrix.

    Only use for 2D matrices.

Args:
        gain: Multiplicative factor to apply to the identity matrix.
        dtype: Default data type, used if no `dtype` argument is provided when
        calling the initializer. Only floating point types are supported.


def np_array(values, dtype=None, copy=True, order='K'):
    if dtype is not None and np.issubdtype(dtype, np.number):
        return np.array(values, copy=copy, order=order).astype(dtype)
    else:
        return np.array(values, dtype=dtype, copy=copy, order=order)

Creates a NumPy array containing input values.

    It will make a copy of the object.

    In NumPy 2.x and later, strict type casting can lead to errors when values
    overflow the specified dtype. This function addresses this by replacing direct
    np.array(..., dtype=...) calls with np.array(...).astype(...). This allows for
    intended overflows, aligning with the behavior of older NumPy versions.

Args:
        values: Array_like objects. E.g., a python list, tuple, or an object whose
        __array__ method returns an array.
        dtype: The desired numpy data type for the array.
        copy: Bool. If True (default), then the array data is copied. If None, a
        copy will only be made if __array__ returns a copy, if obj is a nested
        sequence, or if a copy is needed to satisfy any of the other requirements
        (dtype, order, etc.). Note that any copy of the data is shallow, i.e., for
        arrays with object dtype, the new array will point to the same objects.
        For False it raises a ValueError if a copy cannot be avoided.
        order: {‘K’, ‘A’, ‘C’, ‘F’}.

Returns:
        A NumPy array with the specified data type.


def __init__(self, serialized_tensors: Mapping[base.Trackable, sharding_util.Shard], registered_savers: 'RegisteredSaversDict | None'=None, call_with_mapped_captures: 'MappedCapturesCallable | None'=None):
    self._shardable_tensors_by_task: MutableMapping[device_lib.DeviceSpec, MutableSequence[sharding_util.ShardableTensor]] = {}
    self._keys_to_restore_fn: MutableMapping[TensorKeyAndSliceSpec, RestoreFn] = {}
    self._restore_fn_to_keys: MutableMapping[RestoreFn, MutableSequence[TensorKeyAndSliceSpec]] = {}
    unique_tasks = set()
    for obj, tensor_dict in serialized_tensors.items():
        restore_fn = _restore_noop if obj is None else obj._restore_from_tensors
        for checkpoint_key, tensor_slice_dict in tensor_dict.items():
            if not isinstance(tensor_slice_dict, dict):
                tensor_slice_dict = {'': tensor_slice_dict}
            for slice_spec, tensor_save_spec in tensor_slice_dict.items():
                tensor_value = None
                if not isinstance(tensor_save_spec, saveable_object.SaveSpec):
                    tensor_value = tensor_save_spec
                    tensor_save_spec = saveable_object.SaveSpec(tensor=tensor_value, slice_spec=slice_spec, name=checkpoint_key, dtype=tensor_save_spec.dtype, device=tensor_save_spec.device)
                if (checkpoint_key, slice_spec) in self._keys_to_restore_fn:
                    raise ValueError('Received multiple tensors with the same checkpoint key and slice spec. This is invalid because one will overwrite the other in the checkpoint. This indicates a bug in the Checkpoint key-generation.')
                self._keys_to_restore_fn[checkpoint_key, slice_spec] = restore_fn
                self._restore_fn_to_keys.setdefault(restore_fn, []).append((checkpoint_key, slice_spec))
                if isinstance(tensor_save_spec.device, str):
                    device = device_lib.DeviceSpec.from_string(tensor_save_spec.device)
                    task = device_lib.DeviceSpec.from_string(saveable_object_util.set_cpu0(tensor_save_spec.device))
                else:
                    device = tensor_save_spec.device
                    task = device_lib.DeviceSpec.from_string(saveable_object_util.set_cpu0(device.to_string()))
                self._shardable_tensors_by_task.setdefault(task, []).append(sharding_util.ShardableTensor(_tensor_save_spec=tensor_save_spec, tensor=tensor_value, dtype=tensor_save_spec.dtype, device=device, name=tensor_save_spec.name, shape=None, slice_spec=slice_spec.strip(), checkpoint_key=checkpoint_key, trackable=obj))
                unique_tasks.add(saveable_object_util.set_cpu0(device.to_string()))
    self._num_unique_tasks = len(unique_tasks)
    self._registered_savers = {}
    if registered_savers:
        for registered_name, trackables in registered_savers.items():
            save_fn = _get_mapped_registered_save_fn(registration.get_save_function(registered_name), trackables, call_with_mapped_captures)
            restore_fn = _get_mapped_registered_restore_fn(registration.get_restore_function(registered_name), trackables, call_with_mapped_captures)
            self._registered_savers[registered_name] = (save_fn, restore_fn)

Specify a list of `SaveableObject`s to save and restore.

Args:
        serialized_tensors: A dictionary mapping `Trackable` to a tensor dict,
        which maps checkpoint_key -> (slice_spec ->) -> Tensor/SaveSpec. The
        `Trackable` key is used to get the `restore_from_tensors` function,
        and may be `None` if the tensor is not meant to be restored.
        registered_savers: A dictionary mapping `registration.RegisteredSaver`
        namedtuples to a dictionary of named Trackables. The keys of the
        Trackable dictionary are string names that uniquely identify the
        Trackable in the checkpoint.
        call_with_mapped_captures: TODO


def from_proto(saver_def, import_scope=None):
    return Saver(saver_def=saver_def, name=import_scope)

Returns a `Saver` object created from `saver_def`.

Args:
        saver_def: a `SaverDef` protocol buffer.
        import_scope: Optional `string`. Name scope to use.

Returns:
        A `Saver` built from saver_def.


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'.


def __init__(self, config: Mask2FormerConfig, weight_dict: Dict[str, float]):
    super().__init__()
    requires_backends(self, ['scipy'])
    self.num_labels = config.num_labels
    self.weight_dict = weight_dict
    self.eos_coef = config.no_object_weight
    empty_weight = torch.ones(self.num_labels + 1)
    empty_weight[-1] = self.eos_coef
    self.register_buffer('empty_weight', empty_weight)
    self.num_points = config.train_num_points
    self.oversample_ratio = config.oversample_ratio
    self.importance_sample_ratio = config.importance_sample_ratio
    self.matcher = Mask2FormerHungarianMatcher(cost_class=1.0, cost_dice=config.dice_weight, cost_mask=config.mask_weight, num_points=self.num_points)

The Mask2Former Loss. The loss is computed very similar to DETR. The process happens in two steps: 1) we
    compute hungarian assignment between ground truth masks and the outputs of the model 2) we supervise each pair
    of matched ground-truth / prediction (supervise class and mask)

Args:
        config (`Mask2FormerConfig`):
            The configuration for Mask2Former model also containing loss calculation specific parameters.
        weight_dict (`Dict[str, float]`):
            A dictionary of weights to be applied to the different losses.