def forward(self, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.Tensor]=None, position_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]]=None, inputs_embeds: Optional[torch.FloatTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, cache_position: Optional[torch.LongTensor]=None, labels: Optional[torch.LongTensor]=None, logits_to_keep: Union[int, torch.Tensor]=0, **kwargs) -> Union[Tuple, MoshiCausalLMOutputWithPast]:
    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
    outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, cache_position=cache_position)
    hidden_states = outputs[0]
    if labels is None and (not is_torchdynamo_compiling()):
        logger.warning_once('Starting from v4.46, the `logits` model output will have the same type as the model (except at train time, where it will always be FP32)')
    slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
    logits = self.lm_head(hidden_states[:, slice_indices, :])
    loss = None
    if labels is not None:
        logits = logits.float()
        shift_logits = logits[..., :-1, :].contiguous()
        shift_labels = labels[..., 1:].contiguous()
        shift_logits = shift_logits.view(-1, self.config.vocab_size)
        shift_labels = shift_labels.view(-1)
        shift_labels = shift_labels.to(shift_logits.device)
        loss = self.loss_function(shift_logits, shift_labels, vocab_size=self.config.vocab_size, **kwargs)
    if not return_dict:
        output = (logits, hidden_states) + outputs[1:]
        return (loss,) + output if loss is not None else output
    return MoshiCausalLMOutputWithPast(loss=loss, logits=logits, last_hidden_state=hidden_states, past_key_values=outputs.past_key_values, 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 either be in `[0, ...,
    config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
    (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.

Example:
    ```python
    >>> from transformers import AutoTokenizer, MoshiForCausalLM

    >>> model = MoshiForCausalLM.from_pretrained("kmhf/hf-moshiko")
    >>> tokenizer = AutoTokenizer.from_pretrained("kmhf/hf-moshiko")

    >>> prompt = "What is your favorite condiment?"
    >>> inputs = tokenizer(prompt, return_tensors="pt")

    >>> # Generate
    >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
    >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
    "What is your favorite condiment?"
    ```


def triu(x, k=0):
    if any_symbolic_tensors((x,)):
        return Triu(k=k).symbolic_call(x)
    return backend.numpy.triu(x, k=k)

Return upper triangle of a tensor.

    For tensors with `ndim` exceeding 2, `triu` will apply to the
    final two axes.

Args:
        x: Input tensor.
        k: Diagonal below which to zero elements. Defaults to `0`. the
            main diagonal. `k < 0` is below it, and `k > 0` is above it.

Returns:
        Upper triangle of `x`, of same shape and data type as `x`.


def moments(x, axes, shift=None, name=None, keep_dims=None, keepdims=None):
    keep_dims = deprecated_argument_lookup('keepdims', keepdims, 'keep_dims', keep_dims)
    if keep_dims is None:
        keep_dims = False
    with ops.name_scope(name, 'moments', [x, axes]):
        y = math_ops.cast(x, dtypes.float32) if x.dtype == dtypes.float16 else x
        mean = math_ops.reduce_mean(y, axes, keepdims=True, name='mean')
        variance = math_ops.reduce_mean(math_ops.squared_difference(y, array_ops.stop_gradient(mean)), axes, keepdims=True, name='variance')
        if not keep_dims:
            mean = array_ops.squeeze(mean, axes)
            variance = array_ops.squeeze(variance, axes)
        if x.dtype == dtypes.float16:
            return (math_ops.cast(mean, dtypes.float16), math_ops.cast(variance, dtypes.float16))
        else:
            return (mean, variance)

Calculate the mean and variance of `x`.

    The mean and variance are calculated by aggregating the contents of `x`
    across `axes`.  If `x` is 1-D and `axes = [0]` this is just the mean
    and variance of a vector.

    Note: shift is currently not used; the true mean is computed and used.

    When using these moments for batch normalization (see
    `tf.nn.batch_normalization`):

    * for so-called "global normalization", used with convolutional filters with
        shape `[batch, height, width, depth]`, pass `axes=[0, 1, 2]`.
    * for simple batch normalization pass `axes=[0]` (batch only).

Args:
        x: A `Tensor`.
        axes: Array of ints.  Axes along which to compute mean and
        variance.
        shift: Not used in the current implementation
        name: Name used to scope the operations that compute the moments.
        keep_dims: produce moments with the same dimensionality as the input.
        keepdims: Alias to keep_dims.

Returns:
        Two `Tensor` objects: `mean` and `variance`.


def _get_num_multimodal_tokens(self, image_sizes=None, **kwargs):
    vision_data = {}
    if image_sizes is not None:
        images_kwargs = LlavaNextProcessorKwargs._defaults.get('images_kwargs', {})
        images_kwargs.update(kwargs)
        size = images_kwargs.get('size', None) or self.image_processor.size
        size = (size['shortest_edge'], size['shortest_edge']) if 'shortest_edge' in size else (min(size['height'], size['width']), min(size['height'], size['width']))
        processed_height, processed_width = size
        batch_num_image_tokens = []
        num_image_patches = [1] * len(image_sizes)
        for image_size in image_sizes:
            orig_height, orig_width = image_size
            num_image_tokens = self._get_number_of_features(orig_height, orig_width, processed_height, processed_width)
            if self.vision_feature_select_strategy == 'default':
                num_image_tokens -= 1
            batch_num_image_tokens.append(num_image_tokens)
        vision_data.update({'num_image_tokens': batch_num_image_tokens, 'num_image_patches': num_image_patches})
    return MultiModalData(**vision_data)

Computes the number of placeholder tokens needed for multimodal inputs with the given sizes.

Args:
        image_sizes (List[List[str]], *optional*):
            The input sizes formatted as (height, width) per each image.
        video_sizes (List[List[str]], *optional*):
            The input sizes formatted as (num_frames, height, width) per each video.
        audio_lengths (List[int], *optional*):
            The input length formatted as per each audio.

Returns:
        Dict[str, List[int]]: A dictionary mapping each modality ("image", "video", "audio")
        to a list containing the number of placeholder tokens required. If the model doesn't accept
        a certain modality or no input sizes are provided, the dict value is set to an empty list.


def _VerifyValues(self, input_sizes=None, filter_sizes=None, out_backprop_sizes=None, strides=None, dilations=None, padding=None, data_format_src='NHWC', data_format_dst='NHWC', expected=None):
    total_size_1 = np.prod(input_sizes)
    total_size_2 = np.prod(out_backprop_sizes)
    x1 = np.arange(1, total_size_1 + 1, dtype=np.float32).reshape(input_sizes)
    x2 = np.arange(1, total_size_2 + 1, dtype=np.float32).reshape(out_backprop_sizes)
    strides = [1] + strides + [1]
    if dilations is not None:
        dilations = [1] + dilations + [1]
    expected = np.reshape(expected, filter_sizes)
    x1 = test_utils.ConvertBetweenDataFormats(x1, data_format_src, data_format_dst)
    x2 = test_utils.ConvertBetweenDataFormats(x2, data_format_src, data_format_dst)
    input_sizes = test_utils.PermuteDimsBetweenDataFormats(input_sizes, data_format_src, data_format_dst)
    out_backprop_sizes = test_utils.PermuteDimsBetweenDataFormats(out_backprop_sizes, data_format_src, data_format_dst)
    strides = test_utils.PermuteDimsBetweenDataFormats(strides, data_format_src, data_format_dst)
    if dilations is not None:
        dilations = test_utils.PermuteDimsBetweenDataFormats(dilations, data_format_src, data_format_dst)
    with self.session() as sess:
        t1 = array_ops.placeholder(dtypes.float32, shape=input_sizes)
        t2 = array_ops.placeholder(dtypes.float32, shape=out_backprop_sizes)
        with self.test_scope():
            tensor = gen_nn_ops.conv2d_backprop_filter(input=t1, filter_sizes=filter_sizes, out_backprop=t2, strides=strides, dilations=dilations, padding=padding, data_format=data_format_dst)
        value = sess.run(tensor, {t1: x1, t2: x2})
        self.assertAllEqual(filter_sizes, value.shape)
        self.assertAllClose(expected, value, 0.001)

Tests that gen_nn_ops.conv2d_backprop_filter produces the right output.

Args:
        input_sizes: Input tensor dimensions in
        [batch, input_rows, input_cols, input_depth].
        filter_sizes: Filter tensor dimensions in
        [kernel_rows, kernel_cols, input_depth, output_depth].
        out_backprop_sizes: Output gradients tensor dimensions.
        strides: Stride.
        dilations: Dilations.
        padding: Padding type.
        data_format_src: Data format input is in.
        data_format_dst: Data format verification will run and input is converted
        to.
        expected: Expected output.


def interruptible(inner: Iterable[_T] | Iterator[_T]) -> Iterator[_T]:
    interrupted = threading.Event()

    def handler(unused_signum, unused_frame):
        if interrupted.is_set():
            raise KeyboardInterrupt
        interrupted.set()
    previous_handler = signal.signal(signal.SIGINT, handler)
    try:
        for i in inner:
            yield i
            if interrupted.is_set():
                return
    finally:
        signal.signal(signal.SIGINT, previous_handler)

Catch KeyboardInterrupts to end the loop without raising an Exception.

    While this iterator is running, the first SIGINT (e.g. from interrupting
    the colab runtime, or pressing Ctrl+C ) will not raise an exception but
    instead end the loop after the current iteration.
    A second SIGINT will raise a KeyboardInterrupt Exception as usual.

Example:
    ```python
    for i in interruptible(range(10)):
        print(i)
        time.sleep(3)
    ```

Args:
        inner: arbitrary iterable or iterator

Yields:
        elements from inner


def sharded_filename(filename_tensor: tensor_lib.Tensor, shard: int, num_shards: tensor_lib.Tensor) -> tensor_lib.Tensor:
    return gen_io_ops.sharded_filename(filename_tensor, shard, num_shards)

Append sharding information to a filename.

Args:
        filename_tensor: A string tensor.
        shard: Integer.  The shard for the filename.
        num_shards: An int Tensor for the number of shards.

Returns:
        A string tensor.


def compute_one_match(self, var, other_type, name=None, match_all_views=True, keep_all_views=False, alias_map=None) -> MatchResult:
    bad_matches = []
    good_matches = _UniqueMatches(self._node, keep_all_views)
    views = abstract_utils.get_views([var], self._node)
    skip_future = None
    while True:
        try:
            view = views.send(skip_future)
        except StopIteration:
            break
        subst = datatypes.AliasingDict(aliases=alias_map)
        subst = self.match_var_against_type(var, other_type, subst, view)
        if subst is None:
            if self._node.CanHaveCombination(list(view.values())):
                bad_matches.append(BadMatch(view=view, expected=self._get_bad_type(name, other_type), actual=var))
            skip_future = False
        else:
            skip_future = True
            good_matches.insert(view, subst)
    good_matches = [GoodMatch(view, datatypes.HashableDict(subst)) for view, subst in good_matches.unique()]
    if good_matches and (not match_all_views) or not bad_matches:
        success = True
    elif good_matches:
        bad_matches = [m for m in bad_matches if self._node.HasCombination(list(m.view.values()))]
        success = not bad_matches
    else:
        success = False
    return MatchResult(success=success, good_matches=good_matches, bad_matches=bad_matches)

Match a Variable against a type.

Args:
        var: A cfg.Variable, containing instances.
        other_type: An instance of BaseValue.
        name: Optionally, the variable name.
        match_all_views: If True, every possible match must succeed for the
        overall match to be considered a success. Otherwise, the overall match
        succeeds as long as at least one possible match succeeds.
        keep_all_views: If True, avoid optimizations that discard views.
        alias_map: Optionally, a datatypes.UnionFind, which stores all the type
        renaming information, mapping of type parameter name to its
        representative.

Returns:
        The match result.


def __init__(self, use_resource_alias: bool=False, value_set_codes_table: Optional[str]=None, value_set_codes_definitions: Optional[fhir_package.FhirPackageManager]=None) -> None:
    self._use_resource_alias = use_resource_alias
    self._value_set_codes_table = value_set_codes_table
    self._value_set_codes_definitions = value_set_codes_definitions

Creates a SparkSqlInterpreter.

Args:
        use_resource_alias: Determines whether it is necessary to call the
        resource table directly through an alias.
        value_set_codes_table: The name of the database table containing value set
        code definitions. Used when building SQL for memberOf expressions. If
        given, value set definitions needed for memberOf expressions will be
        retrieved from this table if they can not be found in
        `value_set_codes_definitions`. If neither this nor
        `value_set_codes_definitions` is given, no memberOf SQL will be
        generated.
        value_set_codes_definitions: A package manager containing value set
        definitions which can be used to build SQL for memberOf expressions.
        These value set definitions can be consulted in favor of using an
        external `value_set_codes_table`. If neither this nor
        `value_set_codes_definitions` is given, no memberOf SQL will be
        generated.


def as_str_any(value, encoding='utf-8'):
    if isinstance(value, bytes):
        return as_str(value, encoding=encoding)
    else:
        return str(value)

Converts input to `str` type.

    Uses `str(value)`, except for `bytes` typed inputs, which are converted
    using `as_str`.

Args:
    value: A object that can be converted to `str`.
    encoding: Encoding for `bytes` typed inputs.

Returns:
    A `str` object.


def to(self, *args, **kwargs) -> 'BatchFeature':
    requires_backends(self, ['torch'])
    import torch
    new_data = {}
    device = kwargs.get('device')
    if device is None and len(args) > 0:
        arg = args[0]
        if is_torch_dtype(arg):
            pass
        elif isinstance(arg, str) or is_torch_device(arg) or isinstance(arg, int):
            device = arg
        else:
            raise ValueError(f'Attempting to cast a BatchFeature to type {str(arg)}. This is not supported.')

    def _to(elem):
        if torch.is_floating_point(elem):
            return elem.to(*args, **kwargs)
        if device is not None:
            return elem.to(device=device)
        return elem
    for k, v in self.items():
        if isinstance(v, list) and isinstance(v[0], list):
            new_v = []
            for elems in v:
                new_v.append([_to(elem) for elem in elems])
            new_data[k] = new_v
        elif isinstance(v, list):
            new_data[k] = [_to(elem) for elem in v]
        else:
            new_data[k] = _to(v)
    self.data = new_data
    return self

Send all values to device by calling `v.to(*args, **kwargs)` (PyTorch only). This should support casting in
    different `dtypes` and sending the `BatchFeature` to a different `device`.

Args:
        args (`Tuple`):
            Will be passed to the `to(...)` function of the tensors.
        kwargs (`Dict`, *optional*):
            Will be passed to the `to(...)` function of the tensors.

Returns:
        [`BatchFeature`]: The same instance after modification.


def _unshard_from_sc_to_cpu(stacked_table: tensor.Tensor, from_shard_layouts: Sequence[sparse_core_layout_pb2.SparseCoreTableLayout]) -> Sequence[tensor.Tensor]:
    logging.vlog(1, 'To unshuffle_from_sc_to_cpu on stacked_table.shape: %s', stacked_table[0].shape)
    ret_tensors = []
    for layout in from_shard_layouts:
        padded_table = tpu_embedding_v3_utils.unshuffle_from_sc_to_cpu(stacked_table[0], num_sparse_cores=layout.num_sparse_cores, offset_in_shard=layout.sparse_core_shard_row_offset, size_in_shard=layout.unsharded_padded_shape[0] // layout.num_sparse_cores, shard_rotation=layout.sparse_core_shard_rotation)
        orig_table = tpu_embedding_v3_utils.remove_padding_from_sc(padded_table, layout.unsharded_shape)
        logging.vlog(1, 'orig_tensors.shape[%s]: %s', layout.table_name, orig_table.shape)
        ret_tensors.append(orig_table)
    return ret_tensors

Undo the shard the feature tables into SparseCore stacked table.

Args:
        stacked_table: The value of a SparseCore stacked and sharded table.
        from_shard_layouts: The target layouts for the target hardware.

Returns:
        The unsharded feature tables.


def get_metric(name, constructor, *args, **kwargs):
    metric = _registered_metrics.get(name)
    if metric is not None:
        return metric
    else:
        return constructor(name, *args, **kwargs)

Return an existing metric or create a new one for the given name.

Args:
        name: The name of the metric.
        constructor: A class to instantiate if a new metric is required.
        *args: Additional positional args to pass to the constructor.
        **kwargs: Keyword args for the constructor.

Returns:
        The current metric registered to name, or a new one created by
        invoking constructor(name, *args, **kwargs).


def __init__(self, l1=0.0, l2=0.0, **kwargs):
    super(ActivityRegularization, self).__init__(activity_regularizer=regularizers.L1L2(l1=l1, l2=l2), **kwargs)
    self.supports_masking = True
    self.l1 = l1
    self.l2 = l2

Layer that applies an update to the cost function based input activity.

Args:
        l1: L1 regularization factor (positive float).
        l2: L2 regularization factor (positive float).

    Input shape:
        Arbitrary. Use the keyword argument `input_shape`
        (tuple of integers, does not include the samples axis)
        when using this layer as the first layer in a model.

    Output shape:
        Same shape as input.


def set_tensor(self, tensor_index, value):
    self._interpreter.SetTensor(tensor_index, value)

Sets the value of the input tensor.

    Note this copies data in `value`.

    If you want to avoid copying, you can use the `tensor()` function to get a
    numpy buffer pointing to the input buffer in the tflite interpreter.

Args:
        tensor_index: Tensor index of tensor to set. This value can be gotten from
        the 'index' field in get_input_details.
        value: Value of tensor to set.

Raises:
        ValueError: If the interpreter could not set the tensor.


def __init__(self, model_name: str, columns: list[str], dimension: Optional[int], project: Optional[str]=None, location: Optional[str]=None, credentials: Optional[Credentials]=None, **kwargs):
    self.model_name = model_name
    self.project = project
    self.location = location
    self.credentials = credentials
    if dimension is not None and dimension not in (128, 256, 512, 1408):
        raise ValueError('dimension argument must be one of 128, 256, 512, or 1408')
    self.dimension = dimension
    super().__init__(columns=columns, **kwargs)

Embedding Config for Vertex AI Image Embedding models following
    https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-multimodal-embeddings # pylint: disable=line-too-long
    Image Embeddings are generated for a batch of images using the Vertex AI API.
    Embeddings are returned in a list for each image in the batch. This
    transform makes remote calls to the Vertex AI service and may incur costs
    for use.

Args:
        model_name: The name of the Vertex AI Multi-Modal Embedding model.
        columns: The columns containing the image to be embedded.
        dimension: The length of the embedding vector to generate. Must be one of
        128, 256, 512, or 1408. If not set, Vertex AI's default value is 1408.
        project: The default GCP project for API calls.
        location: The default location for API calls.
        credentials: Custom credentials for API calls.
        Defaults to environment credentials.


def _get_details(self):
    detail_parts = [self._get_full_name(), self._error_message]
    return '\n'.join(filter(None, detail_parts))

Gets the output for the detail section of the TestResultRecord.

Returns:
        A string to set for a TestResultRecord's details.


def __enter__(self) -> str:
    if self._name is None and self._values is not None:
        raise ValueError('At least one of name (%s) and default_name (%s) must be provided.' % (self._name, self._default_name))
    g = get_default_graph()
    if self._values and (not g.building_function):
        g_from_inputs = _get_graph_from_inputs(self._values)
        if g_from_inputs is not g:
            g = g_from_inputs
            self._g_manager = g.as_default()
            self._g_manager.__enter__()
        else:
            self._g_manager = None
    else:
        self._g_manager = None
    try:
        self._name_scope = g.name_scope(self._name)
        return self._name_scope.__enter__()
    except:
        if self._g_manager is not None:
            self._g_manager.__exit__(*sys.exc_info())
        raise

Start the scope block.

Returns:
        The scope name.

Raises:
        ValueError: if neither `name` nor `default_name` is provided
        but `values` are.


def notnan(self: EventSetOrNode) -> EventSetOrNode:
    from temporian.core.operators.unary import notnan
    return notnan(self)

Returns boolean features, `False` in the NaN elements of an
    [`EventSet`][temporian.EventSet].

    Equivalent to `~evset.isnan(...)`.

    Note that for `int` and `bool` this will always be `True` since those types
    don't support NaNs. It only makes actual sense to use on `float` (or
    `tp.float32`) features.

    See also `evset.isnan()`.

Example:
        ```python
        >>> a = tp.event_set(
        ...     timestamps=[1, 2, 3],
        ...     features={"M":[np.nan, 5., np.nan], "N":  [-1, 0, 5]},
        ... )
        >>> b = a.notnan()
        >>> b
        indexes: ...
                'M': [False True False]
                'N': [ True True True]
        ...

        >>> # Filter only rows where "M" is not nan
        >>> a.filter(b["M"])
        indexes: ...
                'M': [5.]
                'N': [0]
        ...

        ```

Returns:
        EventSet with boolean features.


def _to_values_def(self, export_scope=None):
    values_def = control_flow_pb2.ValuesDef()
    values_def.values.extend([ops.strip_name_scope(v, export_scope) for v in sorted(self._values)])
    for k, v in self._external_values.items():
        k = ops.strip_name_scope(k, export_scope)
        values_def.external_values[k] = ops.strip_name_scope(v.name, export_scope)
    return values_def

Converts the values to a `ValuesDef` protocol buffer.

Args:
        export_scope: Optional `string`. Name scope to remove.

Returns:
        A `ValuesDef` protocol buffer.


def __init__(self, source: SourceBase) -> None:
    super().__init__()
    self.source = source

Initializes a Read transform.

Args:
        source: Data source to read from.


def __init__(self, input_size: int, num_experts: int, top_k: int):
    super().__init__()
    self.num_experts = num_experts
    self.input_size = input_size
    self.top_k = top_k
    self.layer = nn.Linear(input_size, num_experts, bias=False)

Initialize the top-k gating mechanism.

Args:
        input_size (`int`):
            Size of the input.
        num_experts (`int`):
            Number of experts.
        top_k (`int`):
            Number of top experts to select.


def get_config(self, obj):
    try:
        shared_object_config = self._shared_objects_config[obj]
    except (TypeError, KeyError):
        return None
    shared_object_config.increment_ref_count()
    return shared_object_config

Gets a `SharedObjectConfig` if one has already been seen for `obj`.

Args:
        obj: The object for which to retrieve the `SharedObjectConfig`.

Returns:
        The SharedObjectConfig for a given object, if already seen. Else,
        `None`.


def __init__(self, rate, seed=None, **kwargs):
    super().__init__(**kwargs)
    if not 0 <= rate <= 1:
        raise ValueError(f'Invalid value received for argument `rate`. Expected a float value between 0 and 1. Received: rate={rate}')
    self.rate = rate
    self.seed = seed
    if rate > 0:
        self.seed_generator = backend.random.SeedGenerator(seed)
    self.supports_masking = True
    self._build_at_init()

Apply multiplicative 1-centered Gaussian noise.

    As it is a regularization layer, it is only active at training time.

Args:
        rate: Float, drop probability (as with `Dropout`).
            The multiplicative noise will have
            standard deviation `sqrt(rate / (1 - rate))`.
        seed: Integer, optional random seed to enable deterministic behavior.

    Call arguments:
        inputs: Input tensor (of any rank).
        training: Python boolean indicating whether the layer should behave in
            training mode (adding dropout) or in inference mode (doing nothing).


def update_scores(self, scores: torch.FloatTensor, g_values: torch.FloatTensor) -> torch.FloatTensor:
    _, _, depth = g_values.shape
    probs = torch.softmax(scores, dim=1)
    for i in range(depth):
        g_values_at_depth = g_values[:, :, i]
        g_mass_at_depth = (g_values_at_depth * probs).sum(axis=1, keepdims=True)
        probs = probs * (1 + g_values_at_depth - g_mass_at_depth)
    log_probs = torch.log(probs)
    log_probs = torch.where(torch.isfinite(log_probs), log_probs, torch.finfo(log_probs.dtype).min)
    return log_probs

Updates scores using the g values.

    We assume that the scores are in the log space.

Args:
        scores (`torch.FloatTensor`): Scores (batch_size, vocab_size).
        g_values (`torch.FloatTensor`): G values (batch_size, vocab_size, depth).

Returns:
        Updated scores (batch_size, vocab_size).


def symlink_force(target, link_name):
    try:
        os.symlink(target, link_name)
    except OSError as e:
        if e.errno == errno.EEXIST:
            os.remove(link_name)
            os.symlink(target, link_name)
        else:
            raise e

Force symlink, equivalent of 'ln -sf'.

Args:
        target: items to link to.
        link_name: name of the link.


def __call__(self):
    output_shapes = self._output_shapes()
    cond_is_stacked = [None]
    indices_to_stack = []

    def cond(not_all_done, *_):
        return not_all_done

    def body(not_all_done, indices, *args):
        num_inputs = self._pfor_input.num_inputs
        inputs = args[:num_inputs]
        output_tas = args[num_inputs:]
        inputs_stacked = [x.is_stacked for x in self._pfor_input.inputs]
        assert len(inputs) >= len(output_tas)
        assert len(inputs) == len(inputs_stacked)
        with ops.name_scope('while_cond'):
            cond_pfor = PFor(loop_var=self._pfor.loop_var, loop_len=array_ops.size(indices), pfor_ops=self._cond_func.graph.get_operations(), fallback_to_while_loop=self._pfor.fallback_to_while_loop, all_indices=indices, all_indices_partitioned=True, pfor_config=self._pfor.pfor_config)
            wrapped_inputs = [wrap(inp, stacked) for inp, stacked in zip(inputs, inputs_stacked)]
            conditions, cond_stacked, _ = _convert_function_call(self._cond_func, cond_pfor, wrapped_inputs)[0]
            cond_is_stacked[0] = cond_stacked
        if not cond_stacked:
            not_all_done, new_indices, new_inputs, new_output_tas = self._process_cond_unstacked(conditions, indices, inputs, output_tas)
        else:
            not_all_done, new_indices, new_inputs, new_output_tas = self._process_cond_stacked(conditions, indices, inputs, inputs_stacked, output_tas)
        with ops.name_scope('while_body'):
            new_outputs, mismatching_stacked_indices = self._process_body(inputs_stacked, new_indices, cond_stacked, new_inputs, not_all_done)
        indices_to_stack[:] = mismatching_stacked_indices
        for i, new_output in enumerate(new_outputs):
            new_output.set_shape(output_shapes[i])
        new_args = [not_all_done, new_indices] + new_outputs + list(new_output_tas)
        return tuple(new_args)

    @def_function.function
    def while_fn():
        init_values = self._init_values()
        ta_shape_invariants = [tensor_shape.TensorShape([]) for _ in self._pfor_input.outputs]
        shape_invariants = [tensor_shape.TensorShape([]), tensor_shape.TensorShape([None])] + output_shapes + ta_shape_invariants
        while_outputs = while_loop.while_loop(cond, body, init_values, shape_invariants=shape_invariants, parallel_iterations=self._parallel_iterations)
        if indices_to_stack:
            return while_outputs
        else:
            num_inputs = self._pfor_input.num_inputs
            new_inputs = while_outputs[2:num_inputs + 2]
            output_tas = while_outputs[num_inputs + 2:]
            assert cond_is_stacked[0] is not None
            outputs = []
            for i, inp in enumerate(new_inputs):
                if cond_is_stacked[0]:
                    if i in self._body_pass_through_indices:
                        outputs.append(init_values[i + 2])
                    else:
                        ta = output_tas[i]
                        if _variant_type_id(inp) == full_type_pb2.TFT_ARRAY:
                            shape_and_type = _parse_variant_shapes_and_types(inp)[0]
                            length = list_ops.tensor_list_length(inp)

                            def _stack_loop_body(index, output_list):
                                current_value = ta.read(index)
                                output_list = list_ops.tensor_list_set_item(output_list, index, list_ops.tensor_list_stack(current_value, shape_and_type.dtype))
                                return (index + 1, output_list)
                            output_list = list_ops.tensor_list_reserve(tensor_shape.TensorShape(shape_and_type.shape), length, shape_and_type.dtype)
                            _, output_list = while_loop.while_loop(lambda index, _: index < length, _stack_loop_body, [0, output_list])
                            outputs.append(output_list)
                        else:
                            outputs.append(ta.stack())
                else:
                    outputs.append(inp)
            return outputs
    _ = while_fn.get_concrete_function()
    if indices_to_stack:
        self._pfor_input.stack_inputs(stack_indices=indices_to_stack, tile_variants=True)
        return self()
    else:
        outputs = while_fn()
        wrapped_outputs = []
        for i, (out, inp) in enumerate(zip(outputs, self._pfor_input.inputs)):
            if i not in self._body_pass_through_indices and cond_is_stacked[0]:
                wrapped_outputs.append(wrap(out, True))
            else:
                wrapped_outputs.append(wrap(out, inp.is_stacked))
        return wrapped_outputs

Converter for the V2 while_loop.

    The conversion of a while_loop is another while_loop.

    The arguments to this converted while_loop are as follows:
    not_all_done: Boolean scalar Tensor indicating if all the pfor iterations
        are done.
    indices: int32 1-D Tensor storing the id of the pfor iterations that are not
        done.
    args: Remaining arguments. These can be divided into 2 categories:
        - The first set of arguments correspond one-to-one to the inputs to the
        unvectorized while_loop.
        - The second set are TensorArrays, corresponding one-to-one to each output
        of the unvectorized while_loop. Each TensorArray has `PFor.loop_len`
        elements, i.e. the number of pfor iterations. At the end, the i'th
        element of each TensorArray will contain the output computed by the i'th
        iteration of pfor. Note that elements can be written into these tensors
        arrays in any order, depending on when the corresponding pfor iteration
        is done.
    In each iteration, the while_loop body recomputes the condition for all
    active pfor iterations to see which of them are now done. It then partitions
    all the inputs and passes them along to the converted body. Values for all
    the iterations that are done are written to TensorArrays indexed by the pfor
    iteration number. When all iterations are done, the TensorArrays are stacked
    to get the final value.

Returns:
        List of converted outputs.


def reduced_shape(input_shape, axes):
    constant_input_shape = tensor_util.constant_value(input_shape)
    if constant_input_shape is not None:
        constant_axes = tensor_util.constant_value(axes)
        if constant_axes is not None:
            constant_axes = np.array(constant_axes, dtype=np.int32)
            constant_input_shape = np.array(constant_input_shape, dtype=np.int32)
            constant_input_shape[constant_axes] = 1
            return constant_input_shape
    axes = ops.convert_to_tensor(axes)
    input_rank = array_ops.size(input_shape, out_type=axes.dtype)
    axes = (axes + input_rank) % input_rank
    axes_shape = array_ops.shape(axes)
    return gen_data_flow_ops.dynamic_stitch([range(input_rank), axes], [input_shape, array_ops.ones(axes_shape, dtype=input_shape.dtype)])

Helper function for reduction ops.

Args:
        input_shape: 1-D Tensor, the shape of the Tensor being reduced.
        axes: 1-D Tensor, the reduction axes.

Returns:
        A 1-D Tensor, the output shape as if keepdims were set to True.


def _objective_and_vega(volatilities):
    vols = volatilities * sqrt_t / normalization
    d1 = (norm_forwards - norm_strikes) / vols
    implied_prices = (norm_forwards - norm_strikes) * _cdf(d1) + vols * _pdf(d1)
    if is_call_options is not None:
        put_prices = implied_prices - norm_forwards + norm_strikes
        implied_prices = tf.where(is_call_options, implied_prices, put_prices)
    vega = _pdf(d1) * sqrt_t / discount_factors / normalization
    return (implied_prices - normalized_prices, vega)

Calculate the Bachelier price and vega for a given volatility.

    This method returns normalized results.

Args:
        volatilities: A real `Tensor` of same shape and dtype as `forwards`. The
        volatility to expiry.

Returns:
        A tuple containing (value, gradient) of the black scholes price, both of
        which are `Tensor`s of the same shape and dtype as `volatilities`.


def range_input_producer(limit, num_epochs=None, shuffle=True, seed=None, capacity=32, shared_name=None, name=None):
    with ops.name_scope(name, 'input_producer', [limit]) as name:
        range_tensor = math_ops.range(limit)
        return input_producer(range_tensor, [], num_epochs, shuffle, seed, capacity, shared_name, 'fraction_of_%d_full' % capacity, name)

Produces the integers from 0 to limit-1 in a queue.

    Note: if `num_epochs` is not `None`, this function creates local counter
    `epochs`. Use `local_variables_initializer()` to initialize local variables.

Args:
        limit: An int32 scalar tensor.
        num_epochs: An integer (optional). If specified, `range_input_producer`
        produces each integer `num_epochs` times before generating an
        OutOfRange error. If not specified, `range_input_producer` can cycle
        through the integers an unlimited number of times.
        shuffle: Boolean. If true, the integers are randomly shuffled within each
        epoch.
        seed: An integer (optional). Seed used if shuffle == True.
        capacity: An integer. Sets the queue capacity.
        shared_name: (optional). If set, this queue will be shared under the given
        name across multiple sessions.
        name: A name for the operations (optional).

Returns:
        A Queue with the output integers.  A `QueueRunner` for the Queue
        is added to the current `Graph`'s `QUEUE_RUNNER` collection.

    @compatibility(eager)
    Input pipelines based on Queues are not supported when eager execution is
    enabled. Please use the `tf.data` API to ingest data under eager execution.
    @end_compatibility


def __init__(self, lr=0.002, beta_1=0.9, beta_2=0.999, epsilon=None, decay=0.0, **kwargs):
    super(Adamax, self).__init__(**kwargs)
    with backend.name_scope(self.__class__.__name__):
        self.iterations = backend.variable(0, dtype='int64', name='iterations')
        self.lr = backend.variable(lr, name='lr')
        self.beta_1 = backend.variable(beta_1, name='beta_1')
        self.beta_2 = backend.variable(beta_2, name='beta_2')
        self.decay = backend.variable(decay, name='decay')
    if epsilon is None:
        epsilon = backend.epsilon()
    self.epsilon = epsilon
    self.initial_decay = decay

Adamax optimizer from Adam paper's Section 7.

    It is a variant of Adam based on the infinity norm.
    Default parameters follow those provided in the paper.

Args:
        lr: float >= 0. Learning rate.
        beta_1/beta_2: floats, 0 < beta < 1. Generally close to 1.
        epsilon: float >= 0. Fuzz factor.
        If `None`, defaults to `backend.epsilon()`.
        decay: float >= 0. Learning rate decay over each update.


def aggregate_and_return_name_for_input(self, out_graphdef):
    del out_graphdef
    raise RuntimeError('Unimplemented abstract method.')

This adds the node(s) to out_graphdef and returns the input node name.

Args:
        out_graphdef: A graphdef that is ready to have this input added.

Returns:
        The output that the stub should use as an input for this operand.

Raises:
        RuntimeError: if the method is not implemented.


def execute_command(self, command: str, read: bool=False) -> None:
    if self.debug:
        print(command)
        return 'SIMULATING VALUE' if read else True
    else:
        print('\nCOMMAND\n', command, '\n' + '-' * 40)
        try:
            cmd = subprocess.run(command, shell=True, capture_output=read, text=True, check=True)
            if read:
                return cmd.stdout.strip()
            return True
        except subprocess.CalledProcessError as e:
            return False

Helper function that either executes or prints each command.

Args:
        command - a command line command, typically a gcloud command.
        read - if True, the commands output is passed back to the caller.

Returns:
        Bool - if command value is not required, simply indicate successor failure.
        String - if read is specified, the command output or error is returned.


def apply(self, score: Optional[float]) -> Optional[int]:
    raise NotImplementedError

Applies the threshold function to a given score to classify it as
    normal or outlier.

Args:
        score: The outlier score generated from the detector (model).

Returns:
        The label assigned to the score, either `self._normal_label`
        or `self._outlier_label`


def preprocess(text):
    if len(text.strip()) > 0:
        question, context = text.split(';')
        yield (question, context)

preprocess separates the text into question and context
    by splitting on semi-colon.

Args:
        text (str): string with question and context separated by semi-colon.

Yields:
        (str, str): yields question and context from text.


def patchify_image(self, image: 'torch.Tensor', patch_size: Optional[Dict[str, int]]=None) -> 'torch.Tensor':
    requires_backends(self, ['torch'])
    patch_size = patch_size if patch_size is not None else self.patch_size
    patch_height, patch_width = (patch_size['height'], patch_size['width'])
    batch_size, channels, _, _ = image.shape
    unfolded_along_height = image.unfold(2, patch_height, patch_height)
    patches = unfolded_along_height.unfold(3, patch_width, patch_width)
    patches = patches.contiguous()
    patches = patches.view(batch_size, channels, -1, patch_height, patch_width)
    patches = patches.permute(0, 2, 3, 4, 1)
    patches = patches.reshape(batch_size, -1, channels * patch_height * patch_width)
    return patches

Convert an image into a tensor of patches.

Args:
        image (`torch.Tensor`):
            Image to convert. Shape: [batch, channels, height, width]
        patch_size (`Dict[str, int]`, *optional*, defaults to `self.patch_size`):
            Dictionary in the format `{"height": int, "width": int}` specifying the size of the patches.


def get_text_features(self, input_ids: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None, position_ids: Optional[torch.Tensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> torch.FloatTensor:
    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
    text_outputs = self.text_model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    pooled_output = text_outputs[1] if return_dict is not None else text_outputs.pooler_output
    text_features = self.text_projection(pooled_output)
    text_features = F.normalize(text_features, dim=-1)
    return text_features

Returns:
    text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by
    applying the projection layer to the pooled output of [`ClapTextModel`].

Example:
    ```python
    >>> from transformers import AutoTokenizer, ClapModel

    >>> model = ClapModel.from_pretrained("laion/clap-htsat-unfused")
    >>> tokenizer = AutoTokenizer.from_pretrained("laion/clap-htsat-unfused")

    >>> inputs = tokenizer(["the sound of a cat", "the sound of a dog"], padding=True, return_tensors="pt")
    >>> text_features = model.get_text_features(**inputs)
    ```


def log_every_n(level, msg, n, *args):
    count = _GetNextLogCountPerToken(_GetFileAndLine())
    log_if(level, msg, not count % n, *args)

Log 'msg % args' at level 'level' once per 'n' times.

    Logs the 1st call, (N+1)st call, (2N+1)st call,  etc.
    Not threadsafe.

Args:
        level: The level at which to log.
        msg: The message to be logged.
        n: The number of times this should be called before it is logged.
        *args: The args to be substituted into the msg.


def __init__(self, name, func, create_scope_now=False, unique_name=None, custom_getter=None, create_graph_function=False):
    if create_graph_function:
        self._func = def_function.function(func)
    else:
        self._func = func
    self._stacktrace = traceback.format_stack()[:-2]
    self._name = name
    self._unique_name = unique_name
    self._custom_getter = custom_getter
    if name is None:
        raise ValueError('name cannot be None.')
    if create_scope_now:
        with variable_scope._pure_variable_scope(self._unique_name or variable_scope._get_unique_variable_scope(self._name), custom_getter=self._custom_getter) as vs:
            self._variable_scope = vs
    else:
        self._variable_scope = None
    self._variables_created = False
    self._first_call = True

Creates a template for the given function.

Args:
        name: A name for the scope created by this template. The name will be made
        unique by appending `_N` to the it (see how
        `tf.compat.v1.variable_scope` treats the `default_name` for details).
        func: The function to apply each time.
        create_scope_now: Whether to create the scope at Template construction
        time, rather than first call. Defaults to false. Creating the scope at
        construction time may be more convenient if the template is to passed
        through much lower level code, and you want to be sure of the scope name
        without knowing exactly where it will be first called. If set to True,
        the scope will be created in the constructor, and all subsequent times
        in `__call__`, leading to a trailing numeral being added to the names of
        all created Tensors. If set to False, the scope will be created at the
        first call location.
        unique_name: When used, it overrides `name` and is not made unique. If a
        template of the same scope/unique_name already exists and reuse is
        false, an error is raised. Defaults to None.
        custom_getter: optional custom getter to pass to `variable_scope()`
        create_graph_function: When True, `func` will be executed as a graph
        function. Enabling this flag gives the caller access to graph-function
        semantics, i.e., accesses to variables are totally ordered and
        side-effecting ops are not pruned.

Raises:
        ValueError: if `name` is None.


def recipe_to_python(name, description, instructions, tasks, parameters={}, project=None, client_credentials=None, user_credentials=None, service_credentials=None):
    tasks = json_expand_queries(tasks)
    code = DISCLAIMER
    code += 'import argparse\n'
    code += 'import textwrap\n\n'
    code += 'from starthinker.util.configuration import Configuration\n'
    imported = set()
    for task in tasks:
        script, task = next(iter(task.items()))
        if script not in imported:
            code += 'from starthinker.task.%s.run import %s\n' % (script, script)
            imported.add(script)
    code += '\n'
    code += '\n'
    fields = json_get_fields(tasks)
    if fields:
        code += 'def recipe_%s(config, %s):\n' % (name, ', '.join([f['name'] for f in fields]))
    else:
        code += 'def recipe_%s(config):\n' % name
    if description or fields:
        code += '  """' + textwrap.fill(description, width=80, subsequent_indent='     ') + '\n'
        if fields:
            code += '\n     Args:\n'
            for field in fields:
                code += '       %s (%s) - %s\n' % (field['name'], field['kind'], field.get('description', 'NA'))
        code += '  """\n\n'
    for task in tasks:
        script, task = next(iter(task.items()))
        code += '  %s(config, %s)\n\n' % (script, dict_to_python(task, indent=1))
    code += '\n'
    code += '\n'
    code += 'if __name__ == "__main__":\n'
    code += parameters_to_argparse(description, instructions, fields)
    code += '\n'
    code += '  args = parser.parse_args()\n'
    code += '\n'
    code += '  config = Configuration(\n    project=args.project,\n    user=args.user,\n    service=args.service,\n    client=args.client,\n    key=args.key,\n    verbose=args.verbose\n  )'
    code += '\n\n'
    if fields:
        code += '  recipe_%s(config, %s)\n' % (name, ', '.join(['args.%s' % f['name'] for f in fields]))
    else:
        code += '  recipe_%s(config)\n' % name
    return code

Converts a JSON recipe into a python stand alone example.

    Sets up multiple steps to execute recipe:
        1. Install starthinker from repository
        2. Get Cloud Project ID.
        3. Get Client Credentials ( optional if User Credentials exist ).
        4. Enter Recipe parameters if fields present.
        5. Execute recipe tasks.

Args:
        * name: (string) The name of the notebook.
        * description: (string) A description fo the recipe.
        * instructions: (string) Recipe manual instructions, for example connecting datastudios.
        * tasks: (list) The task JSON to execute.
        * parameters: (dict) Values for field parameters in tasks, optional.
        * project: (string) The GCP project id.
        * client_credentials: (string) The GCP Desktop Client Credentials in JSON string.
        * user_credentials: (string) Not used, placeholder.
        * service_credentials: (string) Not used, placeholder.

Returns:
        * (string) Rendered example source code to be written to a py file.


def _InsertNodeAt(new_node, target, after=False):
    if new_node.parent is not None:
        raise RuntimeError('inserting node which already has a parent', (new_node, new_node.parent))
    parent_of_target = target.parent
    if parent_of_target is None:
        raise RuntimeError('expected target node to have a parent', (target,))
    for i, child in enumerate(parent_of_target.children):
        if child is target:
            insertion_index = i + 1 if after else i
            parent_of_target.insert_child(insertion_index, new_node)
            return
    raise RuntimeError('unable to find insertion point for target node', (target,))

Underlying implementation for node insertion.

Args:
        new_node: a new node to insert (this node should not be in the tree).
        target: the target node.
        after: if True, new_node is inserted after target. Otherwise, it's inserted
        before target.

Returns:
        nothing

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


def make_parser():

    def add_kythe_field(parser, field):
        parser.add_argument('--' + field, dest=field, type=str, action='store', default='', help="Part of kythe's file-level vname proto.")
    parser = argparse.ArgumentParser(usage='%(prog)s [options] input')
    add_kythe_field(parser, 'kythe_corpus')
    add_kythe_field(parser, 'kythe_root')
    add_kythe_field(parser, 'kythe_path')
    parser.add_argument('--show-types', action='store_true', dest='show_types', default=None, help='Display inferred types.')
    parser.add_argument('--show-kythe', action='store_true', dest='show_kythe', default=None, help='Display kythe facts.')
    parser.add_argument('--show-spans', action='store_true', dest='show_spans', default=None, help='Display kythe spans.')
    parser.add_argument('--skip-stdlib', action='store_true', dest='skip_stdlib', default=None, help='Display inferred types.')
    wrapper = datatypes.ParserWrapper(parser)
    pytype_config.add_basic_options(wrapper)
    with wrapper.add_only(['--imports_info', '--debug']):
        pytype_config.add_infrastructure_options(wrapper)
        pytype_config.add_debug_options(wrapper)
    wrapper.add_argument('input', metavar='input', nargs=1, help='A .py file to index')
    return XrefParser(parser, pytype_single_args=wrapper.actions)

Make parser for command line args.

Returns:
        A Parser object.


def to_dict(self) -> dict[str, Any]:
    output = copy.deepcopy(self.__dict__)
    sig = inspect.signature(self.__init__)
    attrs_to_save = sig.parameters
    attrs_to_save = [x for x in attrs_to_save if x not in self.__class__.attributes]
    attrs_to_save += ['auto_map']
    output = {k: v for k, v in output.items() if k in attrs_to_save}
    output['processor_class'] = self.__class__.__name__
    if 'tokenizer' in output:
        del output['tokenizer']
    if 'image_processor' in output:
        del output['image_processor']
    if 'video_processor' in output:
        del output['video_processor']
    if 'feature_extractor' in output:
        del output['feature_extractor']
    if 'chat_template' in output:
        del output['chat_template']
    output = {k: v for k, v in output.items() if not (isinstance(v, PushToHubMixin) or v.__class__.__name__ == 'BeamSearchDecoderCTC')}
    return output

Serializes this instance to a Python dictionary.

Returns:
        `Dict[str, Any]`: Dictionary of all the attributes that make up this processor instance.


def _get_saver_or_default():
    collection_key = ops.GraphKeys.SAVERS
    savers = ops.get_collection(collection_key)
    if savers:
        if len(savers) > 1:
            raise RuntimeError('More than one item in collection {}. Please indicate which one to use by passing it to the constructor.'.format(collection_key))
        return savers[0]
    saver = Saver(sharded=True, allow_empty=True)
    if saver is not None:
        ops.add_to_collection(collection_key, saver)
    return saver

Returns the saver from SAVERS collection, or creates a default one.

    This method is used by other members of the training module, such as
    `Scaffold`, or `CheckpointSaverHook`.

Returns:
        `Saver`.

Raises:
        RuntimeError: If the SAVERS collection already has more than one items.


def parse_module_content(content: str) -> List[str]:
    objects = []
    current_object = []
    lines = content.split('\n')
    end_markers = [')', ']', '}', '"""']
    for line in lines:
        is_valid_object = len(current_object) > 0
        if is_valid_object and len(current_object) == 1:
            is_valid_object = not current_object[0].startswith('# Copied from')
        if not is_empty_line(line) and find_indent(line) == 0 and is_valid_object:
            if line in end_markers:
                current_object.append(line)
                objects.append('\n'.join(current_object))
                current_object = []
            else:
                objects.append('\n'.join(current_object))
                current_object = [line]
        else:
            current_object.append(line)
    if len(current_object) > 0:
        objects.append('\n'.join(current_object))
    return objects

Parse the content of a module in the list of objects it defines.

Args:
        content (`str`): The content to parse

Returns:
        `List[str]`: The list of objects defined in the module.


def prod(x, axis=None, keepdims=False, dtype=None):
    if any_symbolic_tensors((x,)):
        return Prod(axis=axis, keepdims=keepdims, dtype=dtype).symbolic_call(x)
    return backend.numpy.prod(x, axis=axis, keepdims=keepdims, dtype=dtype)

Return the product of tensor elements over a given axis.

Args:
        x: Input tensor.
        axis: Axis or axes along which a product is performed. The default,
            `axis=None`, will compute the product of all elements
            in the input tensor.
        keepdims: If this is set to `True`, the axes which are reduce
            are left in the result as dimensions with size one.
        dtype: Data type of the returned tensor.

Returns:
        Product of elements of `x` over the given axis or axes.


def __init__(self, cluster_resolver, checkpoint_or_checkpoint_manager, checkpoint_dir=None, termination_config=None):
    if isinstance(checkpoint_or_checkpoint_manager, checkpoint_lib.Checkpoint) and (not checkpoint_dir):
        raise errors.InvalidArgumentError('When a checkpoint is passed, a checkpoint_dir must be passed as well.')
    self._cluster_resolver = cluster_resolver
    self._termination_config = termination_config
    self._checkpoint_or_checkpoint_manager = checkpoint_or_checkpoint_manager
    self._checkpoint_dir = checkpoint_dir
    self._platform_device = failure_handling_util.detect_platform()
    completed_termination_config = _complete_config_for_environment(self._platform_device, self._termination_config)
    self._termination_watcher_fn = completed_termination_config.termination_watcher_fn
    self._exit_fn = completed_termination_config.exit_fn
    self._grace_period = completed_termination_config.grace_period
    self._save_fn = completed_termination_config.save_fn
    self._local_mode = True
    if self._platform_device in (failure_handling_util.PlatformDevice.GCE_TPU, failure_handling_util.PlatformDevice.GCE_CPU):
        logging.warning('PreemptionCheckpointHandler does not support usage with TPU or CPU device on GCP.')
    elif self._platform_device == failure_handling_util.PlatformDevice.INTERNAL_TPU:
        self._initialize_for_tpu_strategy()
    else:
        if cluster_resolver and 'ps' in cluster_resolver.cluster_spec().as_dict():
            raise NotImplementedError('PreemptionCheckpointHandler does not supportusage with tf.distribute.experimental.ParameterServerStrategy.')
        self._initialize_for_mirrored_and_multi_worker_mirrored()
    logging.info('PreemptionCheckpointHandler initialized or restored.')

Creates the `PreemptionCheckpointHandler`.

Args:
        cluster_resolver: a `tf.distribute.cluster_resolver.ClusterResolver`
        object. You may also obtain it through the `cluster_resolver` attribute
        of the distribution strategy in use.
        checkpoint_or_checkpoint_manager: a `tf.train.CheckpointManager` or a
        `tf.train.Checkpoint`. If you are using a `tf.train.CheckpointManager`
        to manage checkpoints outside the `PreemptionCheckpointHandler` for
        backup purpose as well, pass it as `checkpoint_or_checkpoint_manager`
        argument. Otherwise, pass a `tf.train.Checkpoint` and the
        `PreemptionCheckpointHandler` will create
        a `tf.train.CheckpointManager` to manage it in the `checkpoint_dir`.
        checkpoint_dir: a directory where the `PreemptionCheckpointHandler` saves
        and restores checkpoints. When a `PreemptionCheckpointHandler` is
        created, the latest checkpoint in the `checkpoint_dir` will be restored.
        (This is not needed if a `tf.train.CheckpointManager` instead of a
        `tf.train.Checkpoint` is passed as the
        `checkpoint_or_checkpoint_manager` argument.)
        termination_config: optional, a
        `tf.distribute.experimental.TerminationConfig` object to configure for a
        platform other than Google Borg or GCP.


def _preload_simple_restoration(self, name):
    deferred_dependencies_list = self._deferred_dependencies.get(name, ())
    if not deferred_dependencies_list:
        return
    for checkpoint_position in deferred_dependencies_list:
        if not checkpoint_position.is_simple_variable():
            return None
    checkpoint_position = max(deferred_dependencies_list, key=lambda restore: restore.checkpoint.restore_uid)
    return CheckpointInitialValueCallable(checkpoint_position=checkpoint_position)

Return a dependency's value for restore-on-create.

    Note the restoration is not deleted; if for some reason preload is called
    and then not assigned to the variable (for example because a custom getter
    overrides the initializer), the assignment will still happen once the
    variable is tracked (determined based on checkpoint.restore_uid).

Args:
        name: The object-local name of the dependency holding the variable's
        value.

Returns:
        An callable for use as a variable's initializer/initial_value, or None if
        one should not be set (either because there was no variable with this name
        in the checkpoint or because it needs more complex deserialization). Any
        non-trivial deserialization will happen when the variable object is
        tracked.


def __init__(self, weight_shape: Sequence[int]) -> None:
    self.filters = np.random.uniform(low=-1.0, high=1.0, size=weight_shape)
    if bias_fn is not None:
        self.bias = np.random.uniform(low=-1.0, high=1.0, size=weight_shape[-1])

Initializes a MatmulModel.

Args:
        weight_shape: Shape of the weight tensor.


def align_long_axis(self, image: np.ndarray, size: Dict[str, int], data_format: Optional[Union[str, ChannelDimension]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None) -> np.ndarray:
    input_height, input_width = get_image_size(image, channel_dim=input_data_format)
    output_height, output_width = (size['height'], size['width'])
    if input_data_format is None:
        input_data_format = infer_channel_dimension_format(image)
    if input_data_format == ChannelDimension.LAST:
        rot_axes = (0, 1)
    elif input_data_format == ChannelDimension.FIRST:
        rot_axes = (1, 2)
    else:
        raise ValueError(f'Unsupported data format: {input_data_format}')
    if output_width < output_height and input_width > input_height or (output_width > output_height and input_width < input_height):
        image = np.rot90(image, 3, axes=rot_axes)
    if data_format is not None:
        image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
    return image

Align the long axis of the image to the longest axis of the specified size.

Args:
        image (`np.ndarray`):
            The image to be aligned.
        size (`Dict[str, int]`):
            The size `{"height": h, "width": w}` to align the long axis to.
        data_format (`str` or `ChannelDimension`, *optional*):
            The data format of the output image. If unset, the same format as the input image is used.
        input_data_format (`ChannelDimension` or `str`, *optional*):
            The channel dimension format of the input image. If not provided, it will be inferred.

Returns:
        `np.ndarray`: The aligned image.


def process_one_file(options):
    log.info('Process %s => %s', options.input, options.output)
    try:
        ret = check_or_generate_pyi(options)
    except utils.UsageError:
        logging.exception('')
        return 1
    if not options.check:
        if options.pickle_output:
            pyi_output = options.verify_pickle
        else:
            pyi_output = options.output
        if pyi_output:
            _write_pyi_output(options, ret.pyi, pyi_output)
        if options.pickle_output:
            log.info('write pickle %r => %r', options.input, options.output)
            write_pickle(ret.ast, options, ret.context.loader)
    if options.unused_imports_info_files:
        if options.use_rewrite:
            pass
        else:
            cwd = os.getcwd()
            unused_paths = sorted(ret.context.loader.get_unused_imports_map_paths())
            with options.open_function(options.unused_imports_info_files, 'wt', encoding='utf-8') as f:
                for unused_path in unused_paths:
                    f.write(f'{os.path.relpath(unused_path, cwd)}\n')
    exit_status = handle_errors(ret.context.errorlog, options)
    ret.context.program = None
    if options.touch and (not exit_status):
        with options.open_function(options.touch, 'a'):
            os.utime(options.touch, None)
    return exit_status

Check a .py file or generate a .pyi for it, according to options.

Args:
        options: config.Options object.

Returns:
        An error code (0 means no error).


def create_model() -> tf.keras.Model:
    return tf.keras.Sequential(pg.oneof([lambda: [tf.keras.layers.Flatten(), tf.keras.layers.Dense(pg.oneof([64, 128]), pg.oneof(['relu', 'sigmoid']))], lambda: [tf.keras.layers.Lambda(lambda x: tf.reshape(x, (-1, 28, 28, 1))), tf.keras.layers.Conv2D(pg.oneof([64, 128]), pg.oneof([(3, 3), (5, 5)]), activation=pg.oneof(['relu', 'sigmoid'])), tf.keras.layers.Flatten()]]) + [tf.keras.layers.Dense(10, activation='softmax')])

Create model for training.

    Create a simple tf.keras model for training.

Returns:
        The model to use for training.


def normalize(self, image, mean, std, rescale=False):
    self._ensure_format_supported(image)
    if isinstance(image, PIL.Image.Image):
        image = self.to_numpy_array(image, rescale=True)
    elif rescale:
        if isinstance(image, np.ndarray):
            image = self.rescale(image.astype(np.float32), 1 / 255.0)
        elif is_torch_tensor(image):
            image = self.rescale(image.float(), 1 / 255.0)
    if isinstance(image, np.ndarray):
        if not isinstance(mean, np.ndarray):
            mean = np.array(mean).astype(image.dtype)
        if not isinstance(std, np.ndarray):
            std = np.array(std).astype(image.dtype)
    elif is_torch_tensor(image):
        import torch
        if not isinstance(mean, torch.Tensor):
            if isinstance(mean, np.ndarray):
                mean = torch.from_numpy(mean)
            else:
                mean = torch.tensor(mean)
        if not isinstance(std, torch.Tensor):
            if isinstance(std, np.ndarray):
                std = torch.from_numpy(std)
            else:
                std = torch.tensor(std)
    if image.ndim == 3 and image.shape[0] in [1, 3]:
        return (image - mean[:, None, None]) / std[:, None, None]
    else:
        return (image - mean) / std

Normalizes `image` with `mean` and `std`. Note that this will trigger a conversion of `image` to a NumPy array
    if it's a PIL Image.

Args:
        image (`PIL.Image.Image` or `np.ndarray` or `torch.Tensor`):
            The image to normalize.
        mean (`List[float]` or `np.ndarray` or `torch.Tensor`):
            The mean (per channel) to use for normalization.
        std (`List[float]` or `np.ndarray` or `torch.Tensor`):
            The standard deviation (per channel) to use for normalization.
        rescale (`bool`, *optional*, defaults to `False`):
            Whether or not to rescale the image to be between 0 and 1. If a PIL image is provided, scaling will
            happen automatically.


def write(self, dataset):
    if not isinstance(dataset, data_types.DatasetV2):
        raise TypeError(f'Invalid `dataset.` Expected a `tf.data.Dataset` object but got {type(dataset)}.')
    if not dataset_ops.get_structure(dataset).is_compatible_with(tensor_spec.TensorSpec([], dtypes.string)):
        raise TypeError(f'Invalid `dataset`. Expected a`dataset` that produces scalar `tf.string` elements, but got a dataset which produces elements with shapes {dataset_ops.get_legacy_output_shapes(dataset)} and types {dataset_ops.get_legacy_output_types(dataset)}.')
    dataset = dataset._apply_debug_options()
    return gen_experimental_dataset_ops.dataset_to_tf_record(dataset._variant_tensor, self._filename, self._compression_type)

Writes a dataset to a TFRecord file.

    An operation that writes the content of the specified dataset to the file
    specified in the constructor.

    If the file exists, it will be overwritten.

Args:
        dataset: a `tf.data.Dataset` whose elements are to be written to a file

Returns:
        In graph mode, this returns an operation which when executed performs the
        write. In eager mode, the write is performed by the method itself and
        there is no return value.

    Raises
        TypeError: if `dataset` is not a `tf.data.Dataset`.
        TypeError: if the elements produced by the dataset are not scalar strings.


def __init__(self, num_participants):
    self._num_participants = num_participants
    self._counter = 0
    self._flag = False
    self._local_sense = threading.local()
    self._lock = threading.Lock()
    self._condition = threading.Condition()

Initializes the barrier object.

Args:
        num_participants: an integer which is the expected number of calls of
        `wait` pass to through this barrier.


def trace(name, *trace_args):

    def decorator(f):

        def wrapper(*args, **kwargs):
            t = tracer(name)
            if t.getEffectiveLevel() < logging.DEBUG:
                return f(*args, **kwargs)
            argspec = inspect.getfullargspec(f)
            t.debug('%s: {', f.__name__)
            for arg in trace_args:
                if isinstance(arg, int):
                    argname = argspec.args[arg]
                    val = args[arg]
                else:
                    argname = arg
                    val = kwargs[arg]
                t.debug('%s: %s = %s', f.__name__, argname, show(val))
            ret = f(*args, **kwargs)
            t.debug('%s: -> %s', f.__name__, show(ret))
            t.debug('%s: }', f.__name__)
            return ret
        return wrapper
    return decorator

Record args and return value for a function call.

    The trace is of the form
        function name: {
        function name: arg = value
        function name: arg = value
        ...
        function name: -> return
        function name: }

    This will let us write tools to pretty print the traces with indentation etc.

Args:
        name: module name, usually `__name__`
        *trace_args: function arguments to log

Returns:
        a decorator


def generate_dequeue_op(self, tpu_device=0):
    self.freeze()
    if self._generated_dequeue_op and (not ops.inside_function()):
        raise ValueError("Can't generate two dequeue Ops from the same queue")
    self._generated_dequeue_op = True
    full_name = '%s/dequeue' % self._name
    sharded_shapes = [policy.get_sharded_shape(shape) for shape, policy in zip(self._tuple_shapes, self._sharding_policies)]
    with ops.device(tpu_name_util.core(tpu_device)):
        values = tpu_ops.infeed_dequeue_tuple(dtypes=self._tuple_types, shapes=sharded_shapes, name=full_name)
    return tag_sharding_attribute_for_dequeued_tensors(values, self._input_partition_dims)

Generate TPU dequeue ops.

Args:
        tpu_device: The TPU device ordinal where the infeed instruction should be
        placed.

Returns:
        A list of Outputs corresponding to a partition of infeed dequeued
        into XLA, suitable for use within a replicated block.

Raises:
        ValueError: if the types or shapes of the tuple elements have not been
        set; or if a dequeue op has already been generated.


def __init__(self, state_handler: sdk_worker.CachingStateHandler, transform_id: str, key_coder: coders.Coder, window_coder: coders.Coder) -> None:
    self._state_handler = state_handler
    self._transform_id = transform_id
    self._key_coder = key_coder
    self._window_coder = window_coder
    self._timers_info: Dict[str, TimerInfo] = {}
    self._all_states: Dict[tuple, FnApiUserRuntimeStateTypes] = {}

Initialize a ``FnApiUserStateContext``.

Args:
        state_handler: A StateServicer object.
        transform_id: The name of the PTransform that this context is associated.
        key_coder: Coder for the key type.
        window_coder: Coder for the window type.


def create(configs):
    if not configs:
        raise Error(ANDROID_DEVICE_EMPTY_CONFIG_MSG)
    elif configs == ANDROID_DEVICE_PICK_ALL_TOKEN:
        ads = get_all_instances()
    elif not isinstance(configs, list):
        raise Error(ANDROID_DEVICE_NOT_LIST_CONFIG_MSG)
    elif isinstance(configs[0], dict):
        ads = get_instances_with_configs(configs)
    elif isinstance(configs[0], str):
        ads = get_instances(configs)
    else:
        raise Error('No valid config found in: %s' % configs)
    _start_services_on_ads(ads)
    return ads

Creates AndroidDevice controller objects.

Args:
        configs: Represents configurations for Android devices, this can take one of
        the following forms:
        * str, only asterisk symbol is accepted, indicating that all connected
        Android devices will be used
        * A list of dict, each representing a configuration for an Android device.
        * A list of str, each representing the serial number of Android device.

Returns:
        A list of AndroidDevice objects.


def resize(self, image: np.ndarray, size: Dict[str, int], resample: PILImageResampling=PILImageResampling.BILINEAR, data_format: Optional[Union[str, ChannelDimension]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None, **kwargs) -> np.ndarray:
    if 'shortest_edge' in size and 'longest_edge' in size:
        size = get_resize_output_image_size(image, size, input_data_format)
    elif 'height' in size and 'width' in size:
        size = (size['height'], size['width'])
    else:
        raise ValueError("size must be a dictionary with keys 'shortest_edge' and 'longest_edge' or 'height' and 'width'.")
    return resize(image, 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]`):
            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 get_channel_dimension_axis(image: np.ndarray, input_data_format: Optional[Union[ChannelDimension, str]]=None) -> int:
    if input_data_format is None:
        input_data_format = infer_channel_dimension_format(image)
    if input_data_format == ChannelDimension.FIRST:
        return image.ndim - 3
    elif input_data_format == ChannelDimension.LAST:
        return image.ndim - 1
    raise ValueError(f'Unsupported data format: {input_data_format}')

