def _build_select_and_next_from_expressions(self, builders: Tuple[column_expression_builder.ColumnExpressionBuilder, ...], child_builders: MutableSequence[column_expression_builder.ColumnExpressionBuilder], columns_selected: MutableSequence[str]) -> Tuple[MutableSequence[str], MutableSequence[str]]:
    select_expressions, next_from_expressions = ([], [])
    for column_name in columns_selected:
        select_expressions.append(f'(SELECT {column_name}) AS {column_name}')
    for builder in builders:
        child_builders.extend(builder.children)
        column_alias = _get_column_alias(builder)
        if builder.column_name:
            columns_selected.append(column_alias)
        needs_unnest = builder.needs_unnest or builder.children
        select_expression = self._encode(builder=builder, select_scalars_as_array=needs_unnest)
        if needs_unnest:
            select_expression = f'{select_expression} AS {column_alias}_needs_unnest_'
            next_from_expressions.append(self._build_next_from_expression(column_alias))
        else:
            select_expression = f'{select_expression} AS {column_alias}'
        select_expressions.append(select_expression)
    return (select_expressions, next_from_expressions)

Build select expressions and next from expressions from the builders.

Args:
        builders: the immutable current builders to compute select expressions.
        child_builders: collects the current given builders' children for the next
        round.
        columns_selected: accumulatively collects columns which has already been
        handled completely.

Returns:
        The select expressions and next from expressions computed form the given
        builders.


def info(msg: str, *args, **kwargs) -> None:
    _DEFAULT_LOGGER.info(msg, *args, **kwargs)

Logs info message.

Args:
        msg: Message with possible format string.
        *args: Values for variables in the format string.
        **kwargs: Keyword arguments for the logger.


def _decorator(func):
    opname = func.__name__
    cap_sym_name = sym_name.capitalize()
    func.__doc__ = '\n    Assert the condition `x {sym}` holds element-wise.\n\n    When running in graph mode, you should add a dependency on this operation\n    to ensure that it runs. Example of adding a dependency to an operation:\n\n    ```python\n    with tf.control_dependencies([tf.debugging.{opname}(x, y)]):\n      output = tf.reduce_sum(x)\n    ```\n\n    {sym_name} means, for every element `x[i]` of `x`, we have `x[i] {sym}`.\n    If `x` is empty this is trivially satisfied.\n\n    Args:\n      x:  Numeric `Tensor`.\n      data:  The tensors to print out if the condition is False.  Defaults to\n        error message and first few entries of `x`.\n      summarize: Print this many entries of each tensor.\n      message: A string to prefix to the default message.\n      name: A name for this operation (optional).  Defaults to "{opname}".\n\n    Returns:\n      Op that raises `InvalidArgumentError` if `x {sym}` is False.\n      @compatibility(eager)\n        returns None\n      @end_compatibility\n\n    Raises:\n      InvalidArgumentError: if the check can be performed immediately and\n        `x {sym}` is False. The check can be performed immediately during\n        eager execution or if `x` is statically known.\n    '.format(sym=sym, sym_name=cap_sym_name, opname=opname)
    return func

Generated decorator that adds the appropriate docstring to the function for symbol `sym`.

Args:
        func: Function for a TensorFlow op

Returns:
        Version of `func` with documentation attached.


def __init__(self, key_dtype, value_dtype):
    self._key_dtype = dtypes.as_dtype(key_dtype)
    self._value_dtype = dtypes.as_dtype(value_dtype)
    super(LookupInterface, self).__init__()

Construct a lookup table interface.

Args:
        key_dtype: The table key type.
        value_dtype: The table value type.


def call(self, input_values: tf.Tensor, attention_mask: tf.Tensor | None=None, token_type_ids: tf.Tensor | None=None, position_ids: tf.Tensor | None=None, head_mask: tf.Tensor | None=None, inputs_embeds: tf.Tensor | None=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, training: bool=False) -> Union[TFBaseModelOutput, Tuple[tf.Tensor]]:
    output_hidden_states = output_hidden_states if output_hidden_states else self.config.output_hidden_states
    output_attentions = output_attentions if output_attentions else self.config.output_attentions
    return_dict = return_dict if return_dict else self.config.return_dict
    outputs = self.hubert(input_values=input_values, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, training=training)
    return outputs

Returns:

Example:
    ```python
    >>> from transformers import AutoProcessor, TFHubertModel
    >>> from datasets import load_dataset
    >>> import soundfile as sf

    >>> processor = AutoProcessor.from_pretrained("facebook/hubert-large-ls960-ft")
    >>> model = TFHubertModel.from_pretrained("facebook/hubert-large-ls960-ft")

    >>> def map_to_array(batch):
    ...     speech, _ = sf.read(batch["file"])
    ...     batch["speech"] = speech
    ...     return batch

    >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
    >>> ds = ds.map(map_to_array)

    >>> input_values = processor(ds["speech"][0], return_tensors="tf").input_values  # Batch size 1
    >>> hidden_states = model(input_values).last_hidden_state
    ```


def __init__(self, given, enum_type, options):
    super(InvalidEnumValue, self).__init__('Could not parse [{0}] into a valid {1}.  Valid values are [{2}]'.format(given, enum_type, ', '.join(options)))

Constructs a new exception.

Args:
        given: str, The given string that could not be parsed.
        enum_type: str, The human readable name of the enum you were trying to
        parse.
        options: list(str), The valid values for this enum.


def broadcast_to(rt_input, shape: DynamicRaggedShape):
    if not isinstance(shape, DynamicRaggedShape):
        raise TypeError('shape must be a DynamicRaggedShape')
    rt_input = ragged_tensor.convert_to_tensor_or_ragged_tensor(rt_input)
    origin_shape = None
    if ragged_tensor.is_ragged(rt_input):
        if shape.num_row_partitions != 0:
            if rt_input.row_splits.dtype != shape.dtype:
                raise ValueError('Cannot coerce row_splits.dtype')
        else:
            shape = shape.with_dtype(rt_input.row_splits.dtype)
        origin_shape = DynamicRaggedShape.from_tensor(rt_input)
    elif shape.num_row_partitions != 0:
        origin_shape = DynamicRaggedShape.from_tensor(rt_input, dtype=shape.dtype)
    else:
        origin_shape = DynamicRaggedShape.from_tensor(rt_input, dtype=dtypes.int64)
        shape = shape.with_dtype(dtype=dtypes.int64)
    broadcaster = _get_broadcaster(origin_shape, shape)
    return broadcaster.broadcast(rt_input)

Broadcasts a potentially ragged tensor to a ragged shape.

    Tiles `rt_input` as necessary to match the given shape.

    Behavior is undefined if `rt_input` is not broadcast-compatible with `shape`.

Args:
        rt_input: The potentially ragged tensor to broadcast.
        shape: A `DynamicRaggedShape`

Returns:
        A potentially ragged tensor whose values are taken from
        `rt_input`, and whose shape matches `shape`.


def _wait_for_glob(self, pattern, timeout_secs, for_checkpoint=True):
    end_time = time.time() + timeout_secs
    while time.time() < end_time:
        if for_checkpoint:
            if checkpoint_management.checkpoint_exists(pattern):
                return
        elif len(gfile.Glob(pattern)) >= 1:
            return
        time.sleep(0.05)
    self.assertFalse(True, 'Glob never matched any file: %s' % pattern)

Wait for a checkpoint file to appear.

Args:
        pattern: A string.
        timeout_secs: How long to wait for in seconds.
        for_checkpoint: whether we're globbing for checkpoints.


def get_logits_and_probs(logits=None, probs=None, multidimensional=False, validate_args=False, name='get_logits_and_probs', dtype=None):
    with ops.name_scope(name, values=[probs, logits]):
        if (probs is None) == (logits is None):
            raise ValueError('Must pass probs or logits, but not both.')
        if probs is None:
            logits = ops.convert_to_tensor(logits, name='logits', dtype=dtype)
            if not logits.dtype.is_floating:
                raise TypeError('logits must having floating type.')
            if multidimensional:
                if validate_args:
                    logits = embed_check_categorical_event_shape(logits)
                return (logits, nn.softmax(logits, name='probs'))
            return (logits, math_ops.sigmoid(logits, name='probs'))
        probs = ops.convert_to_tensor(probs, name='probs', dtype=dtype)
        if not probs.dtype.is_floating:
            raise TypeError('probs must having floating type.')
        if validate_args:
            with ops.name_scope('validate_probs'):
                one = constant_op.constant(1.0, probs.dtype)
                dependencies = [check_ops.assert_non_negative(probs)]
                if multidimensional:
                    probs = embed_check_categorical_event_shape(probs)
                    dependencies += [check_ops.assert_near(math_ops.reduce_sum(probs, -1), one, message='probs does not sum to 1.')]
                else:
                    dependencies += [check_ops.assert_less_equal(probs, one, message='probs has components greater than 1.')]
                probs = control_flow_ops.with_dependencies(dependencies, probs)
        with ops.name_scope('logits'):
            if multidimensional:
                return (math_ops.log(probs), probs)
            return (math_ops.log(probs) - math_ops.log1p(-1.0 * probs), probs)

Converts logit to probabilities (or vice-versa), and returns both.

Args:
        logits: Floating-point `Tensor` representing log-odds.
        probs: Floating-point `Tensor` representing probabilities.
        multidimensional: Python `bool`, default `False`. If `True`, represents
        whether the last dimension of `logits` or `probs`, a `[N1, N2, ...  k]`
        dimensional tensor, representing the logit or probability of `shape[-1]`
        classes.
        validate_args: Python `bool`, default `False`. When `True`, either assert `0
        <= probs <= 1` (if not `multidimensional`) or that the last dimension of
        `probs` sums to one.
        name: A name for this operation (optional).
        dtype: `tf.DType` to prefer when converting args to `Tensor`s.

Returns:
        logits, probs: Tuple of `Tensor`s. If `probs` has an entry that is `0` or
        `1`, then the corresponding entry in the returned logit will be `-Inf` and
        `Inf` respectively.

Raises:
        ValueError: if neither `probs` nor `logits` were passed in, or both were.


def make_seeds(self, count=1):
    alg = self.algorithm
    if alg in (a.value for a in random_ops_util.Algorithm):
        keys = self._make_int64_keys(shape=[count])
        zeros = array_ops.zeros_like(keys)
        return array_ops_stack.stack([keys, zeros])
    else:
        raise ValueError(stateless_random_ops.unsupported_alg_error_msg(alg))

Generates seeds for stateless random ops.

    For example:

    ```python
    seeds = get_global_generator().make_seeds(count=10)
    for i in range(10):
        seed = seeds[:, i]
        numbers = stateless_random_normal(shape=[2, 3], seed=seed)
        ...
    ```

