def load(self, sess, tags, import_scope=None, **saver_kwargs):
    saved_model_proto = parse_saved_model(self._export_dir)
    metrics.IncrementReadApi(_LOADER_LABEL)
    with sess.graph.as_default():
        saver, _ = self.load_graph(sess.graph, tags, import_scope, **saver_kwargs)
        self.restore_variables(sess, saver, import_scope)
        self.run_init_ops(sess, tags, import_scope)
    meta_graph_def = self.get_meta_graph_def_from_tags(tags)
    if len(saved_model_proto.meta_graphs) == 1 and saved_model_proto.meta_graphs[0].HasField('object_graph_def'):
        metrics.IncrementRead(write_version='2')
    else:
        metrics.IncrementRead(write_version='1')
    return meta_graph_def

Load the MetaGraphDef graph and restore variable values into the session.

Args:
        sess: tf.compat.v1.Session to restore variable values.
        tags: a set of string tags identifying a MetaGraphDef.
        import_scope: Optional `string` -- if specified, prepend this string
        followed by '/' to all loaded tensor names. This scope is applied to
        tensor instances loaded into the passed session, but it is *not* written
        through to the static `MetaGraphDef` protocol buffer that is returned.
        **saver_kwargs: keyword arguments to pass to tf.train.import_meta_graph.

Returns:
        `MetagraphDef` proto of the graph that was loaded.


def session(self, graph: Optional[ops.Graph]=None, config: Optional[config_pb2.ConfigProto]=None, use_gpu: bool=True, force_gpu: bool=False) -> Iterator[s.Session]:
    if context.executing_eagerly():
        yield EagerSessionWarner()
    else:
        with self._create_session(graph, config, force_gpu) as sess:
            with self._constrain_devices_and_set_default(sess, use_gpu, force_gpu):
                yield sess