Returns the channel dimension axis of the image.

Args:
        image (`np.ndarray`):
            The image to get the channel dimension axis of.
        input_data_format (`ChannelDimension` or `str`, *optional*):
            The channel dimension format of the image. If `None`, will infer the channel dimension from the image.

Returns:
        The channel dimension axis of the image.


def _CTCLossGrad(op, grad_loss, _):
    return _CTCLossGradImpl(op, grad_loss, _)

The derivative provided by CTC Loss.

Args:
        op: the CTCLoss op.
        grad_loss: The backprop for cost.

Returns:
        The CTC Loss gradient.


def count_nonzero(x, axis=None):
    if any_symbolic_tensors((x,)):
        return CountNonzero(axis=axis).symbolic_call(x)
    return backend.numpy.count_nonzero(x, axis=axis)

Counts the number of non-zero values in `x` along the given `axis`.

    If no axis is specified then all non-zeros in the tensor are counted.

Args:
        x: Input tensor.
        axis: Axis or tuple of axes along which to count the number of
            non-zeros. Defaults to `None`.

Returns:
        int or tensor of ints.

Example:
    >>> x = keras.ops.array([[0, 1, 7, 0], [3, 0, 2, 19]])
    >>> keras.ops.count_nonzero(x)
    5
    >>> keras.ops.count_nonzero(x, axis=0)
    array([1, 1, 2, 1], dtype=int64)
    >>> keras.ops.count_nonzero(x, axis=1)
    array([2, 3], dtype=int64)