Args:
        count: the number of seed pairs (note that stateless random ops need a
        pair of seeds to invoke).

Returns:
        A tensor of shape [2, count] and dtype int64.


def output_mask(self):
    output = self.output
    if isinstance(output, list):
        return [getattr(x, '_keras_mask', None) for x in output]
    else:
        return getattr(output, '_keras_mask', None)

Retrieves the output mask tensor(s) of a layer.

    Only applicable if the layer has exactly one inbound node,
    i.e. if it is connected to one incoming layer.

Returns:
        Output mask tensor (potentially None) or list of output
        mask tensors.

Raises:
        AttributeError: if the layer is connected to
        more than one incoming layers.


def get_master_port(real_launcher=False):
    master_port_base = os.environ.get('DS_TEST_PORT', DEFAULT_MASTER_PORT)
    if not real_launcher:
        master_port_base = str(int(master_port_base) + 1)
    return master_port_base

When using a single gpu launcher emulation (i.e. not deepspeed or python -m torch.distributed)
    the issue is that once the port is tied it can't be used anywhere else outside of this process,
    since torch.dist doesn't free the port until the process exits. Therefore for the sake of being
    able to run both emulated launcher and normal launcher tests we need 2 distinct ports.

    This function will give the right port in the right context. For real launcher it'll give the
    base port, for emulated launcher it'll give the base port + 1. In both cases a string is
    returned.

Args:
        `real_launcher`: whether a real launcher is going to be used, or the emulated one


def __init__(self, seed=None, name=None, **kwargs):
    if name is None:
        name = auto_name(self.__class__.__name__)
    self.name = name
    custom_backend = kwargs.pop('backend', None)
    if kwargs:
        raise ValueError(f'Unrecognized keyword arguments: {kwargs}')
    if custom_backend is not None:
        self.backend = custom_backend
    else:
        self.backend = backend
    self._initial_seed = seed
    if seed is None:
        seed = make_default_seed()
    if not isinstance(seed, int):
        raise ValueError(f'Argument `seed` must be an integer. Received: seed={seed}')

    def seed_initializer(*args, **kwargs):
        dtype = kwargs.get('dtype', None)
        return self.backend.convert_to_tensor([seed, 0], dtype=dtype)
    with self.backend.name_scope(self.name, caller=self):
        self.state = self.backend.Variable(seed_initializer, shape=(2,), dtype=self.backend.random_seed_dtype(), trainable=False, aggregation='none', name='seed_generator_state')

Generates variable seeds upon each call to a function generating
    random numbers.

    In Keras, all random number generators (such as
    `keras.random.normal()`) are stateless, meaning that if you pass an
    integer seed to them (such as `seed=42`), they will return the same
    values for repeated calls. To get different values for each
    call, a `SeedGenerator` providing the state of the random generator
    has to be used.

    Note that all the random number generators have a default seed of None,
    which implies that an internal global SeedGenerator is used.
    If you need to decouple the RNG from the global state you can provide
    a local `StateGenerator` with either a deterministic or random initial
    state.

    Remark concerning the JAX backen: Note that the use of a local
    `StateGenerator` as seed argument is required for JIT compilation of
    RNG with the JAX backend, because the use of global state is not
    supported.

Example:
    ```python
    seed_gen = keras.random.SeedGenerator(seed=42)
    values = keras.random.normal(shape=(2, 3), seed=seed_gen)
    new_values = keras.random.normal(shape=(2, 3), seed=seed_gen)
    ```

    Usage in a layer:

    ```python
    class Dropout(keras.Layer):
        def __init__(self, **kwargs):
            super().__init__(**kwargs)
            self.seed_generator = keras.random.SeedGenerator(1337)

        def call(self, x, training=False):
            if training:
                return keras.random.dropout(
                    x, rate=0.5, seed=self.seed_generator
                )
            return x
    ```


def nondominated_sort(inputs: List[pg.DNA]) -> List[List[pg.DNA]]:
    dependency_graph = [[] for i in range(len(inputs))]
    indegree = [0] * len(inputs)
    queue = []
    for i in range(len(inputs)):
        for j in range(len(inputs)):
            if dominates(base.get_fitness(inputs[i]), base.get_fitness(inputs[j])):
                dependency_graph[i].append(j)
            elif dominates(base.get_fitness(inputs[j]), base.get_fitness(inputs[i])):
                indegree[i] += 1
        if indegree[i] == 0:
            queue.append(i)
    result = []
    while queue:
        l = len(queue)
        frontier = []
        for i in range(l):
            parent = queue.pop(0)
            frontier.append(inputs[parent])
            for child in dependency_graph[parent]:
                indegree[child] -= 1
                if indegree[child] == 0:
                    queue.append(child)
        result.append(frontier)
    return result

Algorithm fast-non-dominated-sort implementation using topological sort.

    Please see section III A in the original paper.

Args:
        inputs: A list of DNA that need to be sorted.

Returns:
        A list of sorted frontier (a list of DNA).