A context manager for a TensorFlow Session for use in executing tests.

    Note that this will set this session and the graph as global defaults.

    Use the `use_gpu` and `force_gpu` options to control where ops are run. If
    `force_gpu` is True, all ops are pinned to `/device:GPU:0`. Otherwise, if
    `use_gpu` is True, TensorFlow tries to run as many ops on the GPU as
    possible. If both `force_gpu and `use_gpu` are False, all ops are pinned to
    the CPU.

Example:
    ``` python
    class MyOperatorTest(test_util.TensorFlowTestCase):
        def testMyOperator(self):
        with self.session():
        valid_input = [1.0, 2.0, 3.0, 4.0, 5.0]
        result = MyOperator(valid_input).eval()
        self.assertEqual(result, [1.0, 2.0, 3.0, 5.0, 8.0]
        invalid_input = [-1.0, 2.0, 7.0]
        with self.assertRaisesOpError("negative input not supported"):
            MyOperator(invalid_input).eval()
    ```

Args:
        graph: Optional graph to use during the returned session.
        config: An optional config_pb2.ConfigProto to use to configure the
        session.
        use_gpu: If True, attempt to run as many ops as possible on GPU.
        force_gpu: If True, pin all ops to `/device:GPU:0`.

Yields:
        A Session object that should be used as a context manager to surround
        the graph building and execution code in a test case.


def complete_multipart_upload(self, request):
    parts = {'Parts': request.parts}
    try:
        self.client.complete_multipart_upload(Bucket=request.bucket, Key=request.object, UploadId=request.upload_id, MultipartUpload=parts)
    except Exception as e:
        raise messages.S3ClientError(str(e), get_http_error_code(e))

Completes a multipart upload to S3

Args:
        request: (UploadPartRequest) input message

Returns:
        (Void) The response message.


def test_enrichment(pcoll, enrichment_handler: str, handler_config: Dict[str, Any], timeout: Optional[float]=30):
    if enrichment_handler == 'BigTable':
        row_key = handler_config['row_key']
        bt_data = INPUT_TABLES['BigTable', handler_config['instance_id'], handler_config['table_id']]
        products = {str(data[row_key]): data for data in bt_data}

        def _fn(row):
            left = row._asdict()
            right = products[str(left[row_key])]
            left['product'] = left.get('product', None) or right
            return beam.Row(**left)
    elif enrichment_handler == 'BigQuery':
        row_key = handler_config['fields']
        dataset, table = handler_config['table_name'].split('.')[-2:]
        bq_data = INPUT_TABLES['BigQuery', str(dataset), str(table)]
        bq_data = {tuple((str(data[key]) for key in row_key)): data for data in bq_data}

        def _fn(row):
            left = row._asdict()
            right = bq_data[tuple((str(left[k]) for k in row_key))]
            row = {key: left.get(key, None) or right[key] for key in {*left.keys(), *right.keys()}}
            return beam.Row(**row)
    else:
        raise ValueError(f'{enrichment_handler} is not a valid enrichment_handler.')
    return pcoll | beam.Map(_fn)

Mocks the Enrichment transform for testing purposes.

    This PTransform simulates the behavior of the Enrichment transform by
    looking up data from predefined in-memory tables based on the provided
    `enrichment_handler` and `handler_config`.

    Note: The Github action that invokes these tests does not have gcp
    dependencies installed which is a prerequisite to
    apache_beam.transforms.enrichment.Enrichment as a top-level import.

Args:
        pcoll: The input PCollection.
        enrichment_handler: A string indicating the type of enrichment handler
        to simulate (e.g., 'BigTable', 'BigQuery').
        handler_config: A dictionary containing configuration details for the
        simulated handler (e.g., table names, row keys, fields).
        timeout: An optional timeout value (ignored in this mock).

Returns:
        A PCollection containing the enriched data.


def private_map(self):
    return self._private_map

A map from parents to symbols that should not be included at all.

    This map can be edited, but it should not be edited once traversal has
    begun.

Returns:
        The map marking symbols to not include.


def pool_direct_single_axis(input, axis, window_size, pooling_type, padding, dilation_rate, stride):
    effective_window_size = (window_size - 1) * dilation_rate + 1
    input_size = input.shape[axis]
    if padding == 'SAME':
        output_size = int(math.ceil(input_size / stride))
        total_padding_amount = max(0, (output_size - 1) * stride + effective_window_size - input_size)
        before_padding = total_padding_amount // 2
    elif padding == 'VALID':
        output_size = int(math.ceil((input_size - effective_window_size + 1) / stride))
        before_padding = 0
    else:
        raise ValueError('Unsupported padding type: %r' % (padding,))
    output_shape = input.shape[:axis] + (output_size,) + input.shape[axis + 1:]
    output = np.zeros(output_shape, input.dtype)
    initial_dim_selector = tuple((np.s_[:] for _ in range(axis)))
    if pooling_type == 'MAX':
        pooling_func = np.max
    elif pooling_type == 'AVG':
        pooling_func = np.mean
    else:
        raise ValueError('Unsupported pooling type: %r' % (pooling_type,))
    for output_pos in range(output_size):
        input_start_pos = output_pos * stride - before_padding
        input_end_pos = min(input_start_pos + effective_window_size, input_size)
        if input_start_pos < 0:
            input_start_pos += dilation_rate
        input_slice = np.s_[input_start_pos:input_end_pos:dilation_rate]
        output[initial_dim_selector + (output_pos,)] = pooling_func(input[initial_dim_selector + (input_slice,)], axis=axis)
    return output

Numpy implementation of pooling along a single axis.

    This is intended for testing only, and therefore isn't particularly efficient.

    See pool_direct below for the meaning of the arguments.

Args:
        input: numpy array.
        axis: axis along which to perform pooling.
        window_size: int >= 1.  Size of pooling window within axis.
        pooling_type: either "MAX" or "AVG".
        padding: either "SAME" or "VALID".
        dilation_rate: int >= 1.  Dilation factor for window, i.e. stride at which
        to sample input.
        stride: int >= 1.  Stride at which to generate output.

Returns:
        pooling output array of rank N+2.

Raises:
        ValueError: if arguments are invalid.


def forward(self, input_ids: torch.LongTensor, speaker_id: torch.Tensor, lang_id: torch.Tensor) -> Tuple[torch.Tensor]:
    hidden_states = self.unit_embedding(input_ids).transpose(1, 2)
    spkr = self.speaker_embedding(speaker_id).transpose(1, 2)
    lang = self.language_embedding(lang_id).transpose(1, 2)
    log_dur_pred = self.dur_predictor(hidden_states.transpose(1, 2))
    dur_out = torch.clamp(torch.round(torch.expm1(log_dur_pred)).long(), min=1)
    if hidden_states.size(0) == 1:
        hidden_states = torch.repeat_interleave(hidden_states, dur_out.view(-1), dim=2)
    else:
        if hidden_states.shape[0] > 1 and self.training:
            logger.warning('`self.training=True` and you use batching. You lose parallelism during the hifigan\n                               forward pass because the samples are interleaved.')
        hidden_states = [torch.repeat_interleave(hidden_state, duration, dim=-1).transpose(0, 1) for hidden_state, duration in zip(hidden_states, dur_out)]
        hidden_states = nn.utils.rnn.pad_sequence(hidden_states, batch_first=True).transpose(1, 2)
    spkr = spkr.repeat(1, 1, hidden_states.shape[-1])
    lang = lang.repeat(1, 1, hidden_states.shape[-1])
    hidden_states = torch.cat([lang, hidden_states, spkr], dim=1)
    hidden_states = self.hifi_gan(hidden_states)
    unit_lengths = self._get_dur_output_lengths(input_ids, dur_out)
    lengths = self._get_output_hifigan_lengths(unit_lengths)
    return (hidden_states, lengths)

Args:
    input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
        Indices of input sequence tokens in the vocabulary.

        Indices can be obtained using [`SeamlessM4Tv2TextToUnitForConditionalGeneration`]. [What are input
        IDs?](../glossary#input-ids)
    speaker_id (`int`, *optional*):
        The id of the speaker used for speech synthesis. Must be lower than `config.vocoder_num_spkrs`.
    tgt_lang (`str`, *optional*):
        The language id to use as target language for translation.


def forward(self, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.FloatTensor]=None, token_type_ids: Optional[torch.LongTensor]=None, position_ids: Optional[torch.LongTensor]=None, head_mask: Optional[torch.FloatTensor]=None, inputs_embeds: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[MaskedLMOutput, Tuple]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    outputs = self.albert(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    sequence_outputs = outputs[0]
    prediction_scores = self.predictions(sequence_outputs)
    masked_lm_loss = None
    if labels is not None:
        loss_fct = CrossEntropyLoss()
        masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
    if not return_dict:
        output = (prediction_scores,) + outputs[2:]
        return (masked_lm_loss,) + output if masked_lm_loss is not None else output
    return MaskedLMOutput(loss=masked_lm_loss, logits=prediction_scores, hidden_states=outputs.hidden_states, attentions=outputs.attentions)

labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
    Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
    config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
    loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`

Example:
    ```python
    >>> import torch
    >>> from transformers import AutoTokenizer, AlbertForMaskedLM

    >>> tokenizer = AutoTokenizer.from_pretrained("albert/albert-base-v2")
    >>> model = AlbertForMaskedLM.from_pretrained("albert/albert-base-v2")

    >>> # add mask_token
    >>> inputs = tokenizer("The capital of [MASK] is Paris.", return_tensors="pt")
    >>> with torch.no_grad():
    ...     logits = model(**inputs).logits

    >>> # retrieve index of [MASK]
    >>> mask_token_index = (inputs.input_ids == tokenizer.mask_token_id)[0].nonzero(as_tuple=True)[0]
    >>> predicted_token_id = logits[0, mask_token_index].argmax(axis=-1)
    >>> tokenizer.decode(predicted_token_id)
    'france'
    ```

    ```python
    >>> labels = tokenizer("The capital of France is Paris.", return_tensors="pt")["input_ids"]
    >>> labels = torch.where(inputs.input_ids == tokenizer.mask_token_id, labels, -100)
    >>> outputs = model(**inputs, labels=labels)
    >>> round(outputs.loss.item(), 2)
    0.81
    ```


def num_accelerators(self, task_type=None, task_id=None, config_proto=None):
    if self._tpu == 'local':
        return {'TPU': len([d for d in framework_config.list_logical_devices() if d.device_type == 'TPU'])}
    retry_count = 1
    while True:
        try:
            device_details = TPUClusterResolver._get_device_dict_and_cores(cluster_resolver_lib.get_accelerator_devices(self.master(), config_proto=config_proto))
            break
        except errors.DeadlineExceededError:
            error_message = 'Failed to connect to master. The TPU might not be ready (e.g. still scheduling) or the master address is incorrect: got (%s)' % self.master()
            if retry_count <= _TPU_CONN_RETRIES:
                logging.warning(error_message)
                logging.warning('Retrying (%d/%d)...', retry_count, _TPU_CONN_RETRIES)
                retry_count += 1
            else:
                raise RuntimeError(error_message)
    if device_details.total_cores:
        return {'TPU': TPUClusterResolver._verify_and_return_same_core_count(device_details.device_map)}
    return {'TPU': 0}

Returns the number of TPU cores per worker.

    Connects to the master and list all the devices present in the master,
    and counts them up. Also verifies that the device counts per host in the
    cluster is the same before returning the number of TPU cores per host.

Args:
        task_type: Unused.
        task_id: Unused.
        config_proto: Used to create a connection to a TPU master in order to
        retrieve the system metadata.

Raises:
        RuntimeError: If we cannot talk to a TPU worker after retrying or if the
        number of TPU devices per host is different.


def write_build_info(filename, key_value_list):
    build_info = {}
    if cuda_config:
        build_info.update(cuda_config.config)
    if tensorrt_config:
        build_info.update(tensorrt_config.config)
    for arg in key_value_list:
        key, value = arg.split('=')
        if value.lower() == 'true':
            build_info[key] = True
        elif value.lower() == 'false':
            build_info[key] = False
        else:
            build_info[key] = value.format(**build_info)
    sorted_build_info_pairs = sorted(build_info.items())
    contents = '\n# Copyright 2020 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the "License");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an "AS IS" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n"""Auto-generated module providing information about the build."""\nimport collections\n\nbuild_info = collections.OrderedDict(%s)\n' % sorted_build_info_pairs
    open(filename, 'w').write(contents)

Writes a Python that describes the build.

Args:
        filename: filename to write to.
        key_value_list: A list of "key=value" strings that will be added to the
        module's "build_info" dictionary as additional entries.


def expand_value_set_url(self, url: str) -> Optional[value_set_pb2.ValueSet]:
    value_set = self._value_set_from_url(url)
    if value_set is None:
        return None
    concept_sets = itertools.chain(value_set.compose.include, value_set.compose.exclude)
    if any((concept_set.filter for concept_set in concept_sets)):
        return None
    includes = [self._concept_set_to_expansion(value_set, include) for include in value_set.compose.include]
    excludes = [self._concept_set_to_expansion(value_set, exclude) for exclude in value_set.compose.exclude]
    if None in includes or None in excludes:
        return None
    logging.info('Expanding value set url: %s version: %s locally', value_set.url.value, value_set.version.value)
    includes = itertools.chain.from_iterable(includes)
    excludes = itertools.chain.from_iterable(excludes)
    codes_to_remove = set(((concept.version.value, concept.system.value, concept.code.value) for concept in excludes))
    codes = [concept for concept in includes if (concept.version.value, concept.system.value, concept.code.value) not in codes_to_remove]
    expanded_value_set = copy.deepcopy(value_set)
    expanded_value_set.expansion.contains.extend(codes)
    return expanded_value_set

Attempts to expand value sets without contacting a terminology service.

    For value sets with an extensional set of codes, collect all codes
    referenced in the value set's 'compose' field.

    For value sets which reference a code system without specifying the
    extensional set of codes within it, attempt to find the definition of the
    code system in the ValueSetResolver and expand to all codes in the code
    system. If the code system resource is not available locally, returns `None`
    indicating a terminology service should instead be used to find the value
    set expansion.

    If the value set has an intensional set of codes, returns `None` indicating
    a terminology service should instead be used to find the value set
    expansion.

    See https://www.hl7.org/fhir/valueset.html#int-ext
    for more details about value set expansion and intensional versus
    extensional value sets.

Args:
        url: The url of the value set for which to retrieve expanded codes.

Returns:
        The expanded value set or None if a terminology service should be
        consulted instead.


def conv(inp, fltr, window_strides, padding, dimension_numbers, filter_dilation=None):
    input_spec, filter_spec, output_spec = dimension_numbers
    if input_spec != output_spec:
        raise ValueError('Input and output data formats must be the same; got %s and %s' % (input_spec, output_spec))
    supported_filter_spec = ['WIO', 'HWIO', 'DHWIO']
    if filter_spec not in supported_filter_spec:
        raise ValueError('The supported data format for the filter are %s; got %s' % (supported_filter_spec, filter_spec))
    if input_spec[1:-1] != filter_spec[:-2]:
        raise ValueError('Input data format (%s) is not compatible with filter data format (%s)' % (input_spec, filter_spec))
    dtype = tf_np.result_type(inp, fltr)
    inp = tf_np.asarray(inp, dtype)
    fltr = tf_np.asarray(fltr, dtype)
    return tf_np.asarray(nn_ops.convolution_v2(input=inp, filters=fltr, padding=padding, strides=window_strides, dilations=filter_dilation, data_format=input_spec))

Convolution over an N-D array.

    See https://www.tensorflow.org/api_docs/python/tf/nn/convolution and
    https://www.tensorflow.org/xla/operation_semantics#conv_convolution for
    reference.

Args:
        inp: an (N+2)-D array. The input of the convolution.
        fltr: an (N+2)-D array. The filter (i.e. kernel) of the convolution.
        window_strides: a sequence of N ints, the strides for moving the convolution
        window.
        padding: a string, either "VALID" or "SAME". The padding algorithm.
        dimension_numbers: a tuple of three strings encoding the data format of
        input, filter and output. "I" means input; "O" means output; "C" means
        channel; other characters such as "W", "H" and "D" means spatial
        dimensions.
        filter_dilation: the dilation rates for the filter. Dilating the filter
        means adding "holes" to the filter.

Returns:
        An (N+2)-D array. The convolution result.


def generate_block_graph(block_graph: blocks.BlockGraph, loader: jinja2.BaseLoader) -> str:
    return _generate_visualization(template_file=_BLOCKGRAPH_TEMPLATE_NAME, loader=loader, graph_data=block_serializer.encode_merged_graph(block_graph))

Generate the visualization webpage.

Args:
        block_graph: blocks.BlockGraph. The block graph of the code.
        loader: A jinja22 loader

Returns:
        str. The rendered visualization page.


def __init__(self):
    self.all_node_names = []
    self.name_node_values = []

A tree visitor that collects the names of all tree nodes into a list.

Attributes:
        all_node_names: collected list of the names, available when the traversal
        is over.
        name_node_values: collects a list of NAME leaves (in addition to those going
        into all_node_names).


def pack_images(batch_images: List[List[np.ndarray]], max_image_tiles: int) -> Tuple[np.ndarray, List[List[int]]]:
    batch_size = len(batch_images)
    max_num_images = max([len(images) for images in batch_images])
    shapes = [image.shape for images in batch_images for image in images]
    _, channels, tile_height, tile_width = shapes[0]
    stacked_images = np.zeros((batch_size, max_num_images, max_image_tiles, channels, tile_height, tile_width), dtype=np.float32)
    all_num_tiles = []
    for i, images in enumerate(batch_images):
        num_sample_tiles = []
        for j, image in enumerate(images):
            num_tiles = image.shape[0]
            stacked_images[i, j, :num_tiles] = image
            num_sample_tiles.append(num_tiles)
        all_num_tiles.append(num_sample_tiles)
    return (stacked_images, all_num_tiles)

Stack a list of lists of images with variable lengths into a numpy array, applying zero padding as needed.
    Each list in the input represents a batch sample, and each image within a list is expected to be
    pre-split into tiles. The resulting array will have a shape of
    (batch_size, max_num_images, max_image_tiles, channels, tile_height, tile_width).

Args:
        batch_images (`List[List[np.ndarray]]`):
            A list of lists of image tiles. Each inner list represents
            a batch sample containing multiple images, where each image is pre-split into tiles.
            The shape of each tile array is (num_tiles, channels, tile_height, tile_width).
        max_image_tiles (int):
            The maximum number of tiles any image was potantially split.

Returns:
        `Tuple[np.ndarray, List[List[int]]]`: A tuple containing:
            - stacked_images (`np.ndarray`):
                A numpy array of stacked images with shape
                (batch_size, max_num_images, max_image_tiles, channels, tile_height, tile_width).
            - all_num_tiles (`List[List[int]]`):
                A list of lists containing the number of tiles
                for each image in each batch sample.


def unnormalize(self, image: np.array, image_mean: Union[float, Iterable[float]], image_std: Union[float, Iterable[float]], input_data_format: Optional[Union[str, ChannelDimension]]=None) -> np.array:
    num_channels = 3
    if isinstance(image_mean, Iterable):
        if len(image_mean) != num_channels:
            raise ValueError(f'mean must have {num_channels} elements if it is an iterable, got {len(image_mean)}')
    else:
        image_mean = [image_mean] * num_channels
    if isinstance(image_std, Iterable):
        if len(image_std) != num_channels:
            raise ValueError(f'std must have {num_channels} elements if it is an iterable, got {len(image_std)}')
    else:
        image_std = [image_std] * num_channels
    rev_image_mean = tuple((-mean / std for mean, std in zip(image_mean, image_std)))
    rev_image_std = tuple((1 / std for std in image_std))
    image = self.normalize(image=image, mean=rev_image_mean, std=rev_image_std, input_data_format=input_data_format)
    return image

Unnormalizes `image` using the mean and standard deviation specified by `mean` and `std`.
    image = (image * image_std) + image_mean

Args:
        image (`torch.Tensor` of shape `(batch_size, num_channels, image_size, image_size)` or `(num_channels, image_size, image_size)`):
            Batch of pixel values to postprocess.
        image_mean (`float` or `Iterable[float]`):
            The mean to use for unnormalization.
        image_std (`float` or `Iterable[float]`):
            The standard deviation to use for unnormalization.
        input_data_format (`ChannelDimension` or `str`, *optional*):
            The channel dimension format for the input image. If unset, the channel dimension format is inferred
            from the input image. Can be one of:
            - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
            - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
            - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.


def _resolve_if_choice_type(fhir_message: message.Message) -> Optional[message.Message]:
    if annotation_utils.is_choice_type(fhir_message):
        choice_field = fhir_message.WhichOneof('choice')
        if choice_field is None:
            return None
        return cast(message.Message, proto_utils.get_value_at_field(fhir_message, choice_field))
    return fhir_message

Resolve to the proper field if given a choice type, return as-is if not.

    Each value in a FHIR choice type is a different field on the protobuf
    representation wrapped under a proto onoeof field.  Therefore, if
    an expression points to a choice type, we should return the populated
    field -- while just returning the field as-is for non-choice types. This
    way we can simply pass nested messages through this class, and return the
    populated item when appropriate.

Args:
        fhir_message: the evaluation result which may or may not be a choice type

Returns:
        The result value, resolved to the sub-field if it is a choice type.


def _from_components(self, components):
    raise NotImplementedError('%s._from_components()' % type(self).__name__)

Reconstructs a value from a nested structure of Tensor/CompositeTensor.

Args:
        components: A nested structure of `tf.Tensor` or `tf.CompositeTensor`,
        compatible with `self._component_specs`.  (Caller is responsible for
        ensuring compatibility.)

Returns:
        A value that is compatible with this `TypeSpec`.


def __init__(self, element_spec):
    self._element_spec = element_spec

Type specification for `tf.data.Iterator`.

    For instance, `tf.data.IteratorSpec` can be used to define a tf.function that
    takes `tf.data.Iterator` as an input argument:

    >>> @tf.function(input_signature=[tf.data.IteratorSpec(
    ...   tf.TensorSpec(shape=(), dtype=tf.int32, name=None))])
    ... def square(iterator):
    ...   x = iterator.get_next()
    ...   return x * x
    >>> dataset = tf.data.Dataset.from_tensors(5)
    >>> iterator = iter(dataset)
    >>> print(square(iterator))
    tf.Tensor(25, shape=(), dtype=int32)

Attributes:
        element_spec: A (nested) structure of `tf.TypeSpec` objects that represents
        the type specification of the iterator elements.


def _convert_dynamic_dimension_to_zero(shape):
    if shape.rank is None:
        return shape
    return tensor_shape.TensorShape([0 if d is None else d for d in shape.as_list()])

Converts dynamic dimensions in `shape` to zero.

    The fake params created to match the intermediates captured in other branches
    could have dynamic dimensions. But the XLA shape is not able to handle
    dynamic dimensions in TF TensorShape. Setting the dynamic dimensions to
    size zero will help avoid failing safety checks in bridge. When XLA
    DynamicConditional op reconciles branch differences, XLA will replace the
    dimension size 0 with a bounded dimension determined from the shape of
    real argument in the other branch.

    Note: Rank unknown shapes are returned as they are.

Args:
        shape: The TensorShape of fake param.

Returns:
        The new TensorShape with dynamic dimensions set to zero.


def _get_indices_and_values(x, jump_locations, values, side, batch_rank):
    indices = tf.searchsorted(jump_locations, x, side=side)
    num_data_points = tf.shape(values)[batch_rank] - 2
    if side == 'right':
        indices_jump = indices - 1
        indices_jump = tf.maximum(indices_jump, 0)
    else:
        indices_jump = tf.minimum(indices, num_data_points)
    value = tf.gather(values, indices, axis=batch_rank, batch_dims=batch_rank)
    jump_location = tf.gather(jump_locations, indices_jump, axis=batch_rank, batch_dims=batch_rank)
    return (value, jump_location, indices_jump)

Computes values and jump locations of the piecewise constant function.

    Given `jump_locations` and the `values` on the corresponding segments of the
    piecewise constant function, the function identifies the nearest jump to `x`
    from the right or left (which is determined by the `side` argument) and the
    corresponding value of the piecewise constant function at `x`

Args:
        x: A real `Tensor` of shape `batch_shape + [num_points]`. Points at which
        the function has to be evaluated.
        jump_locations: A `Tensor` of the same `dtype` as `x` and shape
        `batch_shape + [num_jump_points]`. The locations where the function
        changes its values. Note that the values are expected to be ordered
        along the last dimension.
        values: A `Tensor` of the same `dtype` as `x` and shape
        `batch_shape + [num_jump_points + 1]`. Defines `values[..., i]` on
        `jump_locations[..., i - 1], jump_locations[..., i]`.
        side: A Python string. Whether the function is left- or right- continuous.
        The corresponding values for side should be `left` and `right`.
        batch_rank: A Python scalar stating the batch rank of `x`.

Returns:
        A tuple of three `Tensor` of the same `dtype` as `x` and shapes
        `batch_shape + [num_points] + event_shape`, `batch_shape + [num_points]`,
        and `batch_shape + [num_points] + [2 * len(batch_shape)]`. The `Tensor`s
        correspond to the values, jump locations at `x`, and the corresponding
        indices used to obtain jump locations via `tf.gather_nd`.


def _save_and_write_assets(self, meta_graph_def, assets_list=None):
    write_fn = functools.partial(_add_asset_to_metagraph, meta_graph_def)
    asset_filename_map = _maybe_save_assets(write_fn, assets_list)
    if not asset_filename_map:
        tf_logging.info('No assets to write.')
        return
    copy_assets_to_destination_dir(asset_filename_map, self._export_dir, self._saved_asset_files)

Saves asset to the meta graph and writes asset files to disk.

Args:
        meta_graph_def: The meta graph def to which the assets will be added.
        assets_list: The list where the asset paths are setup.


def __init__(self, sequence, use_multiprocessing=False):
    self.sequence = sequence
    self.use_multiprocessing = use_multiprocessing
    global _SEQUENCE_COUNTER
    if _SEQUENCE_COUNTER is None:
        try:
            _SEQUENCE_COUNTER = multiprocessing.Value('i', 0)
        except OSError:
            _SEQUENCE_COUNTER = 0
    if isinstance(_SEQUENCE_COUNTER, int):
        self.uid = _SEQUENCE_COUNTER
        _SEQUENCE_COUNTER += 1
    else:
        with _SEQUENCE_COUNTER.get_lock():
            self.uid = _SEQUENCE_COUNTER.value
            _SEQUENCE_COUNTER.value += 1
    self.workers = 0
    self.executor_fn = None
    self.queue = None
    self.run_thread = None
    self.stop_signal = None

Base class to enqueue inputs.

    The task of an Enqueuer is to use parallelism to speed up preprocessing.
    This is done with processes or threads.

Example:
    ```python
        enqueuer = SequenceEnqueuer(...)
        enqueuer.start()
        datas = enqueuer.get()
        for data in datas:
            # Use the inputs; training, evaluating, predicting.
            # ... stop sometime.
        enqueuer.stop()
    ```

    The `enqueuer.get()` should be an infinite stream of datas.


def avg_pool(value, ksize, strides, padding, data_format='NHWC', name=None, input=None):
    with ops.name_scope(name, 'AvgPool', [value]) as name:
        value = deprecation.deprecated_argument_lookup('input', input, 'value', value)
        if data_format is None:
            data_format = 'NHWC'
        channel_index = 1 if data_format.startswith('NC') else 3
        ksize = _get_sequence(ksize, 2, channel_index, 'ksize')
        strides = _get_sequence(strides, 2, channel_index, 'strides')
        return gen_nn_ops.avg_pool(value, ksize=ksize, strides=strides, padding=padding, data_format=data_format, name=name)

Performs the average pooling on the input.

    Each entry in `output` is the mean of the corresponding size `ksize`
    window in `value`.

Args:
        value: A 4-D `Tensor` of shape `[batch, height, width, channels]` and type
        `float32`, `float64`, `qint8`, `quint8`, or `qint32`.
        ksize: An int or list of `ints` that has length `1`, `2` or `4`. The size of
        the window for each dimension of the input tensor.
        strides: An int or list of `ints` that has length `1`, `2` or `4`. The
        stride of the sliding window for each dimension of the input tensor.
        padding: A string, either `'VALID'` or `'SAME'`. The padding algorithm.
        See the "returns" section of `tf.nn.convolution` for details.
        data_format: A string. 'NHWC' and 'NCHW' are supported.
        name: Optional name for the operation.
        input: Alias for value.

Returns:
        A `Tensor` with the same type as `value`.  The average pooled output tensor.


def __mul__(self, right: torch.Tensor) -> Rigid:
    if not isinstance(right, torch.Tensor):
        raise TypeError('The other multiplicand must be a Tensor')
    new_rots = self._rots * right
    new_trans = self._trans * right[..., None]
    return Rigid(new_rots, new_trans)

Pointwise left multiplication of the transformation with a tensor. Can be used to e.g. mask the Rigid.

Args:
        right:
            The tensor multiplicand

Returns:
        The product


def __init__(self, *, indent: int=4):
    self._lines: list[_Line] = []
    self._indent_size = indent
    self._indent_lvl = 0

Util to build multi-line text.

    Useful for pretty-print tools and human readable `__repr__`.

Example:
    ```python
    d = {'a': 1, 'b': 2}

    lines = epy.Lines()
    lines += 'dict('
    with lines.indent():
        for k, v in d.items():
        lines += f'{k}={v},'
    lines += ')'
    text = lines.join()
    ```

    Output:

    ```
    dict(
        a=1,
        b=2,
    )
    ```


def custom_apply(self, path: utils.KeyPath, value_spec: pg_typing.ValueSpec, allow_partial: bool, child_transform: Optional[Callable[[utils.KeyPath, pg_typing.Field, Any], Any]]=None) -> Tuple[bool, 'List']:
    proceed_with_standard_apply = True
    if self._value_spec:
        if value_spec and (not value_spec.is_compatible(self._value_spec)):
            raise ValueError(utils.message_on_path(f'List (spec={self._value_spec!r}) cannot be assigned to an incompatible field (spec={value_spec!r}).', path))
        if self._allow_partial == allow_partial:
            proceed_with_standard_apply = False
        else:
            self._allow_partial = allow_partial
    elif isinstance(value_spec, pg_typing.List):
        self._value_spec = value_spec
    return (proceed_with_standard_apply, self)

Implement pg.typing.CustomTyping interface.

Args:
        path: KeyPath of current object.
        value_spec: Origin value spec of the field.
        allow_partial: Whether allow partial object to be created.
        child_transform: Function to transform child node values in dict_obj into
        their final values. Transform function is called on leaf nodes first,
        then on their containers, recursively.

Returns:
        A tuple (proceed_with_standard_apply, transformed value)


def _scale_loss_for_estimator_enabled(self):
    self._scale_loss_for_estimator = True
    try:
        yield
    finally:
        self._scale_loss_for_estimator = False

Scope which sets a flag used for scaling losses in optimizer.

Yields:
        `_scale_loss_for_estimator_enabled` is a context manager with a
        side effect, but doesn't return a value.


def forward(self, input_ids: Optional[torch.Tensor]=None, token_type_ids: Optional[torch.Tensor]=None, position_ids: Optional[torch.Tensor]=None, inputs_embeds: Optional[torch.Tensor]=None, labels: Optional[torch.Tensor]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, **kwargs) -> Union[Tuple, NextSentencePredictorOutput]:
    if 'next_sentence_label' in kwargs:
        warnings.warn('The `next_sentence_label` argument is deprecated and will be removed in a future version, use `labels` instead.', FutureWarning)
        labels = kwargs.pop('next_sentence_label')
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    outputs = self.fnet(input_ids, token_type_ids=token_type_ids, position_ids=position_ids, inputs_embeds=inputs_embeds, output_hidden_states=output_hidden_states, return_dict=return_dict)
    pooled_output = outputs[1]
    seq_relationship_scores = self.cls(pooled_output)
    next_sentence_loss = None
    if labels is not None:
        loss_fct = CrossEntropyLoss()
        next_sentence_loss = loss_fct(seq_relationship_scores.view(-1, 2), labels.view(-1))
    if not return_dict:
        output = (seq_relationship_scores,) + outputs[2:]
        return (next_sentence_loss,) + output if next_sentence_loss is not None else output
    return NextSentencePredictorOutput(loss=next_sentence_loss, logits=seq_relationship_scores, hidden_states=outputs.hidden_states)

labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
    Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair
    (see `input_ids` docstring). Indices should be in `[0, 1]`:

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

Example:
    ```python
    >>> from transformers import AutoTokenizer, FNetForNextSentencePrediction
    >>> import torch

    >>> tokenizer = AutoTokenizer.from_pretrained("google/fnet-base")
    >>> model = FNetForNextSentencePrediction.from_pretrained("google/fnet-base")
    >>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced."
    >>> next_sentence = "The sky is blue due to the shorter wavelength of blue light."
    >>> encoding = tokenizer(prompt, next_sentence, return_tensors="pt")
    >>> outputs = model(**encoding, labels=torch.LongTensor([1]))
    >>> logits = outputs.logits
    >>> assert logits[0, 0] < logits[0, 1]  # next sentence was random
    ```


def sheets_clear(config, auth, sheet_url_or_name, sheet_tab, sheet_range):
    if config.verbose:
        print('SHEETS CLEAR', sheet_url_or_name, sheet_tab, sheet_range)
    sheet_id = sheets_id(config, auth, sheet_url_or_name)
    if sheet_id:
        API_Sheets(config, auth).spreadsheets().values().clear(spreadsheetId=sheet_id, range=sheets_tab_range(sheet_tab, sheet_range), body={}).execute()
    else:
        raise ValueError('Sheet does not exist for %s: %s' % (config, auth, sheet_url_or_name))

Clear a sheet in the specified range.

Args:
        config - see starthinker/util/configuration.py
        auth - user or service
        url_or_name - one of: URL, document title, or id
        sheet_tab - name of tab to get id for
        sheet_range - A1 notation or blank if whole sheet

    No Return


def log_optimal_transport(scores: torch.Tensor, reg_param: torch.Tensor, iterations: int) -> torch.Tensor:
    batch_size, num_rows, num_columns = scores.shape
    one_tensor = scores.new_tensor(1)
    num_rows_tensor, num_columns_tensor = ((num_rows * one_tensor).to(scores), (num_columns * one_tensor).to(scores))
    source_reg_param = reg_param.expand(batch_size, num_rows, 1)
    target_reg_param = reg_param.expand(batch_size, 1, num_columns)
    reg_param = reg_param.expand(batch_size, 1, 1)
    couplings = torch.cat([torch.cat([scores, source_reg_param], -1), torch.cat([target_reg_param, reg_param], -1)], 1)
    log_normalization = -(num_rows_tensor + num_columns_tensor).log()
    log_source_distribution = torch.cat([log_normalization.expand(num_rows), num_columns_tensor.log()[None] + log_normalization])
    log_target_distribution = torch.cat([log_normalization.expand(num_columns), num_rows_tensor.log()[None] + log_normalization])
    log_source_distribution, log_target_distribution = (log_source_distribution[None].expand(batch_size, -1), log_target_distribution[None].expand(batch_size, -1))
    log_optimal_transport_matrix = log_sinkhorn_iterations(couplings, log_source_distribution, log_target_distribution, num_iterations=iterations)
    log_optimal_transport_matrix = log_optimal_transport_matrix - log_normalization
    return log_optimal_transport_matrix

Perform Differentiable Optimal Transport in Log-space for stability

Args:
        scores: (`torch.Tensor` of shape `(batch_size, num_rows, num_columns)`):
            Cost matrix.
        reg_param: (`torch.Tensor` of shape `(batch_size, 1, 1)`):
            Regularization parameter.
        iterations: (`int`):
            Number of Sinkhorn iterations.

Returns:
        log_optimal_transport_matrix: (`torch.Tensor` of shape `(batch_size, num_rows, num_columns)`): Logarithm of the
        optimal transport matrix.


def convert_to_tensors(self, tensor_type: Optional[Union[str, TensorType]]=None, prepend_batch_axis: bool=False):
    if tensor_type is None:
        return self
    if not isinstance(tensor_type, TensorType):
        tensor_type = TensorType(tensor_type)
    if tensor_type == TensorType.TENSORFLOW:
        if not is_tf_available():
            raise ImportError('Unable to convert output to TensorFlow tensors format, TensorFlow is not installed.')
        import tensorflow as tf
        as_tensor = tf.constant
        is_tensor = tf.is_tensor
    elif tensor_type == TensorType.PYTORCH:
        if not is_torch_available():
            raise ImportError('Unable to convert output to PyTorch tensors format, PyTorch is not installed.')
        import torch
        is_tensor = torch.is_tensor

        def as_tensor(value, dtype=None):
            if isinstance(value, list) and isinstance(value[0], np.ndarray):
                return torch.from_numpy(np.array(value))
            return torch.tensor(value)
    elif tensor_type == TensorType.JAX:
        if not is_flax_available():
            raise ImportError('Unable to convert output to JAX tensors format, JAX is not installed.')
        import jax.numpy as jnp
        as_tensor = jnp.array
        is_tensor = is_jax_tensor
    elif tensor_type == TensorType.MLX:
        if not is_mlx_available():
            raise ImportError('Unable to convert output to MLX tensors format, MLX is not installed.')
        import mlx.core as mx
        as_tensor = mx.array

        def is_tensor(obj):
            return isinstance(obj, mx.array)
    else:

        def as_tensor(value, dtype=None):
            if isinstance(value, (list, tuple)) and isinstance(value[0], (list, tuple, np.ndarray)):
                value_lens = [len(val) for val in value]
                if len(set(value_lens)) > 1 and dtype is None:
                    value = as_tensor([np.asarray(val) for val in value], dtype=object)
            return np.asarray(value, dtype=dtype)
        is_tensor = is_numpy_array
    for key, value in self.items():
        try:
            if prepend_batch_axis:
                value = [value]
            if not is_tensor(value):
                tensor = as_tensor(value)
                self[key] = tensor
        except Exception as e:
            if key == 'overflowing_tokens':
                raise ValueError('Unable to create tensor returning overflowing tokens of different lengths. Please see if a fast version of this tokenizer is available to have this feature available.') from e
            raise ValueError(f"Unable to create tensor, you should probably activate truncation and/or padding with 'padding=True' 'truncation=True' to have batched tensors with the same length. Perhaps your features (`{key}` in this case) have excessive nesting (inputs type `list` where type `int` is expected).") from e
    return self

Convert the inner content to tensors.

Args:
        tensor_type (`str` or [`~utils.TensorType`], *optional*):
            The type of tensors to use. If `str`, should be one of the values of the enum [`~utils.TensorType`]. If
            `None`, no modification is done.
        prepend_batch_axis (`int`, *optional*, defaults to `False`):
            Whether or not to add the batch dimension during the conversion.


def __init__(self, threshold_fn_spec: Spec):
    assert isinstance(threshold_fn_spec.config, dict)
    threshold_fn_spec.config['_run_init'] = True
    threshold_fn = Specifiable.from_spec(threshold_fn_spec)
    assert isinstance(threshold_fn, ThresholdFn)
    assert threshold_fn.is_stateful, 'This DoFn can only take stateful function as threshold_fn'
    self._threshold_fn_spec = threshold_fn_spec

Applies a stateful ThresholdFn to anomaly detection results.

    This DoFn is designed for stateful `ThresholdFn` implementations. It leverages
    Beam's state management to persist and update the state of the `ThresholdFn`
    across multiple elements. This is necessary for `ThresholdFn`s that need to
    accumulate information or adapt over time, such as quantile-based thresholds.

Args:
        threshold_fn_spec (Spec): Specification defining the `ThresholdFn` to be
        used.

Raises:
        AssertionError: If the provided `threshold_fn_spec` leads to the
        creation of a stateless `ThresholdFn`.


def group(self, group_id):
    self._validate_group_id(group_id)
    return self._Context(self, group_id)

Enter a context where the lock is with group `group_id`.

Args:
        group_id: The group for which to acquire and release the lock.

Returns:
        A context manager which will acquire the lock for `group_id`.


def read_fingerprint(export_dir: str) -> Fingerprint:
    try:
        fingerprint = fingerprinting_pywrap.ReadSavedModelFingerprint(export_dir)
    except fingerprinting_pywrap.FileNotFoundException as e:
        raise FileNotFoundError(f'SavedModel Fingerprint Error: {e}') from None
    except fingerprinting_pywrap.FingerprintException as e:
        raise RuntimeError(f'SavedModel Fingerprint Error: {e}') from None
    return Fingerprint.from_proto(fingerprint_pb2.FingerprintDef().FromString(fingerprint))

Reads the fingerprint of a SavedModel in `export_dir`.

    Returns a `tf.saved_model.experimental.Fingerprint` object that contains
    the values of the SavedModel fingerprint, which is persisted on disk in the
    `fingerprint.pb` file in the `export_dir`.

    Read more about fingerprints in the SavedModel guide at
    https://www.tensorflow.org/guide/saved_model.

Args:
        export_dir: The directory that contains the SavedModel.

Returns:
        A `tf.saved_model.experimental.Fingerprint`.

Raises:
        FileNotFoundError: If no or an invalid fingerprint is found.


def __init__(self, input_dataset, target_device, source_device='/cpu:0'):
    self._input_dataset = input_dataset._apply_debug_options()
    self._target_device = target_device
    spec = framework_device.DeviceSpec().from_string(self._target_device)
    self._is_gpu_target = spec.device_type == 'GPU'
    self._source_device_string = source_device
    self._source_device = ops.convert_to_tensor(source_device)
    wrap_ds_variant = gen_dataset_ops.wrap_dataset_variant(self._input_dataset._variant_tensor)

    @def_function.function()
    def _init_func():
        """Creates an iterator for the input dataset.

      Returns:
        A `string` tensor that encapsulates the iterator created.
      """
        ds_variant = gen_dataset_ops.unwrap_dataset_variant(wrap_ds_variant)
        resource = gen_dataset_ops.anonymous_iterator(**self._input_dataset._flat_structure)
        with ops.control_dependencies([gen_dataset_ops.make_iterator(ds_variant, resource)]):
            return gen_dataset_ops.iterator_to_string_handle(resource)
    init_func_concrete = _init_func.get_concrete_function()

    @def_function.function()
    def _remote_init_func():
        return functional_ops.remote_call(target=self._source_device, args=init_func_concrete.captured_inputs, Tout=[dtypes.string], f=init_func_concrete)
    self._init_func = _remote_init_func.get_concrete_function()
    self._init_captured_args = self._init_func.captured_inputs

    @def_function.function(input_signature=[tensor_spec.TensorSpec([], dtypes.string)])
    def _next_func(string_handle):
        """Calls get_next for created iterator.

      Args:
        string_handle: An iterator string handle created by _init_func
      Returns:
        The elements generated from `input_dataset`
      """
        with ops.device(self._source_device_string):
            iterator = iterator_ops.Iterator.from_string_handle(string_handle, dataset_ops.get_legacy_output_types(self), dataset_ops.get_legacy_output_shapes(self), dataset_ops.get_legacy_output_classes(self))
        return structure.to_tensor_list(self.element_spec, iterator.get_next())
    next_func_concrete = _next_func.get_concrete_function()

    @def_function.function(input_signature=[tensor_spec.TensorSpec([], dtypes.string)], experimental_attributes={'experimental_ints_on_device': True})
    def _remote_next_func(string_handle):
        return functional_ops.remote_call(target=self._source_device, args=[string_handle] + next_func_concrete.captured_inputs, Tout=self._input_dataset._flat_types, f=next_func_concrete)
    self._next_func = _remote_next_func.get_concrete_function()
    self._next_captured_args = self._next_func.captured_inputs

    @def_function.function(input_signature=[tensor_spec.TensorSpec([], dtypes.string)])
    def _finalize_func(string_handle):
        """Destroys the iterator resource created.

      Args:
        string_handle: An iterator string handle created by _init_func
      Returns:
        Tensor constant 0
      """
        iterator_resource = gen_dataset_ops.iterator_from_string_handle_v2(string_handle, **self._input_dataset._flat_structure)
        with ops.control_dependencies([resource_variable_ops.destroy_resource_op(iterator_resource, ignore_lookup_error=True)]):
            return array_ops.constant(0, dtypes.int64)
    finalize_func_concrete = _finalize_func.get_concrete_function()

    @def_function.function(input_signature=[tensor_spec.TensorSpec([], dtypes.string)])
    def _remote_finalize_func(string_handle):
        return functional_ops.remote_call(target=self._source_device, args=[string_handle] + finalize_func_concrete.captured_inputs, Tout=[dtypes.int64], f=finalize_func_concrete)
    self._finalize_func = _remote_finalize_func.get_concrete_function()
    self._finalize_captured_args = self._finalize_func.captured_inputs
    g = ops.get_default_graph()
    self._init_func.add_to_graph(g)
    self._next_func.add_to_graph(g)
    self._finalize_func.add_to_graph(g)
    with ops.device(self._target_device):
        variant_tensor = gen_dataset_ops.generator_dataset(self._init_captured_args, self._next_captured_args, self._finalize_captured_args, init_func=self._init_func, next_func=self._next_func, finalize_func=self._finalize_func, **self._input_dataset._flat_structure)
    super(_CopyToDeviceDataset, self).__init__(input_dataset, variant_tensor)

Constructs a _CopyToDeviceDataset.

Args:
        input_dataset: `Dataset` to be copied
        target_device: The name of the device to which elements would be copied.
        source_device: Device where input_dataset would be placed.


def __init__(self, conf, conn=None):
    super(HttpFilesSource, self).__init__(conf)
    self._SetDefaults(conf)
    if not conn:
        conn = pycurl.Curl()
        conn.setopt(pycurl.NOPROGRESS, 1)
        conn.setopt(pycurl.NOSIGNAL, 1)
        conn.setopt(pycurl.TIMEOUT, 60)
        conn.setopt(pycurl.USERAGENT, 'nsscache')
        if self.conf['http_proxy']:
            conn.setopt(pycurl.PROXY, self.conf['http_proxy'])
    self.conn = conn

Initialise the HTTP Data Source.

Args:
        conf: config.Config instance
        conn: pycurl Curl object


def print_generic_type(self, t: types.BaseValue) -> str:

Returns a string of the generic type of t.

    For example, if t is `[0]`, then this method returns "list[int]".

Args:
        t: An abstract value.


def to_numpy(evset: EventSet, timestamp_to_datetime: bool=True, timestamps: bool=True) -> Dict[str, ndarray]:
    timestamp_key = 'timestamp'
    index_names = evset.schema.index_names()
    feature_names = evset.schema.feature_names()
    column_names = index_names + feature_names
    if timestamps:
        column_names += [timestamp_key]
    dst = {column_name: [] for column_name in column_names}
    for index, data in evset.data.items():
        assert isinstance(index, tuple)
        if timestamps:
            if evset.schema.is_unix_timestamp and timestamp_to_datetime:
                dst[timestamp_key].append(data.timestamps.astype('datetime64[s]'))
            else:
                dst[timestamp_key].append(data.timestamps)
        for feature_name, feature in zip(feature_names, data.features):
            dst[feature_name].append(feature)
        num_timestamps = len(data.timestamps)
        for index_name, index_item in zip(index_names, index):
            dst[index_name].append(np.repeat(index_item, num_timestamps))
    dst = {k: np.concatenate(v) for k, v in dst.items()}
    return dst

Converts an [`EventSet`][temporian.EventSet] to a flattened dictionary with
    numpy arrays.

    Usage example:
        ```python
        >>> from datetime import datetime

        >>> evset = tp.event_set(
        ...     timestamps=['2023-11-08T17:14:38', '2023-11-29T21:44:46'],
        ...     features={
        ...         "store": ['STORE_1', 'STORE_2'],
        ...         "revenue": [1571, 6101]
        ...     },
        ...     indexes=["store"],
        ... )

        # Timestamps are exported as datetime64[s] if they were created as datetimes,
        # otherwhise they are floats
        >>> res = tp.to_numpy(evset)
        >>> res
        {'store': array([b'STORE_2', b'STORE_1'], dtype='|S7'), 'revenue': array([6101, 1571]),
        'timestamp': array(['2023-11-29T21:44:46', '2023-11-08T17:14:38'], dtype='datetime64[s]')}

        ```

Args:
        evset: input event set.
        timestamp_to_datetime: If true, cast Temporian timestamps to datetime64
            when is_unix_timestamp is set to True.
        timestamps: If true, the timestamps are included as a column.

Returns:
        object with numpy arrays created from EventSet.


def simple_two_objective_reward(example):
    num = int(example * 10) % 9 + 1
    return [num, 10 - num]

Reward for the trivial search space.

    The reward (i.e. fitness) is a 2-element list. The goal of the search,
    therefore, is to find the pareto frontier in simple_two_objective_pareto
    function.

Args:
        example: a materialized value.

Returns:
        A 2-element list.


def abort_class_if(expr, reason, extras=None):
    if expr:
        abort_class(reason, extras)

Abort all subsequent tests within the same test class in one iteration,
    if expression evaluates to True.

    If one test class is requested multiple times in a test run, this can
    only abort one of the requested executions, NOT all.

Args:
        expr: The expression that is evaluated.
        reason: The reason to abort.
        extras: An optional field for extra information to be included in
        test result.

Raises:
        signals.TestAbortClass: Abort all subsequent tests in a test class.


def get_current_human_time():
    return time.strftime('%m-%d-%Y %H:%M:%S ')

Returns the current time in human readable format.

Returns:
        The current time stamp in Month-Day-Year Hour:Min:Sec format.


def compute_gradient_error(x, x_shape, y, y_shape, x_init_value=None, delta=0.001, init_targets=None, extra_feed_dict=None):
    grad = compute_gradient(x, x_shape, y, y_shape, x_init_value, delta, init_targets, extra_feed_dict=extra_feed_dict)
    return _compute_error(grad)

Computes the gradient error.

    Computes the maximum error for dy/dx between the computed Jacobian and the
    numerically estimated Jacobian.

    This function will modify the tensors passed in as it adds more operations
    and hence changing the consumers of the operations of the input tensors.

    This function adds operations to the current session. To compute the error
    using a particular device, such as a GPU, use the standard methods for
    setting a device (e.g. using with sess.graph.device() or setting a device
    function in the session constructor).

Args:
        x: a tensor or list of tensors
        x_shape: the dimensions of x as a tuple or an array of ints. If x is a list,
        then this is the list of shapes.
        y: a tensor
        y_shape: the dimensions of y as a tuple or an array of ints.
        x_init_value: (optional) a numpy array of the same shape as "x"
        representing the initial value of x. If x is a list, this should be a list
        of numpy arrays.  If this is none, the function will pick a random tensor
        as the initial value.
        delta: (optional) the amount of perturbation.
        init_targets: list of targets to run to initialize model params.
        extra_feed_dict: dict that allows fixing specified tensor values
        during the Jacobian calculation.

Returns:
        The maximum error in between the two Jacobians.


def unravel_index(indices, shape):
    if any_symbolic_tensors((indices,)):
        return UnravelIndex(shape).symbolic_call(indices)
    return backend.numpy.unravel_index(indices, shape)

Convert flat indices to coordinate arrays in a given array shape.

Args:
        indices: An integer or array of integers representing flat indices.
        shape: The shape of the array to unravel into.

Returns:
        Tuple of arrays for each dimension with unraveled indices.

Example:
        >>> indices = 5
        >>> shape = (3, 3)
        >>> unravel_index(indices, shape)
        (1, 2)  # 5 is at row 1, column 2 in a 3x3 array


def sample_g_values(self, ngram_keys: torch.LongTensor) -> torch.LongTensor:
    sampling_table_size, = self.sampling_table.shape
    sampling_table = self.sampling_table.reshape((1, 1, sampling_table_size))
    ngram_keys = ngram_keys % sampling_table_size
    return torch.take_along_dim(sampling_table, indices=ngram_keys, dim=2)

Samples g values from Bernoulli distribution.

    It is not possible to pass random keys in a vectorized way in torch. Instead
    we pre-compute a random sampling table, and use apply modulo table size to
    map from ngram keys (int64) to g values.

Args:
        ngram_keys (`torch.LongTensor`):
            Random keys (batch_size, num_ngrams, depth).

Returns:
        G values (batch_size, num_ngrams, depth).


def forward(self, pixel_values: torch.FloatTensor, pixel_attention_mask: torch.Tensor, spatial_shapes: torch.LongTensor, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None) -> BaseModelOutputWithPooling:
    return self.vision_model(pixel_values=pixel_values, attention_mask=pixel_attention_mask, spatial_shapes=spatial_shapes, output_attentions=output_attentions, output_hidden_states=output_hidden_states)

pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*):
    Mask to avoid performing attention on padding pixel indices.
    spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`):
    Tensor containing the spatial dimensions (height, width) of the input images.

Example:
    ```python
    >>> from PIL import Image
    >>> import requests
    >>> from transformers import AutoProcessor, Siglip2VisionModel

    >>> model = Siglip2VisionModel.from_pretrained("google/siglip2-base-patch16-224")
    >>> processor = AutoProcessor.from_pretrained("google/siglip2-base-patch16-224")

    >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
    >>> image = Image.open(requests.get(url, stream=True).raw)

    >>> inputs = processor(images=image, return_tensors="pt")

    >>> outputs = model(**inputs)
    >>> last_hidden_state = outputs.last_hidden_state
    >>> pooled_output = outputs.pooler_output  # pooled features
    ```


def graph_defs_equal(graph_def_1: graph_pb2.GraphDef, graph_def_2: graph_pb2.GraphDef, treat_nan_as_equal: bool=False) -> bool:
    if not isinstance(graph_def_1, graph_pb2.GraphDef):
        raise TypeError(f'graph_def_1 must be a graph_pb2.GraphDef proto, but got type {type(graph_def_1)}.')
    if not isinstance(graph_def_2, graph_pb2.GraphDef):
        raise TypeError(f'graph_def_2 must be a graph_pb2.GraphDef proto, but got type {type(graph_def_2)}.')
    options = _proto_comparators.ProtoComparisonOptions(treat_nan_as_equal)
    return _proto_comparators.EqualsGraphDef(graph_def_1.SerializeToString(), graph_def_2.SerializeToString(), options)

Returns True iff the graph def arguments are structurally equivalent.

    The notion of equivalence encoded here checks that the set of NodeDefs in
    the GraphDef's function library and main graph body are identical.
    Additionally, it checks that the functions in the function library are equal
    as sets.

    Example usage:

    ```
    with tf.Graph().as_default() as g1:
        tf.constant(1)

    with tf.Graph().as_default() as g2:
        tf.constant(2)

    with tf.Graph().as_default() as g3:
        tf.constant(1)

    assert tf.__internal__.graph_util.graph_defs_equal(g1.as_graph_def(),
                                                    g3.as_graph_def())

    assert not tf.__internal__.graph_util.graph_defs_equal(g1.as_graph_def(),
                                                        g2.as_graph_def())
    ```

Args:
        graph_def_1: Instance of `graph_pb2.GraphDef` to compare.
        graph_def_2: Instance of `graph_pb2.GraphDef` to compare.
        treat_nan_as_equal: Boolean indicating whether or not to treat nan
        floating-point values as equal. This is crucial for any equivalence
        relation defined over GraphDefs, to ensure symmetry.

Returns:
        Boolean indicating structural equivalence as described above.

Raises:
        TypeError: If either of the GraphDefs are not instances of
        `graph_pb2.GraphDef`.


def assertProtoEqual(self, a, b, check_initialized=True, normalize_numbers=False, msg=None, relative_tolerance=None):
    pool = descriptor_pool.Default()
    if isinstance(a, str):
        a = text_format.Parse(a, b.__class__(), descriptor_pool=pool)
    for pb in (a, b):
        if check_initialized:
            errors = pb.FindInitializationErrors()
            if errors:
                self.fail('Initialization errors: %s\n%s' % (errors, pb))
        if normalize_numbers:
            NormalizeNumberFields(pb)
    if relative_tolerance is not None:
        checkFloatEqAndReplace(self, expected=b, actual=a, relative_tolerance=relative_tolerance)
    a_str = text_format.MessageToString(a, descriptor_pool=pool)
    b_str = text_format.MessageToString(b, descriptor_pool=pool)
    if len(a_str) < 2 ** 16 and len(b_str) < 2 ** 16:
        self.assertMultiLineEqual(a_str, b_str, msg=msg)
    else:
        diff = ''.join(difflib.unified_diff(a_str.splitlines(True), b_str.splitlines(True)))
        if diff:
            self.fail('%s :\n%s' % (msg, diff))

Fails with a useful error if a and b aren't equal.

    Comparison of repeated fields matches the semantics of
    unittest.TestCase.assertEqual(), ie order and extra duplicates fields matter.

Args:
        self: googletest.TestCase
        a: proto2 PB instance, or text string representing one.
        b: proto2 PB instance -- message.Message or subclass thereof.
        check_initialized: boolean, whether to fail if either a or b isn't
        initialized.
        normalize_numbers: boolean, whether to normalize types and precision of
        numbers before comparison.
        msg: if specified, is used as the error message on failure.
        relative_tolerance: float, relative tolerance. If this is not provided, then
        all floats are compared using string comparison otherwise, floating point
        comparisons are done using the relative tolerance provided.


def get_top_docs(self, question_hidden_states: np.ndarray, n_docs=5) -> Tuple[np.ndarray, np.ndarray]:
    raise NotImplementedError

For each query in the batch, retrieves `n_docs` documents.

Args:
        question_hidden_states (`np.ndarray` of shape `(batch_size, vector_size)`):
            An array of query vectors.
        n_docs (`int`):
            The number of docs retrieved per query.

Returns:
        `np.ndarray` of shape `(batch_size, n_docs)`: A tensor of indices of retrieved documents. `np.ndarray` of
        shape `(batch_size, vector_size)`: A tensor of vector representations of retrieved documents.


def set_signature(self, signature_map: signature_serialization._SignatureMap, wrapped_functions: Dict[Callable[..., Any], Callable[..., Any]]):
    self.list_children(self.root)
    name = signature_serialization.SIGNATURE_ATTRIBUTE_NAME
    self._children_cache[self.root][name] = signature_map
    self._wrapped_functions.update(wrapped_functions)

Attach signature to the root object.

Args:
        signature_map: An object that contains signature functions.
        wrapped_functions: A dictionary mapping functions to functions that are
        guaranteed to not capture cached variables (functions that capture
        cached variables can't be saved).


def arctan2(x1, x2):
    if any_symbolic_tensors((x1, x2)):
        return Arctan2().symbolic_call(x1, x2)
    return backend.numpy.arctan2(x1, x2)

Element-wise arc tangent of `x1/x2` choosing the quadrant correctly.

    The quadrant (i.e., branch) is chosen so that `arctan2(x1, x2)` is the
    signed angle in radians between the ray ending at the origin and passing
    through the point `(1, 0)`, and the ray ending at the origin and passing
    through the point `(x2, x1)`. (Note the role reversal: the "y-coordinate"
    is the first function parameter, the "x-coordinate" is the second.) By IEEE
    convention, this function is defined for `x2 = +/-0` and for either or both
    of `x1` and `x2` `= +/-inf`.

Args:
        x1: First input tensor.
        x2: Second input tensor.

Returns:
        Tensor of angles in radians, in the range `[-pi, pi]`.

Example:
    Consider four points in different quadrants:
    >>> x = keras.ops.convert_to_tensor([-1, +1, +1, -1])
    >>> y = keras.ops.convert_to_tensor([-1, -1, +1, +1])
    >>> keras.ops.arctan2(y, x) * 180 / numpy.pi
    array([-135., -45., 45., 135.], dtype=float32)

    Note the order of the parameters. `arctan2` is defined also when x2=0 and
    at several other points, obtaining values in the range `[-pi, pi]`:
    >>> keras.ops.arctan2(
    ...     keras.ops.array([1., -1.]),
    ...     keras.ops.array([0., 0.]),
    ... )
    array([ 1.5707964, -1.5707964], dtype=float32)
    >>> keras.ops.arctan2(
    ...     keras.ops.array([0., 0., numpy.inf]),
    ...     keras.ops.array([+0., -0., numpy.inf]),
    ... )
    array([0., 3.1415925, 0.7853982], dtype=float32)


def _validate_paths_for_modules(self, actual_paths, expected_paths, file_to_update_on_error):
    self.assertTrue(actual_paths)
    self.assertTrue(expected_paths)
    missing_paths = expected_paths - actual_paths
    extra_paths = actual_paths - expected_paths
    missing_paths = ["'%s'" % path for path in missing_paths]
    extra_paths = ["'%s'" % path for path in extra_paths]
    self.assertFalse(missing_paths, 'Please add %s to %s.' % (',\n'.join(sorted(missing_paths)), file_to_update_on_error))
    self.assertFalse(extra_paths, 'Redundant paths, please remove %s in %s.' % (',\n'.join(sorted(extra_paths)), file_to_update_on_error))

Validates that actual_paths match expected_paths.

Args:
        actual_paths: */__init__.py file paths listed in file_to_update_on_error.
        expected_paths: */__init__.py file paths that we need to create for
        TensorFlow API.
        file_to_update_on_error: File that contains list of */__init__.py files.
        We include it in error message printed if the file list needs to be
        updated.


def __init__(self, ast, visit_decorators=True):
    self._ast = ast
    maybe_decorators = ['decorator_list'] if visit_decorators else []
    self._node_children = {self._ast.Module: ['body'], self._ast.ClassDef: maybe_decorators + ['keywords', 'bases', 'body'], self._ast.FunctionDef: maybe_decorators + ['body', 'args', 'returns'], self._ast.Assign: ['targets', 'value']}

A base class for writing AST visitors.

    Subclasses should define {visit,enter,leave}_X to process nodes of type X.
    If a visit method returns a non-None value, the visited node is replaced
    with that value.

Attributes:
        _ast: Any module whose interface matches the standard ast library, such as
        typed_ast. The same module must be used to generate the AST to visit.


def run(argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult:
    known_args, pipeline_args = parse_known_args(argv)
    pipeline_options = PipelineOptions(pipeline_args)
    pipeline_options.view_as(SetupOptions).save_main_session = save_main_session
    pipeline = test_pipeline
    if not test_pipeline:
        pipeline = beam.Pipeline(options=pipeline_options)
    model_handler = HuggingFacePipelineModelHandler(task=PipelineTask.QuestionAnswering, model=known_args.model_name, load_model_args={'framework': 'pt', 'revision': known_args.revision})
    if not known_args.input:
        text = pipeline | 'CreateSentences' >> beam.Create(['What does Apache Beam do?;Apache Beam enables batch and streaming data processing.', 'What is the capital of France?;The capital of France is Paris .', 'Where was beam summit?;Apache Beam Summit 2023 was in NYC.'])
    else:
        text = pipeline | 'ReadSentences' >> beam.io.ReadFromText(known_args.input)
    processed_text = text | 'PreProcess' >> beam.ParDo(preprocess) | 'SquadExample' >> beam.ParDo(create_squad_example)
    output = processed_text | 'RunInference' >> RunInference(KeyedModelHandler(model_handler)) | 'ProcessOutput' >> beam.ParDo(PostProcessor())
    _ = output | 'WriteOutput' >> beam.io.WriteToText(known_args.output, shard_name_template='', append_trailing_newlines=True)
    result = pipeline.run()
    result.wait_until_finish()
    return result

Args:
    argv: Command line arguments defined for this example.
    save_main_session: Used for internal testing.
    test_pipeline: Used for internal testing.


def save_counter(self):
    return self.checkpointer().save_counter

An integer variable numbering the checkpoint events.

    This is maintained by the underlying tf.train.Checkpoint object employed by
    AsyncCheckpoint class. The number starts at 0 and gets incremented for each
    checkpoint event.

Returns:
        The save counter variable.


def _get_key_counter(seed, alg):
    if alg == Algorithm.AUTO_SELECT.value:
        key, counter = gen_stateless_random_ops_v2.stateless_random_get_key_counter(seed)
    elif alg == Algorithm.PHILOX.value:
        key, counter = _philox_scramble_seed(seed)
    elif alg == Algorithm.THREEFRY.value:
        key = array_ops.reshape(_uint32s_to_uint64(math_ops.cast(seed, dtypes.uint32)), [1])
        counter = array_ops.zeros([1], dtypes.uint64)
    else:
        raise ValueError(unsupported_alg_error_msg(alg))
    return (key, counter)

Calculates the key and counter to pass to raw RNG ops.

    This function calculates the key and counter that will be passed to
    the raw RNG ops like `StatelessRandomUniformV2`. Depending on the
    input `alg`, the key and counter may be scrambled or copied from
    `seed`. If `alg` is `"auto_select"`, the key and counter will be
    determined at runtime based on device type.

Args:
        seed: An integer tensor of shape [2]. The seed to calculate the key and
        counter from.
        alg: The RNG algorithm. See `tf.random.stateless_uniform` for an
        explanation.

Returns:
        A pair (key, counter) suitable for V2 stateless RNG ops like
        `StatelessRandomUniformV2`.


def convert_ids_to_tokens(self, ids: Union[int, list[int]], skip_special_tokens: bool=False) -> Union[str, list[str]]:
    if isinstance(ids, int):
        if ids in self._added_tokens_decoder:
            return self._added_tokens_decoder[ids].content
        else:
            return self._convert_id_to_token(ids)
    tokens = []
    for index in ids:
        index = int(index)
        if skip_special_tokens and index in self.all_special_ids:
            continue
        if index in self._added_tokens_decoder:
            tokens.append(self._added_tokens_decoder[index].content)
        else:
            tokens.append(self._convert_id_to_token(index))
    return tokens

Converts a single index or a sequence of indices in a token or a sequence of tokens, using the vocabulary and
    added tokens.

Args:
        ids (`int` or `List[int]`):
            The token id (or token ids) to convert to tokens.
        skip_special_tokens (`bool`, *optional*, defaults to `False`):
            Whether or not to remove special tokens in the decoding.

Returns:
        `str` or `List[str]`: The decoded token(s).


def process_data_fn(self, inputs: dict[str, common_types.ConsistentTensorType]) -> dict[str, common_types.ConsistentTensorType]:
    outputs = inputs.copy()
    for transform in self.transforms:
        columns = transform.columns
        for col in columns:
            intermediate_result = transform(outputs[col], output_column_name=col)
            for key, value in intermediate_result.items():
                outputs[key] = value
    return outputs

This method is used in the AnalyzeAndTransformDataset step. It applies
    the transforms to the `inputs` in sequential order on the columns
    provided for a given transform.

Args:
        inputs: A dictionary of column names and data.

Returns:
        A dictionary of column names and transformed data.


def __init__(self, reduction='sum_over_batch_size', name='mean_absolute_error', dtype=None):
    super().__init__(mean_absolute_error, name=name, reduction=reduction, dtype=dtype)

Computes the mean of absolute difference between labels and predictions.

    Formula:

    ```python
    loss = mean(abs(y_true - y_pred))
    ```

Args:
        reduction: Type of reduction to apply to the loss. In almost all cases
            this should be `"sum_over_batch_size"`. Supported options are
            `"sum"`, `"sum_over_batch_size"`, `"mean"`,
            `"mean_with_sample_weight"` or `None`. `"sum"` sums the loss,
            `"sum_over_batch_size"` and `"mean"` sum the loss and divide by the
            sample size, and `"mean_with_sample_weight"` sums the loss and
            divides by the sum of the sample weights. `"none"` and `None`
            perform no aggregation. Defaults to `"sum_over_batch_size"`.
        name: Optional name for the loss instance.
        dtype: The dtype of the loss's computations. Defaults to `None`, which
            means using `keras.backend.floatx()`. `keras.backend.floatx()` is a
            `"float32"` unless set to different value
            (via `keras.backend.set_floatx()`). If a `keras.DTypePolicy` is
            provided, then the `compute_dtype` will be utilized.


def VisitFunction(self, f):
    groups = self._GroupByArguments(f.signatures)
    new_signatures = []
    for stripped_signature, ret_exc in groups.items():
        ret = pytd_utils.JoinTypes(ret_exc.return_types)
        exc = tuple(ret_exc.exceptions)
        new_signatures.append(stripped_signature.Replace(return_type=ret, exceptions=exc))
    return f.Replace(signatures=tuple(new_signatures))

Merge signatures of a function.

    This groups signatures by arguments and then for each group creates a
    single signature that joins the return values / exceptions using "or".

Args:
        f: A pytd.Function instance

Returns:
        Function with simplified / combined signatures.


def saveable_objects_for_op(op, name):
    if not isinstance(name, str):
        raise TypeError(f'names_to_saveables must be a dict mapping string names to trackable operations. Name is not a string: {name}')
    if isinstance(op, saveable_object.SaveableObject):
        yield op
    elif isinstance(op, (list, tuple, variables.PartitionedVariable)):
        if isinstance(op, variables.PartitionedVariable):
            op = list(op)
        slice_name = None
        for variable in op:
            if isinstance(variable, saveable_object.SaveableObject):
                yield variable
                continue
            if not isinstance(variable, variables.Variable):
                raise ValueError(f'Slices must all be Variables: {variable}')
            if not variable._save_slice_info:
                raise ValueError(f'Slices must all be slices: {variable}')
            if slice_name is None:
                slice_name = variable._save_slice_info.full_name
            elif slice_name != variable._save_slice_info.full_name:
                raise ValueError(f'Slices must all be from the same tensor: {slice_name} != {variable._save_slice_info.full_name}')
            if variable.op.type in _REF_VARIABLE_OPS:
                yield ReferenceVariableSaveable(variable, variable._save_slice_info.spec, name)
            else:
                yield ResourceVariableSaveable(variable, variable._save_slice_info.spec, name)
    elif isinstance(op, trackable.Trackable) and (not isinstance(op, variables.Variable)):
        for attr, factory in saveable_objects_from_trackable(op, tf1_saver=True).items():
            if attr == trackable.VARIABLE_VALUE_KEY:
                full_name = name
            elif attr == trackable_utils.SERIALIZE_TO_TENSORS_NAME:
                full_name = name
            else:
                full_name = name + '_' + attr
            op = factory(full_name) if callable(factory) else factory
            for op in saveable_objects_for_op(op, op.name):
                yield op
    elif isinstance(op, resource_variable_ops.BaseResourceVariable):
        if op._in_graph_mode:
            variable = op._graph_element
        else:
            variable = op
        yield ResourceVariableSaveable(variable, '', name)
    else:
        if context.executing_eagerly():
            raise ValueError(f'Can only save/restore ResourceVariables when executing eagerly, got type: {type(op)}.')
        variable = ops.convert_to_tensor(op, as_ref=True)
        if not _tensor_comes_from_variable(variable):
            raise TypeError(f'names_to_saveables must be a dict mapping string names to Tensors/Variables. Not a variable: {variable}')
        if variable.op.type in _REF_VARIABLE_OPS:
            yield ReferenceVariableSaveable(variable, '', name)
        else:
            yield ResourceVariableSaveable(variable, '', name)

Create `SaveableObject`s from an operation.

Args:
        op: A variable, operation, or SaveableObject to coerce into a
        SaveableObject.
        name: A string name for the SaveableObject.

Yields:
        `SaveableObject`s which together save/restore `op`.

Raises:
        TypeError: If `name` is not a string.
        ValueError: For operations with no known conversion to SaveableObject.


def _should_stop(state, stopping_policy_fn):
    return tf.convert_to_tensor(stopping_policy_fn(state.finished), name='should_stop', dtype=tf.bool)

Indicates whether the overall Brent search should continue.

Args:
        state: A Python `_BrentSearchState` namedtuple.
        stopping_policy_fn: Python `callable` controlling the algorithm termination.

Returns:
        A boolean value indicating whether the overall search should continue.


def __init__(self, name, encode_name, is_array, external):
    Message.__init__(self, name, encode_name, is_array)
    if stl.base.IsString(external):
        module, message_type = external.rsplit('.', 1)
        self.descriptor = importlib.import_module(module).__getattribute__(message_type).DESCRIPTOR
    else:
        self.descriptor = external
    self.external = MessageFromExternal._MakeClass(self.descriptor)
    self._ExtractFieldsFromDesc(self.descriptor)

A protocol specification external.

    message mProtocolExample {
        encode "json";
        external "external.class.Name";

Attributes:
        external: External message class name.
        descriptor: google.protobuf.descriptor.Descriptor to generate self.external.


def forward(self, hidden_states: torch.FloatTensor, attention_mask: Optional[torch.Tensor]=None, head_mask: Optional[torch.Tensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[Tuple, BaseModelOutput]:
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    if attention_mask is not None:
        attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)
    hidden_states = self.layer_norm(hidden_states)
    hidden_states = self.dropout(hidden_states)
    position_bias = self.embed_positions(hidden_states)
    synced_gpus = is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self)
    all_hidden_states = () if output_hidden_states else None
    all_self_attentions = () if output_attentions else None
    if head_mask is not None:
        if head_mask.size()[0] != len(self.layers):
            raise ValueError(f'The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}.')
    for idx, encoder_layer in enumerate(self.layers):
        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)
        skip_the_layer = False
        if self.training:
            dropout_probability = torch.rand([])
            skip_the_layer = dropout_probability < self.layerdrop
        if not skip_the_layer or synced_gpus:
            if self.gradient_checkpointing and self.training:
                layer_outputs = self._gradient_checkpointing_func(encoder_layer.__call__, hidden_states, attention_mask, head_mask[idx] if head_mask is not None else None, position_bias, output_attentions)
            else:
                layer_outputs = encoder_layer(hidden_states, attention_mask=attention_mask, position_bias=position_bias, layer_head_mask=head_mask[idx] if head_mask is not None else None, output_attentions=output_attentions)
            hidden_states = layer_outputs[0]
        if skip_the_layer:
            layer_outputs = (None, None)
        if output_attentions:
            all_self_attentions = all_self_attentions + (layer_outputs[1],)
    if output_hidden_states:
        all_hidden_states = all_hidden_states + (hidden_states,)
    if not return_dict:
        return tuple((v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None))
    return BaseModelOutput(last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_self_attentions)

Args:
    hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, feature_size)`):
        Features extracted from the speech or text input by the encoder prenet.
    attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
        Mask to avoid performing convolution and attention on padding token indices. Mask values selected in
        `[0, 1]`:

        - 1 for tokens that are **not masked**,
        - 0 for tokens that are **masked**.

        [What are attention masks?](../glossary#attention-mask)
    output_attentions (`bool`, *optional*):
        Whether or not to return the attentions tensors of all attention layers. See `attentions` under
        returned tensors for more detail.
    head_mask (`torch.Tensor` of shape `(encoder_layers, encoder_attention_heads)`, *optional*):
        Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:

        - 1 indicates the head is **not masked**,
        - 0 indicates the head is **masked**.

    output_hidden_states (`bool`, *optional*):
        Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
        for more detail.
    return_dict (`bool`, *optional*):
        Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.


def get_vocabulary(self, include_special_tokens=True):
    if self.lookup_table.size() == 0:
        vocab, indices = ([], [])
    else:
        keys, values = self.lookup_table.export()
        vocab, indices = (values, keys) if self.invert else (keys, values)
        vocab, indices = (self._tensor_vocab_to_numpy(vocab), indices.numpy())
    lookup = collections.defaultdict(lambda: self.oov_token, zip(indices, vocab))
    vocab = [lookup[x] for x in range(self.vocabulary_size())]
    if self.mask_token is not None and self.output_mode == 'int':
        vocab[0] = self.mask_token
    if not include_special_tokens:
        vocab = vocab[self._token_start_index():]
    if self.vocabulary_dtype == 'string':
        return [i.decode('utf-8') if isinstance(i, bytes) else i for i in vocab]
    else:
        return vocab

Returns the current vocabulary of the layer.

Args:
        include_special_tokens: If `True`, the returned vocabulary
            will include mask and OOV tokens,
            and a term's index in the vocabulary
            will equal the term's index when calling the layer.
            If `False`, the returned vocabulary will not include
            any mask or OOV tokens.


def flip_channel_order(image: np.ndarray, data_format: Optional[ChannelDimension]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None) -> np.ndarray:
    input_data_format = infer_channel_dimension_format(image) if input_data_format is None else input_data_format
    if input_data_format == ChannelDimension.LAST:
        image = image[..., ::-1]
    elif input_data_format == ChannelDimension.FIRST:
        image = image[::-1, ...]
    else:
        raise ValueError(f'Unsupported channel dimension: {input_data_format}')
    if data_format is not None:
        image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
    return image

Flips the channel order of the image.

    If the image is in RGB format, it will be converted to BGR and vice versa.

Args:
        image (`np.ndarray`):
            The image to flip.
        data_format (`ChannelDimension`, *optional*):
            The channel dimension format for the output image. Can be one of:
                - `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
                - `ChannelDimension.LAST`: image in (height, width, num_channels) format.
            If unset, will use same as the input image.
        input_data_format (`ChannelDimension`, *optional*):
            The channel dimension format for the input image. Can be one of:
                - `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
                - `ChannelDimension.LAST`: image in (height, width, num_channels) format.
            If unset, will use the inferred format of the input image.


def GetTermIdentifier(self):
    return self._term

Returns the TERM environment variable for the console.

Returns:
        str: A str that describes the console's text capabilities


def hard_silu(x):
    if any_symbolic_tensors((x,)):
        return HardSilu().symbolic_call(x)
    return backend.nn.hard_silu(x)

Hard SiLU activation function, also known as Hard Swish.

    It is defined as:

    - `0` if `if x < -3`
    - `x` if `x > 3`
    - `x * (x + 3) / 6` if `-3 <= x <= 3`

    It's a faster, piecewise linear approximation of the silu activation.

Args:
        x: Input tensor.

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

Example:
    >>> x = keras.ops.convert_to_tensor([-3.0, -1.0, 0.0, 1.0, 3.0])
    >>> keras.ops.hard_silu(x)
    array([-0.0, -0.3333333, 0.0, 0.6666667, 3.0], shape=(5,), dtype=float32)


def _Inputs(op: ops.Operation, xs_set):
    if _IsFunction(op.graph):
        inputs = []
        for t in op.inputs:
            if t not in xs_set:
                t = _MaybeCaptured(t)
            inputs.append(t)
        return inputs
    else:
        return op.inputs

Returns the inputs of op, crossing closure boundaries where necessary.

Args:
        op: Operation
        xs_set: ObjectIdentitySet of Tensors we are differentiating w.r.t.

Returns:
        A list of tensors. The tensors may be from multiple Graph/FuncGraphs if op
        is in a FuncGraph and has captured inputs.


def __init__(self, agg_func: Callable[[Iterable[int]], int], agg_model_id: Optional[str]=None, include_source_predictions: bool=False, normal_label: int=DEFAULT_NORMAL_LABEL, outlier_label: int=DEFAULT_OUTLIER_LABEL, missing_label: int=DEFAULT_MISSING_LABEL):
    self._agg = agg_func
    self._normal_label = normal_label
    self._outlier_label = outlier_label
    self._missing_label = missing_label
    _AggModelIdMixin.__init__(self, agg_model_id)
    _SourcePredictionMixin.__init__(self, include_source_predictions)

Aggregates anomaly predictions based on their labels.

    This is an abstract base class for `AggregationFn`s that combine multiple
    `AnomalyPrediction` objects into a single `AnomalyPrediction` based on
    the labels of the input predictions.

Args:
        agg_func (Callable[[Iterable[int]], int]): A function that aggregates
        a collection of anomaly labels (integers) into a single label.
        agg_model_id (Optional[str]): The model id used in aggregated predictions.
        Defaults to None.
        include_source_predictions (bool): If True, include the input predictions in
        the `source_predictions` of the output. Defaults to False.


def _format_time(self, time_per_unit, unit_name):
    formatted = ''
    if time_per_unit >= 1 or time_per_unit == 0:
        formatted += f' {time_per_unit:.0f}s/{unit_name}'
    elif time_per_unit >= 0.001:
        formatted += f' {time_per_unit * 1000.0:.0f}ms/{unit_name}'
    else:
        formatted += f' {time_per_unit * 1000000.0:.0f}us/{unit_name}'
    return formatted

format a given duration to display to the user.

    Given the duration, this function formats it in either milliseconds
    or seconds and displays the unit (i.e. ms/step or s/epoch).

Args:
        time_per_unit: the duration to display
        unit_name: the name of the unit to display

Returns:
        A string with the correctly formatted duration and units


def unshuffle_from_sc_to_cpu(t: tensor.Tensor, num_sparse_cores: int, offset_in_shard: int, size_in_shard: int, shard_rotation: int=0) -> tensor.Tensor:
    old_shape = t.shape
    if t.shape[0] % num_sparse_cores != 0:
        raise ValueError('The dim of table ({}) should be multiple of number of sparse cores ({})'.format(t.shape[1], num_sparse_cores))
    shards_t = array_ops.reshape(t, (num_sparse_cores, t.shape[0] // num_sparse_cores, t.shape[1]))
    shards = shards_t[:, offset_in_shard:offset_in_shard + size_in_shard, :]
    if shard_rotation:
        shards = manip_ops.roll(shards, -shard_rotation, axis=0)
    intermediate_tensor = array_ops.transpose(shards, (1, 0, 2))
    new_shape = (size_in_shard * num_sparse_cores, old_shape[1])
    return array_ops.reshape(intermediate_tensor, new_shape)

Unshuffles the sparse core sharded embedding tables to unsharded.

    This converts an input tensor respresenting stacked and sharded embedding
    table into a specific embedding table variable by using the provided
    metadata about the said table within the stacked, sharded embedding table.

Args:
        t: The input stacked and sharded embedding table from sparsecore.
        num_sparse_cores: The number of sparsecores, this determines the number of
        shards that are present in the input t.
        offset_in_shard: Offset within a shard where the queried table starts.
        size_in_shard: size (number of rows) of this queried table within each shard
        of the input t.
        shard_rotation: The rotation of this table's shards.

Returns:
        An embedding table which is part of the stacked embedding table t.


def isanytargetmethod(object):
    decorators, target = tf_decorator.unwrap(object)
    for decorator in decorators:
        if _inspect.ismethod(decorator.decorated_target):
            return True
    while isinstance(target, functools.partial):
        target = target.func
    return callable(target) and (not _inspect.isfunction(target))

Checks if `object` or a TF Decorator wrapped target contains self or cls.

    This function could be used along with `tf_inspect.getfullargspec` to
    determine if the first argument of `object` argspec is self or cls. If the
    first argument is self or cls, it needs to be excluded from argspec when we
    compare the argspec to the input arguments and, if provided, the tf.function
    input_signature.

    Like `tf_inspect.getfullargspec` and python `inspect.getfullargspec`, it
    does not unwrap python decorators.

Args:
        obj: An method, function, or functool.partial, possibly decorated by
        TFDecorator.

Returns:
        A bool indicates if `object` or any target along the chain of TF decorators
        is a method.


def attr_chain(obj, attr):
    next = getattr(obj, attr)
    while next:
        yield next
        next = getattr(next, attr)

Follow an attribute chain.

    If you have a chain of objects where a.foo -> b, b.foo-> c, etc, use this to
    iterate over all objects in the chain. Iteration is terminated by getattr(x,
    attr) is None.

Args:
        obj: the starting object
        attr: the name of the chaining attribute

Yields:
        Each successive object in the chain.


def forward(self, pixel_values: Optional[torch.Tensor]=None, bool_masked_pos: Optional[torch.BoolTensor]=None, head_mask: Optional[torch.Tensor]=None, labels: Optional[torch.Tensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, interpolate_pos_encoding: bool=False, return_dict: Optional[bool]=None) -> Union[tuple, MaskedLMOutput]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    outputs = self.beit(pixel_values, bool_masked_pos=bool_masked_pos, head_mask=head_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, interpolate_pos_encoding=interpolate_pos_encoding, return_dict=return_dict)
    sequence_output = outputs[0]
    sequence_output = self.layernorm(sequence_output)
    prediction_scores = self.lm_head(sequence_output[:, 1:])
    masked_lm_loss = None
    if labels is not None:
        loss_fct = CrossEntropyLoss()
        masked_lm_loss = loss_fct(prediction_scores[bool_masked_pos], labels)
    if not return_dict:
        output = (prediction_scores,) + outputs[1:]
        return (masked_lm_loss,) + output if masked_lm_loss is not None else output
    return MaskedLMOutput(loss=masked_lm_loss, logits=prediction_scores, hidden_states=outputs.hidden_states, attentions=outputs.attentions)

bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`):
    Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
    labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
    Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
    config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
    `config.num_labels > 1` a classification loss is computed (Cross-Entropy).

Example:
    ```python
    >>> from transformers import AutoImageProcessor, BeitForMaskedImageModeling
    >>> import torch
    >>> from PIL import Image
    >>> import requests

    >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
    >>> image = Image.open(requests.get(url, stream=True).raw)

    >>> image_processor = AutoImageProcessor.from_pretrained("microsoft/beit-base-patch16-224-pt22k")
    >>> model = BeitForMaskedImageModeling.from_pretrained("microsoft/beit-base-patch16-224-pt22k")

    >>> num_patches = (model.config.image_size // model.config.patch_size) ** 2
    >>> pixel_values = image_processor(images=image, return_tensors="pt").pixel_values
    >>> # create random boolean mask of shape (batch_size, num_patches)
    >>> bool_masked_pos = torch.randint(low=0, high=2, size=(1, num_patches)).bool()

    >>> outputs = model(pixel_values, bool_masked_pos=bool_masked_pos)
    >>> loss, logits = outputs.loss, outputs.logits
    >>> list(logits.shape)
    [1, 196, 8192]
    ```


def __init__(self, bos_token_id: int):
    if bos_token_id < 0:
        raise ValueError(f'The forced bos token id  must be a non-negative integer, got {bos_token_id}')
    self.bos_token_id = bos_token_id

[`TFLogitsProcessor`] that enforces the specified token as the first generated token.

Args:
        bos_token_id (`int`):
            The id of the token to force as the first generated token.