def __init__(self, parameter_name=None):
    self._parameter_name = parameter_name

Construct a parameter modifier that may be specific to a parameter.

Args:
        parameter_name:  A `ParameterModifier` instance may operate on a class of
        parameters or on a parameter with a particular name.  Only
        `ParameterModifier` instances that are of a unique type or were
        initialized with a unique `parameter_name` will be executed.
        See `__eq__` and `__hash__`.


def log_sinkhorn_iterations(log_cost_matrix: torch.Tensor, log_source_distribution: torch.Tensor, log_target_distribution: torch.Tensor, num_iterations: int) -> torch.Tensor:
    log_u_scaling = torch.zeros_like(log_source_distribution)
    log_v_scaling = torch.zeros_like(log_target_distribution)
    for _ in range(num_iterations):
        log_u_scaling = log_source_distribution - torch.logsumexp(log_cost_matrix + log_v_scaling.unsqueeze(1), dim=2)
        log_v_scaling = log_target_distribution - torch.logsumexp(log_cost_matrix + log_u_scaling.unsqueeze(2), dim=1)
    return log_cost_matrix + log_u_scaling.unsqueeze(2) + log_v_scaling.unsqueeze(1)

Perform Sinkhorn Normalization in Log-space for stability

Args:
        log_cost_matrix (`torch.Tensor` of shape `(batch_size, num_rows, num_columns)`):
            Logarithm of the cost matrix.
        log_source_distribution (`torch.Tensor` of shape `(batch_size, num_rows)`):
            Logarithm of the source distribution.
        log_target_distribution (`torch.Tensor` of shape `(batch_size, num_columns)`):
            Logarithm of the target distribution.

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