def resize(self, image: np.ndarray, size: Dict[str, int], patch_size: Dict[str, int], resample: PILImageResampling=PILImageResampling.BICUBIC, data_format: Optional[Union[str, ChannelDimension]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None, **kwargs) -> np.ndarray:
    if 'longest_edge' in size:
        size = (size['longest_edge'], size['longest_edge'])
    elif 'height' in size and 'width' in size:
        size = (size['height'], size['width'])
    else:
        raise ValueError("size must contain either 'longest_edge' or 'height' and 'width'.")
    if 'height' in patch_size and 'width' in patch_size:
        patch_size = (patch_size['height'], patch_size['width'])
    else:
        raise ValueError("patch_size must contain either 'shortest_edge' or 'height' and 'width'.")
    output_size = get_resize_output_image_size(image, size=size, patch_size=patch_size, input_data_format=input_data_format)
    return resize(image, size=output_size, resample=resample, data_format=data_format, input_data_format=input_data_format, **kwargs)

Resize an image. The shortest edge of the image is resized to size["shortest_edge"], with the longest edge
    resized to keep the input aspect ratio.

Args:
        image (`np.ndarray`):
            Image to resize.
        size (`Dict[str, int]`):
            Dict containing the longest possible edge of the image.
        patch_size (`Dict[str, int]`):
            Patch size used to calculate the size of the output image.
        resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BICUBIC`):
            Resampling filter to use when resiizing the image.
        data_format (`str` or `ChannelDimension`, *optional*):
            The channel dimension format of the image. If not provided, it will be the same as the input image.
        input_data_format (`ChannelDimension` or `str`, *optional*):
            The channel dimension format of the input image. If not provided, it will be inferred.


def nrows(self):
    if self.rank == 0:
        return None
    return self._ragged_shape[0]

The number of rows in this StructuredTensor (if rank>0).

    This means the length of the outer-most dimension of the StructuredTensor.

    Notice that if `self.rank > 1`, then this equals the number of rows
    of the first row partition. That is,
    `self.nrows() == self.row_partitions[0].nrows()`.

    Otherwise `self.nrows()` will be the first dimension of the field values.

Returns:
        A scalar integer `Tensor` (or `None` if `self.rank == 0`).


def preprocess(train_data_path: str, feature_thres: int, val_data_path: typing.Optional[str]=None) -> typing.Tuple[Dataset, typing.List[str], typing.Optional[Dataset]]:
    features = extract_features(train_data_path, feature_thres)
    feature_index = dict(((feature, i) for i, feature in enumerate(features)))
    train_dataset = load_dataset(train_data_path, feature_index)
    val_dataset = load_dataset(val_data_path, feature_index) if val_data_path else None
    return (train_dataset, features, val_dataset)

Loads entries and translates them into JAX arrays. The boolean matrix of
    the input data is represented by row indices and column indices of True values
    instead of the matrix itself for memory efficiency, assuming the matrix is
    highly sparse. Row and column indices are not guaranteed to be sorted.

Args:
        train_data_path (str): A file path to the training data file.
        feature_thres (str): A threshold to filter out features whose number of
        occurances does not exceed the value.
        val_data_path (str, optional): A file path to the validation data file.

Returns:
        A tuple of following items:
        - train_dataset (Dataset): The training dataset.
        - features (List[str]): The list of features.
        - val_dataset (Optional[Dataset]): The validation dataset.
        This becomes None if val_data_path is None.


def hard_tanh(x):
    return ops.hard_tanh(x)

HardTanh activation function.

    It is defined as:
    `hard_tanh(x) = -1 for x < -1`,
    `hard_tanh(x) = x for -1 <= x <= 1`,
    `hard_tanh(x) = 1 for x > 1`.

Args:
        x: Input tensor.


def __init__(self, lr=1.0, rho=0.95, epsilon=None, decay=0.0, **kwargs):
    super(Adadelta, self).__init__(**kwargs)
    with backend.name_scope(self.__class__.__name__):
        self.lr = backend.variable(lr, name='lr')
        self.decay = backend.variable(decay, name='decay')
        self.iterations = backend.variable(0, dtype='int64', name='iterations')
    if epsilon is None:
        epsilon = backend.epsilon()
    self.rho = rho
    self.epsilon = epsilon
    self.initial_decay = decay

Adadelta optimizer.

    Adadelta is a more robust extension of Adagrad
    that adapts learning rates based on a moving window of gradient updates,
    instead of accumulating all past gradients. This way, Adadelta continues
    learning even when many updates have been done. Compared to Adagrad, in the
    original version of Adadelta you don't have to set an initial learning
    rate. In this version, initial learning rate and decay factor can
    be set, as in most other Keras optimizers.

    It is recommended to leave the parameters of this optimizer
    at their default values.

Args:
        lr: float >= 0. Initial learning rate, defaults to 1.
        It is recommended to leave it at the default value.
        rho: float >= 0. Adadelta decay factor, corresponding to fraction of
        gradient to keep at each time step.
        epsilon: float >= 0. Fuzz factor.
        If `None`, defaults to `backend.epsilon()`.
        decay: float >= 0. Initial learning rate decay.

    References:
        - [Adadelta - an adaptive learning rate
        method](http://arxiv.org/abs/1212.5701)


def _process_new(self, feed_item):
    return {'assetIdentifier': {'name': feed_item.get(FieldMap.CREATIVE_ASSET_FILE_NAME, None), 'type': feed_item.get(FieldMap.CREATIVE_TYPE, None)}}

Creates a new creative asset DCM object from a feed item representing a creative asset from the Bulkdozer feed.

    This function simply creates the object to be inserted later by the BaseDAO
    object.

Args:
        feed_item: Feed item representing the creative asset from the Bulkdozer
        feed.

Returns:
        A creative asset object ready to be inserted in DCM through the API.


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
    model_loader = TFModelHandlerTensor(model_uri=known_args.model_path).with_preprocess_fn(lambda image_name: read_image(image_name, known_args.image_dir))
    pipeline = test_pipeline
    if not test_pipeline:
        pipeline = beam.Pipeline(options=pipeline_options)
    image = pipeline | 'ReadImageNames' >> beam.io.ReadFromText(known_args.input) | 'FilterEmptyLines' >> beam.ParDo(filter_empty_lines)
    predictions = image | 'RunInference' >> RunInference(model_loader) | 'PostProcessOutputs' >> beam.ParDo(PostProcessor())
    _ = predictions | '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 forward(self, input_ids: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None, token_type_ids: Optional[torch.Tensor]=None, position_ids: Optional[torch.Tensor]=None, head_mask: Optional[torch.Tensor]=None, inputs_embeds: Optional[torch.Tensor]=None, labels: Optional[torch.Tensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[Tuple[torch.Tensor], ElectraForPreTrainingOutput]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    discriminator_hidden_states = self.electra(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)
    discriminator_sequence_output = discriminator_hidden_states[0]
    logits = self.discriminator_predictions(discriminator_sequence_output)
    loss = None
    if labels is not None:
        loss_fct = nn.BCEWithLogitsLoss()
        if attention_mask is not None:
            active_loss = attention_mask.view(-1, discriminator_sequence_output.shape[1]) == 1
            active_logits = logits.view(-1, discriminator_sequence_output.shape[1])[active_loss]
            active_labels = labels[active_loss]
            loss = loss_fct(active_logits, active_labels.float())
        else:
            loss = loss_fct(logits.view(-1, discriminator_sequence_output.shape[1]), labels.float())
    if not return_dict:
        output = (logits,) + discriminator_hidden_states[1:]
        return (loss,) + output if loss is not None else output
    return ElectraForPreTrainingOutput(loss=loss, logits=logits, hidden_states=discriminator_hidden_states.hidden_states, attentions=discriminator_hidden_states.attentions)

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

    - 0 indicates the token is an original token,
    - 1 indicates the token was replaced.

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

    >>> discriminator = ElectraForPreTraining.from_pretrained("google/electra-base-discriminator")
    >>> tokenizer = AutoTokenizer.from_pretrained("google/electra-base-discriminator")

    >>> sentence = "The quick brown fox jumps over the lazy dog"
    >>> fake_sentence = "The quick brown fox fake over the lazy dog"

    >>> fake_tokens = tokenizer.tokenize(fake_sentence, add_special_tokens=True)
    >>> fake_inputs = tokenizer.encode(fake_sentence, return_tensors="pt")
    >>> discriminator_outputs = discriminator(fake_inputs)
    >>> predictions = torch.round((torch.sign(discriminator_outputs[0]) + 1) / 2)

    >>> fake_tokens
    ['[CLS]', 'the', 'quick', 'brown', 'fox', 'fake', 'over', 'the', 'lazy', 'dog', '[SEP]']

    >>> predictions.squeeze().tolist()
    [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0]
    ```


def affine_transform(self, image: np.array, center: Tuple[float], scale: Tuple[float], rotation: float, size: Dict[str, int], data_format: Optional[ChannelDimension]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None) -> np.array:
    data_format = input_data_format if data_format is None else data_format
    size = (size['width'], size['height'])
    transformation = get_warp_matrix(rotation, center * 2.0, np.array(size) - 1.0, scale * 200.0)
    image = image if input_data_format == ChannelDimension.LAST else to_channel_dimension_format(image, ChannelDimension.LAST, input_data_format)
    image = scipy_warp_affine(src=image, M=transformation, size=(size[1], size[0]))
    image = to_channel_dimension_format(image, data_format, ChannelDimension.LAST)
    return image

Apply an affine transformation to an image.

Args:
        image (`np.array`):
            Image to transform.
        center (`Tuple[float]`):
            Center of the bounding box (x, y).
        scale (`Tuple[float]`):
            Scale of the bounding box with respect to height/width.
        rotation (`float`):
            Rotation angle in degrees.
        size (`Dict[str, int]`):
            Size of the destination image.
        data_format (`ChannelDimension`, *optional*, defaults to `ChannelDimension.FIRST`):
            The channel dimension format of the output image.
        input_data_format (`str` or `ChannelDimension`, *optional*):
            The channel dimension format of the input image.


def copy_function_def_to_graph_def_recursively(func_name, graph_def, copied_functions, default_graph=None):
    if func_name and (not is_function(func_name, default_graph)):
        raise ValueError(f'Function {func_name} was not found. Please make sure the FunctionDef `fdef` is correct.')
    if func_name in copied_functions:
        return
    copied_functions.add(func_name)
    func_def = get_function_def(func_name, default_graph)
    graph_def.library.function.add().CopyFrom(func_def)
    for node_def in func_def.node_def:
        op_def = default_graph.op_def_for_type(node_def.op)
        for attr in op_def.attr:
            if attr.type == 'func':
                func_name = node_def.attr[attr.name].func.name
                copy_function_def_to_graph_def_recursively(func_name, graph_def, copied_functions, default_graph)
            elif attr.type == 'list(func)':
                for fn in node_def.attr[attr.name].list.func:
                    func_name = fn.name
                    copy_function_def_to_graph_def_recursively(func_name, graph_def, copied_functions, default_graph)

Recursively copies `FunctionDef`s to `GraphDef`.

    It copies the outermost `FunctionDef` and all nested `FunctionDef`s to
    `graph_def`. The `copied_function` enforces that every `FunctionDef` will be
    copied at most once. The `FunctionDef`s will be found from `default_graph` if
    this function was called in graph mode or from eager context if this function
    was called in eager mode.

Args:
        func_name: The signature name of FunctionDef to be copied to `graph_def`.
        graph_def: The GraphDef that will contain all `FunctionDef`s in its library.
        copied_functions: A set contains all copied function names.
        default_graph: The `tf.Graph` where all `FunctionDef`s will be found
        in graph mode. Not used in eager mode.


def add_object(self, object_path, weights):
    if not isinstance(weights, dict):
        raise ValueError(f"Argument `weights` should be a dict where keys are weight names (usually '0', '1', etc.) and values are NumPy arrays. Received: type(weights)={type(weights)}")
    if '/' in object_path:
        elements = object_path.split('/')
        partial_path = '/'.join(elements[:-1])
        weights_dict = self.weights_dict
        for e in elements[:-1]:
            if e not in weights_dict:
                raise ValueError(f"Path '{partial_path}' not found in model.")
            weights_dict = weights_dict[e]
        weights_dict[elements[-1]] = weights
    else:
        self.weights_dict[object_path] = weights

Add a new object to the file (e.g. a layer).

Args:
        object_path: String, full path of the
            object to add (e.g. `"layers/dense_2"`).
        weights: Dict mapping weight names to weight
            values (arrays),
            e.g. `{"0": kernel_value, "1": bias_value}`.


def grad(f, has_aux=False):

    def check_loss_shape(np_loss):
        if not isinstance(np_loss, tf_np.ndarray):
            raise ValueError('The result of the function to take gradient must be an ndarray.')
        if not np_loss.shape.is_compatible_with([]):
            raise ValueError('The result of the function to take gradient must be a scalar.')

    def _f(params, *args):
        """The gradient function to be returned."""
        with backprop.GradientTape() as g:
            g.watch(nest.flatten(params))
            outputs = f(params, *args)
            if has_aux:
                np_loss, aux = outputs
            else:
                np_loss = outputs
            check_loss_shape(np_loss)
            tf_grads = g.gradient(np_loss, params)
            if has_aux:
                res = (tf_grads, aux)
            else:
                res = tf_grads
            return _tf_to_np(res)
    return _f

Returns a function that computes gradient of f.

    Gradients can only be computed through numpy and tensorflow operations and not
    through python float operations and values.

Args:
        f: a function of type (params, *args) -> scalar. 'params' can be a nested
        structure (made of lists and tuples) of ndarrays and the gradient is
        evaluated against it. `scalar` is a scalar ndarray.
        has_aux: bool, indicates whether fun returns a pair where the first element
        is considered the output of the mathematical function to be differentiated
        and the second element is auxiliary data.

Returns:
        A gradient function of type (params, *args) -> gradients, where the result
        'gradients' has the same structure and shapes as 'params'.


def _build_map(outputs):
    finished_nodes = set()
    nodes_in_progress = set()
    nodes_in_decreasing_depth = []
    layer_indices = {}
    for output in nest.flatten(outputs):
        _build_map_helper(output, finished_nodes, nodes_in_progress, nodes_in_decreasing_depth, layer_indices)
    return (nodes_in_decreasing_depth, layer_indices)

This method topologically sorts nodes in order from inputs to outputs.

    It uses a depth-first search to topologically sort nodes that appear in the
    _keras_history connectivity metadata of `outputs`.

Args:
        outputs: the output tensors whose _keras_history metadata should be walked.
        This may be an arbitrary nested structure.

Returns:
        A tuple like (ordered_nodes, layer_to_first_traversal_index)
        ordered_nodes: list of nodes appearing in the keras history, topologically
        sorted from original inputs to the `outputs`.
        (If outputs have different sets of ancestors, the inputs to one output
        may appear after a different output).
        layer_to_first_traversal_index:
        A dict mapping layer to the traversal index in the DFS where it is
        seen. Note: if a layer is shared by several nodes, the dict will only
        store the index corresponding to the *first* time the layer seen.


def limit_string_length(string, max_len=50):
    if max_len is None or len(string) <= max_len:
        return string
    else:
        return '...' + string[len(string) - max_len:]

Limit the length of input string.

Args:
        string: Input string.
        max_len: (int or None) If int, the length limit. If None, no limit.

Returns:
        Possibly length-limited string.


def parameters(self, use_literal_values: bool=False) -> Dict[str, str]:
    value_type = 'choice_and_literal' if use_literal_values else 'choice'
    return self.to_dict(value_type=value_type)

Returns parameters for this DNA to emit based on its spec.

    Deprecated: use `to_dict` instead.

Args:
        use_literal_values: If True, use literal values from DNASpec for Choices,
        otherwise returns '{choice}/{num_candidates} ({literal})'. Otherwise
        returns '{choice}/{num_candidates}'.

Returns:
        Dict of parameter names to their values mapped from this DNA.

Raises:
        RuntimeError: If DNA is not associated with a DNASpec.


def _get_context_id(self, context):
    if context in self._context_to_id:
        return self._context_to_id[context]
    graph_is_new = False
    with self._context_lock:
        if context not in self._context_to_id:
            graph_is_new = True
            context_id = _get_id()
            self._context_to_id[context] = context_id
    if graph_is_new:
        self.get_writer().WriteDebuggedGraph(debug_event_pb2.DebuggedGraph(graph_id=context_id, graph_name=getattr(context, 'name', None), outer_context_id=self._get_outer_context_id(context)))
    return self._context_to_id[context]

Get a unique ID for an op-construction context (e.g., a graph).

    If the graph has been encountered before, reuse the same unique ID.
    When encountering a new context (graph), this methods writes a DebugEvent
    proto with the debugged_graph field to the proper DebugEvent file.

Args:
        context: A context to get the unique ID for. Must be hashable. E.g., a
        Graph object.

Returns:
        A unique ID for the context.


def on_element(self, element, window, context):
    pass

Called when a new element arrives in a window.

Args:
        element: the element being added
        window: the window to which the element is being added
        context: a context (e.g. a TriggerContext instance) for managing state
        and setting timers


def cached_session(self, graph=None, config=None, target=None):
    config = self._create_config(config)
    if target is None:
        target = self._default_target
    if getattr(self._thread_local, 'cached_session', None) is None:
        self._thread_local.cached_session = session.Session(graph=None, config=config, target=target)
    sess = self._thread_local.cached_session
    with sess.graph.as_default(), sess.as_default():
        yield sess

Create a test session with master target set to the testing cluster.

    Creates a test session that connects to the local testing cluster.
    The session is only created once per test and then reused.

Args:
        graph: Optional graph to use during the returned session.
        config: An optional config_pb2.ConfigProto to use to configure the
        session.
        target: the target of session to connect to.

Yields:
        A Session object that should be used as a context manager to surround
        the graph building and execution code in a test case. Note that the
        session will live until the end of the test.


def create_input_data_based_on_hw_requirement(num_chip, max_unique_ids_per_partition, per_sc_vocab_size, per_sc_sample_count, num_minibatches_per_physical_sparse_core):
    num_sc_per_chip = 4
    num_physical_replica = num_chip * num_sc_per_chip
    col_ids = []
    row_ids = []
    gains = []
    smallest_num_division = np.power(2, np.ceil(np.log2(num_minibatches_per_physical_sparse_core)))
    division_size = (per_sc_vocab_size + smallest_num_division - 1) // smallest_num_division
    assert division_size >= max_unique_ids_per_partition, f'The max_unique_ids_per_partition is set to {max_unique_ids_per_partition} and the number of minibatches per sparse core is set to {num_minibatches_per_physical_sparse_core}. But the vocab size per sparse core is {per_sc_vocab_size} which is not going to fit that many minibatches, consider setting the number of minibatches smaller.'
    per_sc_per_minibatch_id_nums_for_each_replica = np.random.randint(max_unique_ids_per_partition * (num_minibatches_per_physical_sparse_core - 1) + 1, max_unique_ids_per_partition * num_minibatches_per_physical_sparse_core + 1, size=num_physical_replica)
    per_chip_sample_count = per_sc_sample_count * num_sc_per_chip
    for chip_id in range(num_chip):
        for sc_id in range(num_sc_per_chip):
            np.random.shuffle(per_sc_per_minibatch_id_nums_for_each_replica)
            for physical_replica_id in range(num_physical_replica):
                physical_replica_id_nums = per_sc_per_minibatch_id_nums_for_each_replica[physical_replica_id]
                local_col_ids = np.array([])
                for i in range(num_minibatches_per_physical_sparse_core):
                    local_col_ids_minibatch_size = max_unique_ids_per_partition
                    if i == num_minibatches_per_physical_sparse_core - 1:
                        local_col_ids_minibatch_size = physical_replica_id_nums - i * max_unique_ids_per_partition
                    local_col_ids = np.append(local_col_ids, np.random.choice(np.arange(division_size), size=local_col_ids_minibatch_size, replace=False) + i * division_size)
                local_row_ids = np.random.randint(low=0, high=per_sc_sample_count, size=physical_replica_id_nums)
                row_ids += list(local_row_ids + chip_id * per_chip_sample_count + sc_id * per_sc_sample_count)
                col_ids += list(local_col_ids * num_physical_replica + physical_replica_id)
                gains += list(np.random.random(size=physical_replica_id_nums))
    return (np.array(row_ids), np.array(col_ids), np.array(gains))

Create the coo tensor based on hardware requirements.

Args:
        num_chip: number of chips in the tpu system.
        max_unique_ids_per_partition: max unique ids per physical replica
        per_sc_vocab_size: per sc shard of table size.
        per_sc_sample_count: per sc sample count.
        num_minibatches_per_physical_sparse_core: per sc minibatch number.

Returns:
        row_ids, col_ids, gains and splits


def solve_triangular(a, b, lower=False):
    if any_symbolic_tensors((a, b)):
        return SolveTriangular(lower).symbolic_call(a, b)
    return _solve_triangular(a, b, lower)

Solves a linear system of equations given by `a x = b`.

Args:
        a: A tensor of shape `(..., M, M)` representing the coefficients matrix.
        b: A tensor of shape `(..., M)` or `(..., M, N)` representing the
        right-hand side or "dependent variable" matrix.

Returns:
        A tensor of shape `(..., M)` or `(..., M, N)` representing the solution
        of the linear system. Returned shape is identical to `b`.


def update_state_wrapper(update_state_fn):

    def decorated(metric_obj, *args, **kwargs):
        """Decorated function with `add_update()`."""
        strategy = distribute_lib.get_strategy()
        for weight in metric_obj.weights:
            if backend.is_tpu_strategy(strategy) and (not strategy.extended.variable_created_in_scope(weight)) and (not distribute_lib.in_cross_replica_context()):
                raise ValueError('Trying to run metric.update_state in replica context when the metric was not created in TPUStrategy scope. Make sure the keras Metric is created in TPUstrategy scope. ')
        with tf_utils.graph_context_for_symbolic_tensors(*args, **kwargs):
            update_op = update_state_fn(*args, **kwargs)
        if update_op is not None:
            metric_obj.add_update(update_op)
        return update_op
    return tf_decorator.make_decorator(update_state_fn, decorated)

Decorator to wrap metric `update_state()` with `add_update()`.

Args:
        update_state_fn: function that accumulates metric statistics.

Returns:
        Decorated function that wraps `update_state_fn()` with `add_update()`.


def __init__(self, underlying_result, pipeline_instrument):
    super().__init__(underlying_result.state)
    self._underlying_result = underlying_result
    self._pipeline_instrument = pipeline_instrument

Constructor of PipelineResult.

Args:
        underlying_result: (PipelineResult) the result returned by the underlying
        runner running the pipeline.
        pipeline_instrument: (PipelineInstrument) pipeline instrument describing
        the pipeline being executed with interactivity applied and related
        metadata including where the interactivity-backing cache lies.


def recipe_dv360_segmentology(config, auth_read, recipe_timezone, auth_write, recipe_name, date_range, recipe_slug, partners, advertisers):
    dataset(config, {'description': 'Create a dataset for bigquery tables.', 'hour': [4], 'auth': auth_write, 'dataset': recipe_slug})
    bigquery(config, {'auth': auth_write, 'function': 'Pearson Significance Test', 'to': {'dataset': recipe_slug}})
    dbm(config, {'auth': auth_read, 'report': {'filters': {'FILTER_PARTNER': {'values': partners}, 'FILTER_ADVERTISER': {'values': advertisers}}, 'body': {'timezoneCode': recipe_timezone, 'metadata': {'title': recipe_name, 'dataRange': date_range, 'format': 'CSV'}, 'params': {'type': 'TYPE_CROSS_PARTNER', 'groupBys': ['FILTER_PARTNER', 'FILTER_PARTNER_NAME', 'FILTER_ADVERTISER', 'FILTER_ADVERTISER_NAME', 'FILTER_MEDIA_PLAN', 'FILTER_MEDIA_PLAN_NAME', 'FILTER_ZIP_POSTAL_CODE'], 'metrics': ['METRIC_BILLABLE_IMPRESSIONS', 'METRIC_CLICKS', 'METRIC_TOTAL_CONVERSIONS']}, 'schedule': {'frequency': 'WEEKLY'}}}})
    dbm(config, {'auth': auth_read, 'report': {'name': recipe_name}, 'out': {'bigquery': {'auth': auth_write, 'dataset': recipe_slug, 'table': 'DV360_KPI', 'header': True, 'schema': [{'name': 'Partner_Id', 'type': 'INTEGER', 'mode': 'REQUIRED'}, {'name': 'Partner', 'type': 'STRING', 'mode': 'REQUIRED'}, {'name': 'Advertiser_Id', 'type': 'INTEGER', 'mode': 'REQUIRED'}, {'name': 'Advertiser', 'type': 'STRING', 'mode': 'REQUIRED'}, {'name': 'Campaign_Id', 'type': 'INTEGER', 'mode': 'REQUIRED'}, {'name': 'Campaign', 'type': 'STRING', 'mode': 'REQUIRED'}, {'name': 'Zip', 'type': 'STRING', 'mode': 'NULLABLE'}, {'name': 'Impressions', 'type': 'FLOAT', 'mode': 'NULLABLE'}, {'name': 'Clicks', 'type': 'FLOAT', 'mode': 'NULLABLE'}, {'name': 'Conversions', 'type': 'FLOAT', 'mode': 'NULLABLE'}]}}})
    bigquery(config, {'auth': auth_write, 'from': {'query': 'SELECT\n          Partner_Id,\n          Partner,\n          Advertiser_Id,\n          Advertiser,\n          Campaign_Id,\n          Campaign,\n          Zip,\n          SAFE_DIVIDE(Impressions, SUM(Impressions) OVER(PARTITION BY Advertiser_Id)) AS Impression,\n          SAFE_DIVIDE(Clicks, Impressions) AS Click,\n          SAFE_DIVIDE(Conversions, Impressions) AS Conversion,\n          Impressions AS Impressions          FROM\n          `{dataset}.DV360_KPI`;        ', 'parameters': {'dataset': recipe_slug}, 'legacy': False}, 'to': {'dataset': recipe_slug, 'view': 'DV360_KPI_Normalized'}})
    census(config, {'auth': auth_write, 'normalize': {'census_geography': 'zip_codes', 'census_year': '2018', 'census_span': '5yr'}, 'to': {'dataset': recipe_slug, 'type': 'view'}})
    census(config, {'auth': auth_write, 'correlate': {'join': 'Zip', 'pass': ['Partner_Id', 'Partner', 'Advertiser_Id', 'Advertiser', 'Campaign_Id', 'Campaign'], 'sum': ['Impressions'], 'correlate': ['Impression', 'Click', 'Conversion'], 'dataset': recipe_slug, 'table': 'DV360_KPI_Normalized', 'significance': 80}, 'to': {'dataset': recipe_slug, 'type': 'view'}})

DV360 funnel analysis using Census data.

Args:
        auth_read (authentication) - Credentials used for reading data.
        recipe_timezone (timezone) - Timezone for report dates.
        auth_write (authentication) - Authorization used for writing data.
        recipe_name (string) - Name of report, not needed if ID used.
        date_range (choice) - Timeframe to run the report for.
        recipe_slug (string) - Name of Google BigQuery dataset to create.
        partners (integer_list) - DV360 partner id.
        advertisers (integer_list) - Comma delimited list of DV360 advertiser ids.


def get_cross_replica_context():
    return _get_per_thread_mode().cross_replica_context

Returns the current tf.distribute.Strategy if in a cross-replica context.

    DEPRECATED: Please use `in_cross_replica_context()` and
    `get_strategy()` instead.

Returns:
        Returns the current `tf.distribute.Strategy` object in a cross-replica
        context, or `None`.

        Exactly one of `get_replica_context()` and `get_cross_replica_context()`
        will return `None` in a particular block.


def __init__(self, cache_config: CacheConfig) -> None:
    super().__init__(cache_config)
    if is_optimum_quanto_available():
        optimum_quanto_version = version.parse(importlib.metadata.version('optimum-quanto'))
        if optimum_quanto_version <= version.parse('0.2.5'):
            raise ImportError(f'You need optimum-quanto package version to be greater or equal than 0.2.5 to use `QuantoQuantizedCache`. Detected version {optimum_quanto_version}.')
        from optimum.quanto import MaxOptimizer, qint2, qint4
    if self.nbits not in [2, 4]:
        raise ValueError(f'`nbits` for `quanto` backend has to be one of [`2`, `4`] but got {self.nbits}')
    if self.axis_key not in [0, -1]:
        raise ValueError(f'`axis_key` for `quanto` backend has to be one of [`0`, `-1`] but got {self.axis_key}')
    if self.axis_value not in [0, -1]:
        raise ValueError(f'`axis_value` for `quanto` backend has to be one of [`0`, `-1`] but got {self.axis_value}')
    self.qtype = qint4 if self.nbits == 4 else qint2
    self.optimizer = MaxOptimizer()

Quantized Cache class that uses `quanto` as a backend to perform quantization. Current implementation supports `int2` and `int4` dtypes only.

    Parameters:
        cache_config (`QuantizedCacheConfig`):
            A configuration containing all the arguments to be used by the quantizer, including axis, qtype and group size.

Example:
        ```python
        >>> # Run pip install quanto first if you don't have it yet
        >>> from transformers import AutoTokenizer, AutoModelForCausalLM, QuantoQuantizedCache, QuantizedCacheConfig

        >>> model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
        >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")

        >>> inputs = tokenizer(text="My name is Qwen2", return_tensors="pt")

        >>> # Prepare a cache class and pass it to model's forward
        >>> cache_config = QuantizedCacheConfig(nbits=4)
        >>> past_key_values = QuantoQuantizedCache(cache_config=cache_config)
        >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True)
        >>> outputs.past_key_values # access cache filled with key/values from generation
        QuantoQuantizedCache()
        ```


def get_config(self):
    all_args = tf_inspect.getfullargspec(self.__init__).args
    config = {'name': self.name, 'trainable': self.trainable}
    if hasattr(self, '_batch_input_shape'):
        config['batch_input_shape'] = self._batch_input_shape
    config['dtype'] = policy.serialize(self._dtype_policy)
    if hasattr(self, 'dynamic'):
        if self.dynamic:
            config['dynamic'] = self.dynamic
        elif 'dynamic' in all_args:
            all_args.remove('dynamic')
    expected_args = config.keys()
    extra_args = [arg for arg in all_args if arg not in expected_args]
    if len(extra_args) > 1 and hasattr(self.get_config, '_is_default'):
        raise NotImplementedError('Layer %s has arguments in `__init__` and therefore must override `get_config`.' % self.__class__.__name__)
    return config

Returns the config of the layer.

    A layer config is a Python dictionary (serializable)
    containing the configuration of a layer.
    The same layer can be reinstantiated later
    (without its trained weights) from this configuration.

    The config of a layer does not include connectivity
    information, nor the layer class name. These are handled
    by `Network` (one layer of abstraction above).

    Note that `get_config()` does not guarantee to return a fresh copy of dict
    every time it is called. The callers should make a copy of the returned dict
    if they want to modify it.

Returns:
        Python dictionary.


def resource_path(package: Union[str, types.ModuleType]) -> abstract_path.Path:
    try:
        path = importlib_resources.files(package)
    except AttributeError:
        is_adhoc = True
    else:
        if isinstance(path, importlib_resources._adapters.CompatibilityFiles.SpecPath):
            is_adhoc = True
        else:
            is_adhoc = False
    if is_adhoc:
        if isinstance(package, types.ModuleType):
            package = getattr(package.__spec__, 'name', package.__name__)
        path = pathlib.Path(sys.modules[package].__file__)
        if path.name == '__init__.py':
            path = path.parent
    if isinstance(path, pathlib.Path):
        return abstract_path.Path(path)
    elif isinstance(path, zipfile.Path):
        path = ResourcePath(path.root, path.at)
        return typing.cast(abstract_path.Path, path)
    elif isinstance(path, importlib_resources.abc.Traversable):
        return typing.cast(abstract_path.Path, path)
    else:
        raise TypeError(f'Unknown resource path: {type(path)}: {path}')

Returns read-only root directory path of the module.

    Used to access module resource files.

    Usage:

    ```python
    path = epath.resource_path('tensorflow_datasets') / 'README.md'
    content = path.read_text()
    ```

    This is compatible with everything, including zipapp (`.par`).

    Resource files should be in the `data=` of the `py_library(` (when using
    bazel).

    To write to your project (e.g. automatically update your code), read-only
    resource paths can be converted to read-write paths with
    `epath.to_write_path(path)`.

Args:
        package: Module or module name.

Returns:
        The read-only path to the root module directory


def register_command_handler(self, prefix, handler, help_info, prefix_aliases=None):
    if not prefix:
        raise ValueError('Empty command prefix')
    if prefix in self._handlers:
        raise ValueError('A handler is already registered for command prefix "%s"' % prefix)
    if not callable(handler):
        raise ValueError('handler is not callable')
    if not isinstance(help_info, str):
        raise ValueError('help_info is not a str')
    if prefix_aliases:
        for alias in prefix_aliases:
            if self._resolve_prefix(alias):
                raise ValueError('The prefix alias "%s" clashes with existing prefixes or aliases.' % alias)
            self._alias_to_prefix[alias] = prefix
        self._prefix_to_aliases[prefix] = prefix_aliases
    self._handlers[prefix] = handler
    self._prefix_to_help[prefix] = help_info

Register a callable as a command handler.

Args:
        prefix: Command prefix, i.e., the first word in a command, e.g.,
        "print" as in "print tensor_1".
        handler: A callable of the following signature:
        foo_handler(argv, screen_info=None),
        where argv is the argument vector (excluding the command prefix) and
        screen_info is a dictionary containing information about the screen,
        such as number of columns, e.g., {"cols": 100}.
        The callable should return:
        1) a RichTextLines object representing the screen output.

        The callable can also raise an exception of the type CommandLineExit,
        which if caught by the command-line interface, will lead to its exit.
        The exception can optionally carry an exit token of arbitrary type.
        help_info: A help string.
        prefix_aliases: Aliases for the command prefix, as a list of str. E.g.,
        shorthands for the command prefix: ["p", "pr"]