def _run_test_class(self, config, test_class, tests=None):
    test_instance = test_class(config)
    logging.debug('Executing test class "%s" with config: %s', test_class.__name__, config)
    try:
        cls_result = test_instance.run(tests)
        self.results += cls_result
    except signals.TestAbortAll as e:
        self.results += e.results
        raise e

Instantiates and executes a test class.

    If tests is None, the tests listed in self.tests will be executed
    instead. If self.tests is empty as well, every test in this test class
    will be executed.

Args:
        config: A config_parser.TestRunConfig object.
        test_class: class, test class to execute.
        tests: Optional list of test names within the class to execute.


def benchmarks_main(true_main, argv=None):
    if argv is None:
        argv = sys.argv
    found_arg = [arg for arg in argv if arg.startswith('--benchmark_filter=') or arg.startswith('-benchmark_filter=')]
    if found_arg:
        argv.remove(found_arg[0])
        regex = found_arg[0].split('=')[1]
        app.run(lambda _: _run_benchmarks(regex), argv=argv)
    else:
        true_main()

Run benchmarks as declared in argv.

Args:
        true_main: True main function to run if benchmarks are not requested.
        argv: the command line arguments (if None, uses sys.argv).


def run_saved_model_with_feed_dict(saved_model_dir, tag_set, signature_def_key, input_tensor_key_feed_dict, outdir, overwrite_flag, worker=None, init_tpu=False, use_tfrt=False, tf_debug=False):
    meta_graph_def = saved_model_utils.get_meta_graph_def(saved_model_dir, tag_set)
    inputs_tensor_info = _get_inputs_tensor_info_from_meta_graph_def(meta_graph_def, signature_def_key)
    for input_key_name in input_tensor_key_feed_dict.keys():
        if input_key_name not in inputs_tensor_info:
            raise ValueError('"%s" is not a valid input key. Please choose from %s, or use --show option.' % (input_key_name, '"' + '", "'.join(inputs_tensor_info.keys()) + '"'))
    inputs_feed_dict = {inputs_tensor_info[key].name: tensor for key, tensor in input_tensor_key_feed_dict.items()}
    outputs_tensor_info = _get_outputs_tensor_info_from_meta_graph_def(meta_graph_def, signature_def_key)
    output_tensor_keys_sorted = sorted(outputs_tensor_info.keys())
    output_tensor_names_sorted = [outputs_tensor_info[tensor_key].name for tensor_key in output_tensor_keys_sorted]
    config = None
    if use_tfrt:
        logging.info('Using TFRT session.')
        config = config_pb2.ConfigProto(experimental=config_pb2.ConfigProto.Experimental(use_tfrt=True))
    with session.Session(worker, graph=ops_lib.Graph(), config=config) as sess:
        if init_tpu:
            print('Initializing TPU System ...')
            sess.run(tpu.initialize_system())
        loader.load(sess, tag_set.split(','), saved_model_dir)
        if tf_debug:
            sess = local_cli_wrapper.LocalCLIDebugWrapperSession(sess)
        outputs = sess.run(output_tensor_names_sorted, feed_dict=inputs_feed_dict)
        for i, output in enumerate(outputs):
            output_tensor_key = output_tensor_keys_sorted[i]
            print('Result for output key %s:\n%s' % (output_tensor_key, output))
            if outdir:
                if not os.path.isdir(outdir):
                    os.makedirs(outdir)
                output_full_path = os.path.join(outdir, output_tensor_key + '.npy')
                if not overwrite_flag and os.path.exists(output_full_path):
                    raise RuntimeError('Output file %s already exists. Add "--overwrite" to overwrite the existing output files.' % output_full_path)
                np.save(output_full_path, output)
                print('Output %s is saved to %s' % (output_tensor_key, output_full_path))