Raises:
        ValueError: If
        1) the prefix is empty, or
        2) handler is not callable, or
        3) a handler is already registered for the prefix, or
        4) elements in prefix_aliases clash with existing aliases.
        5) help_info is not a str.


def tpu_devices(devices=None):
    return find_devices('TPU', devices)

Gets TPU devices out of `devices`.

Args:
        devices: A device list (as a list of strings). If None, the list of all
        available devices will be used for it.

Returns:
        Those in `devices` that are TPUs.


def run_suite(test_classes, argv=None):
    args = _parse_cli_args(argv)
    for test_class in test_classes:
        if not issubclass(test_class, base_test.BaseTestClass):
            logging.error('Test class %s does not extend mobly.base_test.BaseTestClass', test_class)
            sys.exit(1)
    if args.list_tests:
        _print_test_names(test_classes)
        sys.exit(0)
    test_configs = config_parser.load_test_config_file(args.config, args.test_bed)
    selected_tests = compute_selected_tests(test_classes, args.tests)
    console_level = logging.DEBUG if args.verbose else logging.INFO
    ok = True
    for config in test_configs:
        runner = test_runner.TestRunner(config.log_path, config.testbed_name)
        with runner.mobly_logger(console_level=console_level):
            for test_class, tests in selected_tests.items():
                runner.add_test_class(config, test_class, tests)
            try:
                runner.run()
                ok = runner.results.is_all_pass and ok
            except signals.TestAbortAll:
                pass
            except Exception:
                logging.exception('Exception when executing %s.', config.testbed_name)
                ok = False
    if not ok:
        sys.exit(1)