Runs SavedModel and fetch all outputs.

    Runs the input dictionary through the MetaGraphDef within a SavedModel
    specified by the given tag_set and SignatureDef. Also save the outputs to file
    if outdir is not None.

Args:
        saved_model_dir: Directory containing the SavedModel to execute.
        tag_set: Group of tag(s) of the MetaGraphDef with the SignatureDef map, in
        string format, separated by ','. For tag-set contains multiple tags, all
        tags must be passed in.
        signature_def_key: A SignatureDef key string.
        input_tensor_key_feed_dict: A dictionary maps input keys to numpy ndarrays.
        outdir: A directory to save the outputs to. If the directory doesn't exist,
        it will be created.
        overwrite_flag: A boolean flag to allow overwrite output file if file with
        the same name exists.
        worker: If provided, the session will be run on the worker.  Valid worker
        specification is a bns or gRPC path.
        init_tpu: If true, the TPU system will be initialized after the session
        is created.
        use_tfrt: If true, TFRT session will be used.
        tf_debug: A boolean flag to use TensorFlow Debugger (TFDBG) to observe the
        intermediate Tensor values and runtime GraphDefs while running the
        SavedModel.

Raises:
        ValueError: When any of the input tensor keys is not valid.
        RuntimeError: An error when output file already exists and overwrite is not
        enabled.