Executes multiple test classes as a suite.

    This is the default entry point for running a test suite script file
    directly.

Args:
        test_classes: List of python classes containing Mobly tests.
        argv: A list that is then parsed as cli args. If None, defaults to cli
        input.


def clean_code(content: str) -> str:
    splits = content.split('"""')
    content = ''.join(splits[::2])
    splits = content.split("'''")
    content = ''.join(splits[::2])
    lines_to_keep = []
    for line in content.split('\n'):
        line = re.sub('#.*$', '', line)
        if len(line) != 0 and (not line.isspace()):
            lines_to_keep.append(line)
    return '\n'.join(lines_to_keep)

Remove docstrings, empty line or comments from some code (used to detect if a diff is real or only concern
    comments or docstings).

Args:
        content (`str`): The code to clean

Returns:
        `str`: The cleaned code.


def resize_position_embeddings(self, new_num_position_embeddings: int):
    self.distilbert.resize_position_embeddings(new_num_position_embeddings)

Resizes position embeddings of the model if `new_num_position_embeddings != config.max_position_embeddings`.

Args:
        new_num_position_embeddings (`int`):
            The number of new position embedding matrix. If position embeddings are learned, increasing the size
            will add newly initialized vectors at the end, whereas reducing the size will remove vectors from the
            end. If position embeddings are not learned (*e.g.* sinusoidal position embeddings), increasing the
            size will add correct vectors at the end following the position encoding algorithm, whereas reducing
            the size will remove vectors from the end.


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

labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
    Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
    config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked),
    the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
    next_sentence_label (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
    Labels for computing the next sequence prediction (classification) loss. 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, BertForPreTraining
    >>> import torch

    >>> tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
    >>> model = BertForPreTraining.from_pretrained("google-bert/bert-base-uncased")

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

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


def make_rhs(self, operator, adjoint, with_batch=True):
    raise NotImplementedError('make_rhs is not defined.')

Make a rhs appropriate for calling operator.solve(rhs).

Args:
        operator:  A `LinearOperator`
        adjoint:  Python `bool`.  If `True`, we are making a 'rhs' value for the
        adjoint operator.
        with_batch: Python `bool`. If `True`, create `rhs` with the same batch
        shape as operator, and otherwise create a matrix without any batch
        shape.

Returns:
        A `Tensor`


def generate_link(flag, np_fun_name):
    if flag == 'dev':
        template = 'https://numpy.org/devdocs/reference/generated/numpy.%s.html'
    elif flag == 'stable':
        template = 'https://numpy.org/doc/stable/reference/generated/numpy.%s.html'
    elif re.match('\\d+(\\.\\d+(\\.\\d+)?)?$', flag):
        template = f'https://numpy.org/doc/{flag}/reference/generated/numpy.%s.html'
    else:
        return None
    return template % np_fun_name

Generates link from numpy function name.

Args:
        flag: the flag to control link form. See `set_np_doc_form`.
        np_fun_name: the numpy function name.

Returns:
        A string.


def __init__(self, text='', font_attr=None):
    """Construct a RichLine with no rich attributes or a single attribute.

Args:
  text: Raw text string
  font_attr: If specified, a single font attribute to be applied to the
    entire text.  Extending this object via concatenation allows creation
    of text with varying attributes.
"""
    self.text = text
    if font_attr:
        self.font_attr_segs = [(0, len(text), font_attr)]
    else:
        self.font_attr_segs = []

Rich single-line text.

Attributes:
        text: A plain string, the raw text represented by this object.  Should not
        contain newlines.
        font_attr_segs: A list of (start, end, font attribute) triples, representing
        richness information applied to substrings of text.


def _get_normalizations(prices, forwards, strikes, discount_factors):
    strikes_abs = tf.abs(strikes)
    forwards_abs = tf.abs(forwards)
    orientations = strikes_abs >= forwards_abs
    normalization = tf.where(orientations, strikes_abs, forwards_abs)
    normalization = tf.where(tf.equal(normalization, 0), tf.ones_like(normalization), normalization)
    normalized_prices = prices / normalization
    if discount_factors is not None:
        normalized_prices /= discount_factors
    else:
        discount_factors = tf.ones_like(normalized_prices)
    return (normalized_prices, normalization, discount_factors)

Returns the normalized prices, normalization factors, and discount_factors.

    The normalization factors is the larger of strikes and forwards.
    If `discount_factors` is not None, these are the discount factors to expiry.
    If None, no discounting is applied and 1's are returned.

Args:
        prices: A real `Tensor` of any shape. The observed market prices of the
        assets.
        forwards: A real `Tensor` of the same shape and dtype as `prices`. The
        current forward prices to expiry.
        strikes: A real `Tensor` of the same shape and dtype as `prices`. The strike
        prices of the options.
        discount_factors: A real `Tensor` of same dtype as the `prices`.

Returns:
        the normalized prices, normalization factors, and discount_factors.


def get_rot_mats(self) -> torch.Tensor:
    if self._rot_mats is not None:
        return self._rot_mats
    elif self._quats is not None:
        return quat_to_rot(self._quats)
    else:
        raise ValueError('Both rotations are None')

Returns the underlying rotation as a rotation matrix tensor.

Returns:
        The rotation as a rotation matrix tensor


def __init__(self, proxy, include_indexes=False):
    self._proxy = proxy
    self._include_indexes = include_indexes

A transform that explodes a PCollection of DataFrame or Series. DataFrame
    is converterd to a schema-aware PCollection, while Series is converted to its
    underlying type.

Args:
        include_indexes: (optional, default: False) When unbatching a DataFrame
        if include_indexes=True, attempt to include index columns in the output
        schema for expanded DataFrames. Raises an error if any of the index
        levels are unnamed (name=None), or if any of the names are not unique
        among all column and index names.


def _split_list_into_bundles(self, output_pcollection, elements, max_element_per_bundle, element_size_fn):
    bundle = self._evaluation_context.create_bundle(output_pcollection)
    bundle_size = 0
    bundles = [bundle]
    for element in elements:
        if max_element_per_bundle and bundle_size >= max_element_per_bundle:
            bundle = self._evaluation_context.create_bundle(output_pcollection)
            bundle_size = 0
            bundles.append(bundle)
        bundle.output(element)
        bundle_size += element_size_fn(element)
    return bundles

Splits elements, an iterable, into multiple output bundles.

Args:
        output_pcollection: PCollection that the elements belong to.
        elements: elements to be chunked into bundles.
        max_element_per_bundle: (approximately) the maximum element per bundle.
        If it is None, only a single bundle will be produced.
        element_size_fn: Function to return the size of a given element.

Returns:
        List of output uncommitted bundles with at least one bundle.


def __init__(self):
    object.__setattr__(self, '_options', {})
    object.__setattr__(self, '_mutable', True)

Base class for representing a set of tf.data options.

Attributes:
        _options: Stores the option values.


def copy(self, stateful=False):
    raise NotImplementedError(f'{self.__class__} is an abstract class. Only classes inheriting this class can be called.')

Creates a new instance of this constraint.

Args:
        stateful(`bool`): Whether to not only copy the constraint for new instance, but also its state.

Returns:
        constraint(`Constraint`): The same constraint as the one being called from.


def _init_from_converter(self, options: QuantizationDebugOptions, converter: TFLiteConverter, calibrated_model: Optional[bytes]=None, float_model: Optional[bytes]=None) -> None:
    self.quant_model = convert.mlir_quantize(calibrated_model, disable_per_channel=converter._experimental_disable_per_channel, fully_quantize=options.fully_quantize, enable_numeric_verify=True, denylisted_ops=options.denylisted_ops, denylisted_nodes=options.denylisted_nodes)
    self._quant_interpreter = _interpreter.Interpreter(model_content=self.quant_model)
    self._float_interpreter = None
    if float_model is not None:
        self._float_interpreter = _interpreter.Interpreter(model_content=float_model)

Convert the model and apply options.

    Converts the quantized model and initializes a quantized model interpreter
    with the quantized model. Returns a float model interpreter if float model
    is provided.

Args:
        options: a QuantizationDebugOptions object.
        converter: an initialized tf.lite.TFLiteConverter.
        calibrated_model: Calibrated model bytes.
        float_model: Float model bytes.


def _block_orth(self, p1, p2, p3):
    p1_shape = p1.shape.as_list()
    if p1_shape != p2.shape.as_list() or p1_shape != p3.shape.as_list():
        raise ValueError(f'The dimension of the matrices must be the same. Received p1.shape={p1.shape}, p2.shape={p2.shape} and p3.shape={p3.shape}.')
    n = p1_shape[0]
    eye = linalg_ops_impl.eye(n, dtype=self.dtype)
    kernel2x2x2 = {}

    def matmul(p1, p2, p3):
        return math_ops.matmul(math_ops.matmul(p1, p2), p3)

    def cast(i, p):
        """Return p or (1-p)."""
        return i * p + (1 - i) * (eye - p)
    for i in [0, 1]:
        for j in [0, 1]:
            for k in [0, 1]:
                kernel2x2x2[i, j, k] = matmul(cast(i, p1), cast(j, p2), cast(k, p3))
    return kernel2x2x2

Construct a 3 x 3 kernel.

    Used to construct orthgonal kernel.

Args:
        p1: A symmetric projection matrix.
        p2: A symmetric projection matrix.
        p3: A symmetric projection matrix.

Returns:
        A 2 x 2 x 2 kernel.

Raises:
        ValueError: If the dimensions of p1, p2 and p3 are different.


def l2_normalize(x, axis=None):
    return nn.l2_normalize(x, axis=axis)

Normalizes a tensor wrt the L2 norm alongside the specified axis.

Args:
        x: Tensor or variable.
        axis: axis along which to perform normalization.

Returns:
        A tensor.


def parser(self, column: str, parser: str, error: str, value: Any) -> None:
    log = self._build_parser_message(column, parser, error, value)
    self.queue_log_message(log)

Adds parser error information to base log message and
    sends it to the logger for writing.

Args:
        * column: column where the rule is applied
        * parser: parser function that failed and raises this message
        * error: error that occurred
        * value: value that fails to parse

Returns:
        * None


def actual_360(*, start_date, end_date, schedule_info=None, dtype=None, name=None):
    del schedule_info
    with tf.name_scope(name or 'actual_360'):
        end_date = dt.convert_to_date_tensor(end_date)
        start_date = dt.convert_to_date_tensor(start_date)
        dtype = dtype or tf.constant(0.0).dtype
        actual_days = tf.cast(start_date.days_until(end_date), dtype=dtype)
        return actual_days / 360

Computes the year fraction between the specified dates.

    The actual/360 convention specifies the year fraction between the start and
    end date as the actual number of days between the two dates divided by 360.

    Note that the schedule info is not needed for this convention and is ignored
    if supplied.

    For more details see:
    https://en.wikipedia.org/wiki/Day_count_convention#Actual/360

Args:
        start_date: A `DateTensor` object of any shape.
        end_date: A `DateTensor` object of compatible shape with `start_date`.
        schedule_info: The schedule info. Ignored for this convention.
        dtype: The dtype of the result. Either `tf.float32` or `tf.float64`. If not
        supplied, `tf.float32` is returned.
        name: Python `str` name prefixed to ops created by this function. If not
        supplied, `actual_360` is used.

Returns:
        A real `Tensor` of supplied `dtype` and shape of `start_date`. The year
        fraction between the start and end date as computed by Actual/360
        convention.


def save_to_file(json_data, filename):
    if filename[-5:] != '.json':
        print('filename: %s' % filename)
        filename += '.json'
    with open(PATH_TO_DIR + '/' + filename, 'w') as f:
        json.dump(json_data, f, sort_keys=True, indent=4)
    print(' Successfully wrote configs to file `%s`.\n' % filename)

Saves all detected configuration(s) into a JSON file.

Args:
        json_data: Dict of all configurations found.
        filename: String that is the name of the output JSON file.


def stop(self):
    self.log.debug('Stopping snippet package %s.', self.package)
    self.close_connection()
    self._stop_server()
    self._destroy_event_client()
    self.log.debug('Snippet package %s stopped.', self.package)

Releases all the resources acquired in `initialize`.

    This function releases following resources:
    * Close the socket connection.
    * Stop forwarding the device port to host.
    * Stop the standing server subprocess running on the host side.
    * Stop the snippet server running on the device side.
    * Stop the event client and set `self._event_client` to None.

Raises:
        android_device_lib_errors.DeviceError: if the server exited with errors on
        the device side.


def _CheckLineJoining(self, code, join_lines):
    llines = yapf_test_helper.ParseAndUnwrap(code)
    self.assertCodeEqual(line_joiner.CanMergeMultipleLines(llines), join_lines)

Check that the given LogicalLines are joined as expected.

Args:
        code: The code to check to see if we can join it.
        join_lines: True if we expect the lines to be joined.


def attach_tracer(tracer_name_template=None):
    if not _has_opentelemetry:
        return lambda cls: cls

    def decorator(cls):
        original_init = cls.__init__

        @functools.wraps(original_init)
        def init_with_tracer(self, *args, **kwargs):
            original_init(self, *args, **kwargs)
            module_name = cls.__module__
            class_name = cls.__qualname__
            if tracer_name_template is None:
                if module_name.startswith('transformers.'):
                    tracer_name = f'{module_name}.{class_name}'
                else:
                    tracer_name = f'transformers.{module_name}.{class_name}'
            else:
                tracer_name = tracer_name_template.format(module=module_name, class_name=class_name)
            self.tracer = get_tracer(tracer_name)
        cls.__init__ = init_with_tracer
        return cls
    return decorator

Decorator that attaches a tracer to a class.

    This decorator should be applied to classes that need OpenTelemetry tracing.
    It adds a tracer attribute to the class instance that can be used by the traced decorator.

Args:
        tracer_name_template: Optional template string for the tracer name.
            If provided, it should contain {module} which will be replaced with the class's full module path
            and {class_name} for the class name.
            If None, a default naming scheme will be used where:
            - If the module already starts with "transformers.", it will use that directly
            - Otherwise, it will prepend "transformers." to the module name

Returns:
        Class decorator function


def get_init_tokens_op(self, num_tokens=-1):
    if self._gradients_applied is False:
        raise ValueError('get_init_tokens_op() should be called after apply_gradients().')
    tokens_needed = self._replicas_to_aggregate - self._total_num_replicas
    if num_tokens == -1:
        num_tokens = self._replicas_to_aggregate
    elif num_tokens < tokens_needed:
        raise ValueError('Too few tokens to finish the first step: %d (given) vs %d (needed)' % (num_tokens, tokens_needed))
    if num_tokens > 0:
        with ops.device(self._global_step.device), ops.name_scope(''):
            tokens = array_ops.fill([num_tokens], self._global_step)
            init_tokens = self._sync_token_queue.enqueue_many((tokens,))
    else:
        init_tokens = control_flow_ops.no_op(name='no_init_tokens')
    return init_tokens

Returns the op to fill the sync_token_queue with the tokens.

    This is supposed to be executed in the beginning of the chief/sync thread
    so that even if the total_num_replicas is less than replicas_to_aggregate,
    the model can still proceed as the replicas can compute multiple steps per
    variable update. Make sure:
    `num_tokens >= replicas_to_aggregate - total_num_replicas`.

Args:
        num_tokens: Number of tokens to add to the queue.

Returns:
        An op for the chief/sync replica to fill the token queue.

Raises:
        ValueError: If this is called before apply_gradients().
        ValueError: If num_tokens are smaller than replicas_to_aggregate -
        total_num_replicas.


def get_extended_attention_mask(self, attention_mask: tf.Tensor, input_shape: Tuple[int], is_decoder: bool) -> tf.Tensor:
    if not isinstance(attention_mask, tf.Tensor):
        attention_mask = tf.convert_to_tensor(attention_mask)
    if attention_mask.shape.rank == 3:
        extended_attention_mask = attention_mask[:, None, :, :]
    elif attention_mask.shape.rank == 2:
        if is_decoder:
            batch_size, seq_length = input_shape
            seq_ids = tf.range(seq_length, dtype=attention_mask.dtype)
            causal_mask = tf.broadcast_to(seq_ids, (batch_size, seq_length, seq_length)) <= seq_ids[None, :, None]
            if shape_list(causal_mask)[1] < shape_list(attention_mask)[1]:
                prefix_seq_len = tf.shape(attention_mask)[1] - tf.shape(causal_mask)[1]
                causal_mask = tf.concat([tf.ones((batch_size, seq_length, prefix_seq_len), dtype=causal_mask.dtype), causal_mask], axis=-1)
            extended_attention_mask = tf.cast(causal_mask[:, None, :, :], attention_mask.dtype) * attention_mask[:, None, None, :]
        else:
            extended_attention_mask = attention_mask[:, None, None, :]
    else:
        raise ValueError('Wrong shape for input_ids (shape {}) or attention_mask (shape {})'.format(input_shape, attention_mask.shape))
    extended_attention_mask = tf.cast(extended_attention_mask, self.dtype)
    extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
    return extended_attention_mask

Makes broadcastable attention and causal masks so that future and masked tokens are ignored.

Args:
        attention_mask (`tf.Tensor`):
            Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
        input_shape (`Tuple[int]`):
            The shape of the input to the model.
        is_decoder (`bool`):
            Whether the model is used as a decoder.

Returns:
        `tf.Tensor` The extended attention mask, with the same dtype as `attention_mask.dtype`.


def _add_state_variable(self, name, shape, dtype, initializer=None, partitioner=None, use_resource=None, **kwargs):
    weight = self.add_weight(name=name, shape=shape, dtype=dtype, initializer=initializer, regularizer=None, trainable=False, constraint=None, partitioner=partitioner, use_resource=use_resource, **kwargs)
    self.state_variables[name] = weight
    return weight

Add a variable that can hold state which is updated during adapt().

Args:
        name: Variable name.
        shape: Variable shape. Defaults to scalar if unspecified.
        dtype: The type of the variable. Defaults to `self.dtype` or `float32`.
        initializer: initializer instance (callable).
        partitioner: Partitioner to be passed to the `Trackable` API.
        use_resource: Whether to use `ResourceVariable`
        **kwargs: Additional keyword arguments. Accepted values are `getter` and
        `collections`.

Returns:
        The created variable.


def get_calibration_min_max_value(self, calibration_statistics_serialized: bytes, calibration_options_serialized: bytes) -> Optional[tuple[float, float]]:
    statistics = calibration_statistics_pb2.CalibrationStatistics.FromString(calibration_statistics_serialized)
    options = stablehlo_quant_config_pb2.CalibrationOptions.FromString(calibration_options_serialized)
    return _call_and_return_none_on_error(functools.partial(calibration_algorithm.get_min_max_value, statistics, options), error_msg=f'Retrieving calibrated min / max failed. Options: {options}.')

Calculates min and max values from statistics.

Args:
        calibration_statistics_serialized: Serialized `CalibrationStatistics`.
        This will be the source to calculate min and max values from.
        calibration_options_serialized: Serialized `CalibrationOptions`. Specifies
        how the min / max should be calculated.

Returns:
        (min_value, max_value): Min and max calculated using calib_opts. `None`
        upon error.


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
    requirements_dir = os.path.dirname(os.path.realpath(__file__))
    pipeline_options.view_as(SetupOptions).requirements_file = f'{requirements_dir}/sklearn_examples_requirements.txt'
    model_loader = KeyedModelHandler(SklearnModelHandlerNumpy(model_file_type=ModelFileType.PICKLE, model_uri=known_args.model_path, large_model=known_args.large_model))
    pipeline = test_pipeline
    if not test_pipeline:
        pipeline = beam.Pipeline(options=pipeline_options)
    label_pixel_tuple = pipeline | 'ReadFromInput' >> beam.io.ReadFromText(known_args.input) | 'PreProcessInputs' >> beam.Map(process_input)
    predictions = label_pixel_tuple | 'RunInference' >> RunInference(model_loader) | 'PostProcessOutputs' >> beam.ParDo(PostProcessor())
    _ = predictions | '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 _dedup_strings(device_strs):
    new_device_strs = []
    for device_str, vals in itertools.groupby(device_strs):
        num = len(list(vals))
        if num == 1:
            new_device_strs.append(device_str)
        else:
            new_device_strs.append('%s (x%d)' % (device_str, num))
    return new_device_strs

Groups together consecutive identical strings.

    For example, given:
        ['GPU 1', 'GPU 2', 'GPU 2', 'GPU 3', 'GPU 3', 'GPU 3']
    This function returns:
        ['GPU 1', 'GPU 2 (x2)', 'GPU 3 (x3)']

Args:
        device_strs: A list of strings, each representing a device.

Returns:
        A copy of the input, but identical consecutive strings are merged into a
        single string.


def forward(self, input_embeds: torch.FloatTensor) -> torch.FloatTensor:
    hidden_states = self.conv_pre(input_embeds)
    for i in range(self.num_upsamples):
        hidden_states = nn.functional.leaky_relu(hidden_states, self.leaky_relu_slope)
        hidden_states = self.upsampler[i](hidden_states)
        res_state = self.resblocks[i * self.num_kernels](hidden_states)
        for j in range(1, self.num_kernels):
            res_state += self.resblocks[i * self.num_kernels + j](hidden_states)
        hidden_states = res_state / self.num_kernels
    hidden_states = nn.functional.leaky_relu(hidden_states)
    hidden_states = self.conv_post(hidden_states)
    hidden_states = torch.tanh(hidden_states)
    waveform = hidden_states.squeeze(1)
    return waveform

Converts a log-mel spectrogram into a speech waveform. Passing a batch of log-mel spectrograms returns a batch
    of speech waveforms. Passing a single, un-batched log-mel spectrogram returns a single, un-batched speech
    waveform.

Args:
        spectrogram (`torch.FloatTensor`):
            Tensor containing the log-mel spectrograms. Can be batched and of shape `(batch_size, sequence_length,
            model_in_dim)`, or un-batched and of shape `(sequence_length, model_in_dim)`. Note that `model_in_dim`
            is the sum of `config.unit_embed_dim`, `config.lang_embed_dim` and `config.spkr_embed_dim`.

Returns:
        `torch.FloatTensor`: Tensor containing the speech waveform. If the input spectrogram is batched, will be of
        shape `(batch_size, num_frames,)`. If un-batched, will be of shape `(num_frames,)`.


def _finish(self, update_ops, name_scope):
    return control_flow_ops.group(*update_ops, name=name_scope)

Do what is needed to finish the update.

    This is called with the `name_scope` using the "name" that
    users have chosen for the application of gradients.

Args:
        update_ops: List of `Operation` objects to update variables.  This list
        contains the values returned by the `_apply_dense()` and
        `_apply_sparse()` calls.
        name_scope: String.  Name to use for the returned operation.

Returns:
        The operation to apply updates.


def recipe_dbm_to_bigquery(config, auth_read, auth_write, dbm_report_id, dbm_report_name, dbm_dataset, dbm_table, dbm_schema, is_incremental_load):
    dbm(config, {'auth': auth_read, 'report': {'report_id': dbm_report_id, 'name': dbm_report_name}, 'out': {'bigquery': {'auth': auth_write, 'dataset': dbm_dataset, 'table': dbm_table, 'schema': dbm_schema, 'header': True, 'is_incremental_load': is_incremental_load}}})

Move existing DV360 reports into a BigQuery table.

Args:
        auth_read (authentication) - Credentials used for reading data.
        auth_write (authentication) - Authorization used for writing data.
        dbm_report_id (integer) - DV360 report ID given in UI, not needed if name used.
        dbm_report_name (string) - Name of report, not needed if ID used.
        dbm_dataset (string) - Existing BigQuery dataset.
        dbm_table (string) - Table to create from this report.
        dbm_schema (json) - Schema provided in JSON list format or empty value to auto detect.
        is_incremental_load (boolean) - Clear data in destination table during this report's time period, then append report data to destination table.


def get_generating_ops(ts):
    ts = make_list_of_t(ts, allow_graph=False)
    return [t.op for t in ts]

Return all the generating ops of the tensors in `ts`.

Args:
        ts: a list of `tf.Tensor`

Returns:
        A list of all the generating `tf.Operation` of the tensors in `ts`.

Raises:
        TypeError: if `ts` cannot be converted to a list of `tf.Tensor`.


def forward(self, hidden_states: torch.Tensor, freqs_ci: torch.Tensor, attention_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
    encoder_states = () if output_hidden_states else None
    all_attentions = () if output_attentions else None
    for encoder_layer in self.layers:
        if output_hidden_states:
            encoder_states = encoder_states + (hidden_states,)
        if self.gradient_checkpointing and self.training:
            layer_outputs = self._gradient_checkpointing_func(encoder_layer.__call__, hidden_states, freqs_ci, attention_mask, output_attentions)
        else:
            layer_outputs = encoder_layer(hidden_state=hidden_states, attention_mask=attention_mask, output_attentions=output_attentions, freqs_ci=freqs_ci)
        if output_attentions:
            all_attentions = all_attentions + (layer_outputs[1],)
        hidden_states = layer_outputs[0]
    if output_hidden_states:
        encoder_states = encoder_states + (hidden_states,)
    if not return_dict:
        return tuple((v for v in [hidden_states, encoder_states, all_attentions] if v is not None))
    return BaseModelOutput(last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions)

Args:
    inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
        Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
        This is useful if you want more control over how to convert `input_ids` indices into associated vectors
        than the model's internal embedding lookup matrix.
    attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
        Mask to avoid performing 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.
    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 random_dna(self, random_generator: Union[types.ModuleType, random.Random, None]=None, attach_spec: bool=True, previous_dna: Optional['DNA']=None) -> 'DNA':
    random_generator = random_generator or random
    dna = self._random_dna(random_generator, previous_dna)
    if attach_spec:
        dna.use_spec(self)
    return dna

Returns a random DNA based on current spec.

Args:
        random_generator: An optional Random object. If None, the global random
        module will be used.
        attach_spec: If True, current spec will be attached to the returned DNA.
        previous_dna: An optional DNA representing previous DNA. This field might
        be useful for generating stateful random DNAs.

Returns:
        A random DNA based on current spec.


def session_creator(self, scaffold=None, config=None, checkpoint_dir=None, checkpoint_filename_with_path=None, max_wait_secs=7200):
    if config:
        session_config = copy.deepcopy(config)
        session_config.MergeFrom(self._session_config)
    else:
        session_config = self._session_config
    if not self._strategy or self._strategy.extended.experimental_should_init:
        logging.info('Creating chief session creator with config: %r', config)
        return monitored_session.ChiefSessionCreator(scaffold, master=self.master_target, config=session_config, checkpoint_dir=checkpoint_dir, checkpoint_filename_with_path=checkpoint_filename_with_path)
    else:
        logging.info('Creating worker session creator with config: %r', config)
        return monitored_session.WorkerSessionCreator(scaffold, master=self.master_target, config=session_config, max_wait_secs=max_wait_secs)

Returns a session creator.

    The returned session creator will be configured with the correct master
    target and session configs. It will also run either init ops or ready ops
    by querying the `strategy` object when `create_session` is called on it.

Args:
        scaffold: A `Scaffold` used for gathering or building supportive ops. If
        not specified a default one is created. It's used to finalize the graph.
        config: `ConfigProto` proto used to configure the session.
        checkpoint_dir: A string. Optional path to a directory where to restore
        variables.
        checkpoint_filename_with_path: Full file name path to the checkpoint file.
        Only one of `checkpoint_dir` or `checkpoint_filename_with_path` can be
        specified.
        max_wait_secs: Maximum time to wait for the session to become available.

Returns:
        a descendant of SessionCreator.


def get_tensor_layout(self, path):
    raise NotImplementedError()

Retrieve the `TensorLayout` for the intermediate tensor.

Args:
        path: a string path for the corresponding tensor.

Returns:
        The `TensorLayout` for the intermediate tensor, which can be used
        by `backend.relayout()` to reshard the tensor. Could also return
        None.


def train_model(preprocessed_dataset_path: str, trained_model_path: str, base_artifact_path: str):
    timestamp = time.time()
    model = torch.hub.load('pytorch/vision:v0.10.0', 'vgg16', pretrained=True)
    target_path = f'{base_artifact_path}/training/trained_model_{timestamp}.pt'
    target_path_gcsfuse = target_path.replace('gs://', '/gcs/')
    Path(target_path_gcsfuse).parent.mkdir(parents=True, exist_ok=True)
    torch.save(model.state_dict(), target_path_gcsfuse)
    Path(trained_model_path).parent.mkdir(parents=True, exist_ok=True)
    with open(trained_model_path, 'w') as f:
        f.write(target_path)

Placeholder method to load a model from the torch hub and save it.

Args:
        preprocessed_dataset_path (str): Path to the preprocessed dataset
        trained_model_path (str): Output path for the trained model
        base_artifact_path (str): path to the base directory of where artifacts can be stored for
        this component


def symmetric_linear_quantization_params(num_bits, saturation_min, saturation_max, per_channel=False):
    with torch.no_grad():
        n = 2 ** (num_bits - 1) - 1
        if per_channel:
            scale, _ = torch.max(torch.stack([saturation_min.abs(), saturation_max.abs()], dim=1), dim=1)
            scale = torch.clamp(scale, min=1e-08) / n
        else:
            scale = max(saturation_min.abs(), saturation_max.abs())
            scale = torch.clamp(scale, min=1e-08) / n
    return scale

Compute the scaling factor with the given quantization range for symmetric quantization.

Args:
        saturation_min (`torch.Tensor`):
            Lower bound for quantization range.
        saturation_max (`torch.Tensor`):
            Upper bound for quantization range.
        per_channel (`bool`, *optional*, defaults to `False`):
            Whether to or not use channel-wise quantization.

Returns:
        `torch.Tensor`: Scaling factor that linearly quantizes the given range between *saturation_min* and
        *saturation_max*.


def _snapshot_device_function_stack_metadata(self) -> list[traceable_stack.TraceableObject]:
    snapshot = []
    for obj in self._device_function_stack.peek_traceable_objs():
        obj_copy = obj.copy_metadata()
        obj_copy.obj = obj.obj.display_name
        snapshot.append(obj_copy)
    return snapshot

Return device function stack as a list of TraceableObjects.

Returns:
        [traceable_stack.TraceableObject, ...] where each TraceableObject's .obj
        member is a displayable name for the user's argument to Graph.device, and
        the filename and lineno members point to the code location where
        Graph.device was called directly or indirectly by the user.


def node_name(self):
    return self._node_name

Name of the node from which the tensor value was dumped.

Returns:
        (`str`) name of the node watched by the debug op.


def set_session(session):
    global _SESSION
    _SESSION.session = session

Sets the global TensorFlow session.

Args:
        session: A TF Session.


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, encoder_hidden_states: Optional[torch.FloatTensor]=None, encoder_attention_mask: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, past_key_values: Optional[Tuple[Tuple[torch.Tensor]]]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, **kwargs) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    if labels is not None:
        use_cache = False
    outputs = self.bert(input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, past_key_values=past_key_values, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    sequence_output = outputs[0]
    prediction_scores = self.cls(sequence_output)
    lm_loss = None
    if labels is not None:
        lm_loss = self.loss_function(prediction_scores, labels, vocab_size=self.config.vocab_size, **kwargs)
    if not return_dict:
        output = (prediction_scores,) + outputs[2:]
        return (lm_loss,) + output if lm_loss is not None else output
    return CausalLMOutputWithCrossAttentions(loss=lm_loss, logits=prediction_scores, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, cross_attentions=outputs.cross_attentions)

labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
    Labels for computing the left-to-right language modeling loss (next word prediction). 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 n `[0, ..., config.vocab_size]`

Example:
    ```python
    >>> from transformers import AutoTokenizer, MegatronBertForCausalLM, MegatronBertConfig
    >>> import torch

    >>> tokenizer = AutoTokenizer.from_pretrained("nvidia/megatron-bert-cased-345m")
    >>> model = MegatronBertForCausalLM.from_pretrained("nvidia/megatron-bert-cased-345m", is_decoder=True)

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

    >>> prediction_logits = outputs.logits
    ```


def recipe_hello(config, auth_read, say_first, say_second, error, sleep):
    hello(config, {'auth': auth_read, 'say': say_first, 'error': error, 'sleep': sleep})
    hello(config, {'auth': auth_read, 'say': say_second, 'sleep': sleep})

Recipe template for say hello.

Args:
        auth_read (authentication) - Credentials used for reading data.
        say_first (string) - Type in a greeting.
        say_second (string) - Type in a greeting.
        error (string) - Optional error for testing.
        sleep (integer) - Seconds to sleep.