def resize_position_embeddings(self, new_num_position_embeddings: int):
    logger.info(f'Setting `config.max_position_embeddings={new_num_position_embeddings}`...')
    self.config.max_position_embeddings = new_num_position_embeddings
    self.embed_positions = PegasusXSinusoidalPositionalEmbedding(self.config.d_model)
    self.embed_positions.to(self.device)

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

Args:
        new_num_position_embeddings (`int`):
            The number of new position embeddings. 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, pixel_values: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> ImageClassifierOutputWithNoAttention:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    outputs = self.textnet(pixel_values, output_hidden_states=output_hidden_states, return_dict=return_dict)
    last_hidden_state = outputs[0]
    for layer in self.classifier:
        last_hidden_state = layer(last_hidden_state)
    logits = self.fc(last_hidden_state)
    loss = None
    if labels is not None:
        if self.config.problem_type is None:
            if self.num_labels == 1:
                self.config.problem_type = 'regression'
            elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
                self.config.problem_type = 'single_label_classification'
            else:
                self.config.problem_type = 'multi_label_classification'
        if self.config.problem_type == 'regression':
            loss_fct = MSELoss()
            if self.num_labels == 1:
                loss = loss_fct(logits.squeeze(), labels.squeeze())
            else:
                loss = loss_fct(logits, labels)
        elif self.config.problem_type == 'single_label_classification':
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
        elif self.config.problem_type == 'multi_label_classification':
            loss_fct = BCEWithLogitsLoss()
            loss = loss_fct(logits, labels)
    if not return_dict:
        output = (logits,) + outputs[2:]
        return (loss,) + output if loss is not None else output
    return ImageClassifierOutputWithNoAttention(loss=loss, logits=logits, hidden_states=outputs.hidden_states)

labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
    Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
    config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
    `config.num_labels > 1` a classification loss is computed (Cross-Entropy).

Example:
    ```python
    >>> import torch
    >>> import requests
    >>> from transformers import TextNetForImageClassification, TextNetImageProcessor
    >>> from PIL import Image

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

    >>> processor = TextNetImageProcessor.from_pretrained("czczup/textnet-base")
    >>> model = TextNetForImageClassification.from_pretrained("czczup/textnet-base")

    >>> inputs = processor(images=image, return_tensors="pt")
    >>> with torch.no_grad():
    ...     outputs = model(**inputs)
    >>> outputs.logits.shape
    torch.Size([1, 2])
    ```


def _compute_gradients_wrt_embedding_table(self, gradient_wrt_activation, embedding_table, feature_indices, feature_values, combiner):
    if combiner not in ('mean', 'sum'):
        raise ValueError('`combiner` must be mean or sum; got {}.'.format(combiner))
    grads_shape = gradient_wrt_activation.shape[:-1] + embedding_table.shape
    grads = np.zeros(shape=grads_shape)
    count = np.zeros(shape=grads_shape)
    for feature_indice, vocabulary_id in zip(feature_indices, feature_values):
        batch_index = tuple(feature_indice[:-1])
        grads[batch_index][vocabulary_id] += gradient_wrt_activation[batch_index]
        count[batch_index] += 1
    count[count == 0] = 1
    if combiner == 'mean':
        grads = grads / count
    return np.reshape(grads, (-1, *embedding_table.shape))

Compute gradients wrt embedding_table.

Args:
        gradient_wrt_activation: `np.array` with shape `batch_size` by embedding
        `dimension`.
        embedding_table: `np.array` with shape `vocabulary_size` by embedding
        `dimension`.
        feature_indices: `indices` as used to construct `SparseTensor`.
        feature_values: `values` as used to construct `SparseTensor`.
        combiner: `String`, 'mean' or 'sum'.

Returns:
        Gradients wrt `embedding_table`, an `np.array`s with shape
        `batch_size` by `vocabulary_size` by
        embedding `dimension`.

Raises:
        ValueError: if `combiner` is not one of 'mean' or 'sum'.


def assert_scalar_v2(tensor, message=None, name=None):
    assert_scalar(tensor=tensor, message=message, name=name)

Asserts that the given `tensor` is a scalar.

    This function raises `ValueError` unless it can be certain that the given
    `tensor` is a scalar. `ValueError` is also raised if the shape of `tensor` is
    unknown.

    This is always checked statically, so this method returns nothing.

Args:
        tensor: A `Tensor`.
        message: A string to prefix to the default message.
        name:  A name for this operation. Defaults to "assert_scalar"

Raises:
        ValueError: If the tensor is not scalar (rank 0), or if its shape is
        unknown.


def _next_power_of_two(x):
    return 1 if x == 0 else 2 ** (int(x) - 1).bit_length()

Calculates the smallest enclosing power of two for an input.

Args:
        x: Positive float or integer number.

Returns:
        Next largest power of two integer.


def is_emulator(self):
    if EMULATOR_SERIAL_REGEX.match(self.serial):
        return True
    elif self.build_info['build_characteristics'] == 'emulator':
        return True
    elif self.build_info['hardware'] in ['ranchu', 'goldfish', 'cutf_cvm']:
        return True
    else:
        return False

Whether this device is probably an emulator.

Returns:
        True if this is probably an emulator.


def get_sequence_length_feature_key_name_from_feature_key_name(feature_name):
    return feature_name + _SEQUENCE_FEATURE_LENGTH_POSTFIX

Gets the name of the sequence length feature from that of the base feature.

Args:
        feature_name: The feature key of a sequence column.

Returns:
        A string which is the feature key for the associated feature length column.


def __init__(self, quantum_h_terms: List[cirq.PauliSum], classical_h_terms: List[cirq.PauliSum], num_layers: int, initializer=tf.keras.initializers.RandomUniform(0, 2 * np.pi), name=None):
    quantum_symbols = []
    classical_symbols = []
    for j in range(num_layers):
        quantum_symbols.append([])
        classical_symbols.append([])
        for k, _ in enumerate(quantum_h_terms):
            quantum_symbols[-1].append(f'gamma_{j}_{k}')
        for k, _ in enumerate(classical_h_terms):
            classical_symbols[-1].append(f'eta_{j}_{k}')
    pqc = cirq.Circuit()
    flat_symbols = []
    for q_symb, c_symb in zip(quantum_symbols, classical_symbols):
        pqc += tfq.util.exponential(quantum_h_terms, coefficients=q_symb)
        pqc += tfq.util.exponential(classical_h_terms, coefficients=c_symb)
        flat_symbols.extend(q_symb + c_symb)
    symbol_names = tf.constant(flat_symbols)
    value_layers_inputs = [[tf.Variable(initializer(shape=[num_layers])), tf.Variable(initializer(shape=[len(classical_h_terms)])), tf.Variable(initializer(shape=[num_layers, len(quantum_h_terms)]))]]

    def embed_params(inputs):
        """Tiles up the variables to properly tie QAIA parameters."""
        exp_etas = tf.expand_dims(inputs[0], 1)
        tiled_thetas = tf.tile(tf.expand_dims(inputs[1], 0), [tf.shape(inputs[0])[0], 1])
        classical_params = exp_etas * tiled_thetas
        return tf.reshape(tf.concat([classical_params, inputs[2]], 1), [-1])
    value_layers = [[tf.keras.layers.Lambda(embed_params)]]
    super().__init__(tfq.convert_to_tensor([pqc]), pqc.all_qubits(), symbol_names, value_layers_inputs, value_layers)

Initializes a QAIA.

    The ansatz is QAOA-like, with the exponential of the EBM ansatz in place
    of the usual "problem Hamiltonian". Mathematically, it is represented as:

    $$\prod_{\ell=1}^P \left[
        \left(
        \prod_{\bm{b} \in \mathcal{B}_K}
            e^{i\eta_\ell \theta_{\bm{b}}\bm{\hat{Z}}^{\bm{b}}}
        \right)\left(
        \prod_{r\in \mathcal{I}}
            e^{i\gamma_{r\ell}\hat{H}_r}
        \right)
        \right],$$

    where $\hat{H}_r$ is `quantum_h_terms`, $\bm{\hat{Z}}^{\bm{b}}$ is
    `classical_h_terms`, and $P$ is `num_layers`.

    # TODO(#119): add link to new version of the paper.
    For further discussion, see the section "Physics-Inspired Architecture:
    Quantum Adiabatic-Inspired Ansatz" in the QHBM paper.

Args:
        quantum_h_terms: Non-commuting terms of the target thermal state
        assumed in the QAIA ansatz.
        classical_h_terms: Hamiltonian representation of the EBM chosen to model
        the target thermal state.
        num_layers: How many layers of the ansatz to apply.
        initializer: A `tf.keras.initializers.Initializer` which specifies how to
        initialize the values of the parameters in `circuit`.
        name: Optional name for the model.


def parse_float(value: Any) -> Numeric:
    return float(value)

Attempts to parse a valid floating point value from the provided value.

Args:
        * value: of Any type

Returns:
        * float value: if valid

Raises:
        * ValueError: if parsing failed


def ops_used_by_graph_def(graph_def):
    name_to_function = {}
    for fun in graph_def.library.function:
        name_to_function[fun.signature.name] = fun
    used_ops = set()
    functions_to_process = []

    def mark_op_as_used(op):
        if op not in used_ops and op in name_to_function:
            functions_to_process.append(name_to_function[op])
        used_ops.add(op)

    def process_node(node):
        mark_op_as_used(node.op)
        if node.op in ['PartitionedCall', 'StatefulPartitionedCall']:
            mark_op_as_used(node.attr['f'].func.name)
    for node in graph_def.node:
        process_node(node)
    while functions_to_process:
        fun = functions_to_process.pop()
        for node in fun.node_def:
            process_node(node)
    return [op for op in used_ops if op not in name_to_function]

Collect the list of ops used by a graph.

    Does not validate that the ops are all registered.

Args:
        graph_def: A `GraphDef` proto, as from `graph.as_graph_def()`.

Returns:
        A list of strings, each naming an op used by the graph.


def _try_load_par_source(source_file_path):
    prefix_path = source_file_path
    while True:
        prefix_path, basename = os.path.split(prefix_path)
        if not basename:
            break
        suffix_path = os.path.normpath(os.path.relpath(source_file_path, start=prefix_path))
        if prefix_path.endswith('.par') and os.path.isfile(prefix_path):
            with zipfile.ZipFile(prefix_path) as z:
                norm_names = [os.path.normpath(name) for name in z.namelist()]
                if suffix_path in norm_names:
                    with z.open(z.namelist()[norm_names.index(suffix_path)]) as zf:
                        source_text = zf.read().decode('utf-8')
                        return source_text.split('\n')

Try loading the source code inside a .par file.

    A .par file is a zip-compressed, self-contained Python executable.
    It contains the content of individual Python source files that can
    be read only through extracting from the zip file.

Args:
        source_file_path: The full path to the file inside the .par file. This
        path should include the path to the .par file itself, followed by the
        intra-par path, e.g.,
        "/tmp/my_executable.par/org-tensorflow/tensorflow/python/foo/bar.py".

Returns:
        If successful, lines of the source file as a `list` of `str`s.
        Else, `None`.


def __init__(self):
    self.return_types = []
    self.exceptions = []

Mutable class for collecting return types and exceptions of functions.

    The collecting is stable: Items are kept in the order in which they were
    encountered.

Attributes:
        return_types: Return types seen so far.
        exceptions: Exceptions seen so far.


def shape_equal(shape1, shape2, axis=None, allow_none=True):
    if len(shape1) != len(shape2):
        return False
    shape1 = list(shape1)
    shape2 = list(shape2)
    if axis is not None:
        if isinstance(axis, int):
            axis = [axis]
        for ax in axis:
            shape1[ax] = -1
            shape2[ax] = -1
    if allow_none:
        for i in range(len(shape1)):
            if shape1[i] is None:
                shape1[i] = shape2[i]
            if shape2[i] is None:
                shape2[i] = shape1[i]
    return shape1 == shape2

Check if two shapes are equal.

Args:
        shape1: A list or tuple of integers for first shape to be compared.
        shape2: A list or tuple of integers for second shape to be compared.
        axis: An integer, list, or tuple of integers (optional):
            Axes to ignore during comparison. Defaults to `None`.
        allow_none (bool, optional): If `True`, allows `None` in a shape
            to match any value in the corresponding position of the other shape.
            Defaults to `True`.

Returns:
        bool: `True` if shapes are considered equal based on the criteria,
        `False` otherwise.

Example:
    >>> shape_equal((32, 64, 128), (32, 64, 128))
    True
    >>> shape_equal((32, 64, 128), (32, 64, 127))
    False
    >>> shape_equal((32, 64, None), (32, 64, 128), allow_none=True)
    True
    >>> shape_equal((32, 64, None), (32, 64, 128), allow_none=False)
    False
    >>> shape_equal((32, 64, 128), (32, 63, 128), axis=1)
    True
    >>> shape_equal((32, 64, 128), (32, 63, 127), axis=(1, 2))
    True
    >>> shape_equal((32, 64, 128), (32, 63, 127), axis=[1,2])
    True
    >>> shape_equal((32, 64), (32, 64, 128))
    False


def _kernel_label_map(self, op_to_kernel_label_map) -> Iterator[None]:
    if not isinstance(op_to_kernel_label_map, dict):
        raise TypeError('op_to_kernel_label_map must be a dictionary mapping strings to strings')
    saved_labels = {}
    for op_type, label in op_to_kernel_label_map.items():
        if not (isinstance(op_type, str) and isinstance(label, str)):
            raise TypeError('op_to_kernel_label_map must be a dictionary mapping strings to strings')
        try:
            saved_labels[op_type] = self._op_to_kernel_label_map[op_type]
        except KeyError:
            pass
        self._op_to_kernel_label_map[op_type] = label
    try:
        yield
    finally:
        for op_type, label in op_to_kernel_label_map.items():
            try:
                self._op_to_kernel_label_map[op_type] = saved_labels[op_type]
            except KeyError:
                del self._op_to_kernel_label_map[op_type]

EXPERIMENTAL: A context manager for setting kernel labels.

    This context manager can be used to select particular
    implementations of kernels within the scope of the context.

    For example:

        with ops.Graph().as_default() as g:
        f_1 = Foo()  # Uses the default registered kernel for the Foo op.
        with g.kernel_label_map({"Foo": "v_2"}):
            f_2 = Foo()  # Uses the registered kernel with label "v_2"
                        # for the Foo op.
            with g.kernel_label_map({"Foo": "v_3"}):
            f_3 = Foo()  # Uses the registered kernel with label "v_3"
                        # for the Foo op.
            with g.kernel_label_map({"Foo": ""}):
                f_4 = Foo()  # Uses the default registered kernel
                            # for the Foo op.

Args:
        op_to_kernel_label_map: A dictionary mapping op type strings to kernel
        label strings.

Returns:
        A context manager that sets the kernel label to be used for one or more
        ops created in that context.

Raises:
        TypeError: If op_to_kernel_label_map is not a dictionary mapping
        strings to strings.


def __init__(self, l2=0.01):
    l2 = 0.01 if l2 is None else l2
    validate_float_arg(l2, name='l2')
    self.l2 = l2

A regularizer that applies a L2 regularization penalty.

    The L2 regularization penalty is computed as:
    `loss = l2 * reduce_sum(square(x))`

    L2 may be passed to a layer as a string identifier:

    >>> dense = Dense(3, kernel_regularizer='l2')

    In this case, the default value used is `l2=0.01`.

Args:
        l2: float, L2 regularization factor.