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

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


def init_cache(self, batch_size, max_length):
    input_ids = jnp.ones((batch_size, max_length), dtype='i4')
    attention_mask = jnp.ones_like(input_ids)
    init_variables = self.module.init(jax.random.PRNGKey(0), input_ids, attention_mask, return_dict=False, init_cache=True)
    return unfreeze(init_variables['cache'])

Args:
    batch_size (`int`):
        batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache.
    max_length (`int`):
        maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized
        cache.


def _resolve_grad_captures(body_graph, body_grad_graph, while_op):
    new_capture_inputs = []
    for t in body_grad_graph.external_captures:
        if t.graph == body_graph:
            for i, output in enumerate(t.graph.outputs):
                if output is t:
                    t = while_op.outputs[i]
                    break
            assert t.graph == body_graph.outer_graph
        new_capture_inputs.append(t)
    return new_capture_inputs

Returns the tensors to pass as captured inputs to `body_grad_graph`.

    `body_grad_graph` may have external references to:
    1. Its outer graph containing the input gradients. These are left as-is.
    2. Accumulators captured from the forward-pass graph. These should have been
        added as `while_op` outputs after the gradient graph was built. We replace
        these with the corresponding output of `while_op`, i.e. a tensor in
        `body_graph.outer_graph`. In the case of nested control flow or functions,
        the gradient logic handling `body_grad_graph.outer_graph` will make sure
        the tensor from `body_graph.outer_graph` is also correctly captured.

Args:
        body_graph: FuncGraph. The forward-pass body function.
        body_grad_graph: FuncGraph. The body gradients function.
        while_op: The forward-pass While Operation calling `body_graph`.

Returns:
        A list of input tensors to be passed as the captured inputs to
        `body_grad_graph`.


def _from_pyval(cls, pyval, typespec, path_so_far):
    if isinstance(pyval, dict):
        return cls._from_pydict(pyval, typespec, path_so_far)
    elif isinstance(pyval, (list, tuple)):
        keys = set()
        rank = _pyval_find_struct_keys_and_depth(pyval, keys)
        if rank is not None:
            return cls._from_pylist_of_dict(pyval, keys, rank, typespec, path_so_far)
        else:
            return cls._from_pylist_of_value(pyval, typespec, path_so_far)
    else:
        return cls._from_pyscalar(pyval, typespec, path_so_far)

Helper function for from_pyval.

Args:
        pyval: The nested Python structure that should be used to create the new
        `StructuredTensor`.
        typespec: A `StructuredTensor.Spec` specifying the expected type for each
        field. If not specified, then all nested dictionaries are turned into
        StructuredTensors, and all nested lists are turned into Tensors (if
        rank<2) or RaggedTensors (if rank>=2).
        path_so_far: the path of fields that led here (for error messages).

Returns:
        A `StructuredTensor`.


def _key2seed(a):

    def int64_to_int32s(a):
        """Converts an int64 tensor of shape [] to an int32 tensor of shape [2]."""
        a = math_ops.cast(a, dtypes.uint64)
        fst = math_ops.cast(a, dtypes.uint32)
        snd = math_ops.cast(gen_bitwise_ops.right_shift(a, constant_op.constant(32, dtypes.uint64)), dtypes.uint32)
        a = [fst, snd]
        a = nest.map_structure(lambda x: math_ops.cast(x, dtypes.int32), a)
        a = array_ops_stack.stack(a)
        return a
    return int64_to_int32s(a)

Converts an RNG key to an RNG seed.

Args:
        a: an RNG key, an ndarray of shape [] and dtype `np.int64`.

Returns:
        an RNG seed, a tensor of shape [2] and dtype `tf.int32`.


def embedding_lookup(self, features: Any, weights: Optional[Any]=None) -> Any:
    if not self._built:
        self.build()
    nest.assert_same_structure(features, self._feature_config)
    flat_inputs = nest.flatten(features)
    flat_weights = [None] * len(flat_inputs)
    if weights is not None:
        nest.assert_same_structure(features, weights)
        flat_weights = nest.flatten(weights)
    flat_features = nest.flatten_with_joined_string_paths(self._feature_config)
    outputs = []
    for inp, weight, (path, feature) in zip(flat_inputs, flat_weights, flat_features):
        table = self.embedding_tables[feature.table]
        if weight is not None:
            if isinstance(inp, tensor.Tensor):
                raise ValueError('Weight specified for {}, but input is dense.'.format(path))
            elif type(weight) is not type(inp):
                raise ValueError('Weight for {} is of type {} but it does not match type of the input which is {}.'.format(path, type(weight), type(inp)))
            elif feature.max_sequence_length > 0:
                raise ValueError('Weight specified for {}, but this is a sequence feature.'.format(path))
        if isinstance(inp, tensor.Tensor):
            if feature.max_sequence_length > 0:
                raise ValueError('Feature {} is a sequence feature but a dense tensor was passed.'.format(path))
            outputs.append(embedding_ops.embedding_lookup_v2(table, inp))
        elif isinstance(inp, sparse_tensor.SparseTensor):
            outputs.append(self._embedding_lookup_for_sparse_tensor(inp, weight, table, feature))
        elif isinstance(inp, ragged_tensor.RaggedTensor):
            outputs.append(self._embedding_lookup_for_ragged_tensor(inp, weight, table, feature))
        else:
            raise ValueError('Input {} is type {}. Tensor, SparseTensor or RaggedTensor expected.'.format(path, type(inp)))
    return nest.pack_sequence_as(self._feature_config, outputs)

Apply embedding lookup on TPUs using Tensorcore.

    Note that all the sparse and ragged tensors will be converted to dense
    tensors on CPU and then passed to the TPU to do embedding look up. Large
    embedding lookup is not supported by this API, use the TPUEmbedding mid
    level api instead.

Args:
        features: a nested structure of Tensors, SparseTensors or RaggedTensors.
        weights: a nested structure of Tensors, SparseTensors or RaggedTensors or
        None for no weights. If not None, structure must match that of inputs,
        but entries are allowed to be None.

Returns:
        A nested structure of Tensors with the same structure as inputs.


def __init__(self, model: PreTrainedModel, max_batch_size: int=1, max_cache_len: int=4096):
    super().__init__()
    if not hasattr(model.config, 'use_cache') or model.config.use_cache is False:
        raise ValueError('The model must have caching enabled to be performant.')
    if hasattr(model.config, 'layer_types') and getattr(model.config, 'sliding_window', None) is not None:
        self.model = TorchExportableModuleWithHybridCache(model, max_batch_size, max_cache_len)
    else:
        logging.info('Using `StaticCache` for export as `layer_types` is not specified or `sliding_window` is `null` in the config.')
        self.model = TorchExportableModuleWithStaticCache(model)

Initializes the exportable module with `HybridCache`.

Args:
        model (`PreTrainedModel`): The pretrained model to wrap.
        max_batch_size (int): Maximum batch size for the cache.
        max_cache_len (int): Maximum sequence length for the cache.

Raises:
        ValueError: If the model is configured with a unsupported cache implementation.


def xxd_output_to_object(input_cc_file):
    model_bytes = xxd_output_to_bytes(input_cc_file)
    return convert_bytearray_to_object(model_bytes)

Converts xxd output C++ source file to object.

Args:
        input_cc_file: Full path name to th C++ source file dumped by xxd

Raises:
        RuntimeError: If input_cc_file path is invalid.
        IOError: If input_cc_file cannot be opened.

Returns:
        A python object corresponding to the input tflite file.


def update_state(self, y_true, y_pred, sample_weight=None):
    y_true = math_ops.cast(y_true, self._dtype)
    y_pred = math_ops.cast(y_pred, self._dtype)
    [y_true, y_pred], sample_weight = metrics_utils.ragged_assert_compatible_and_get_flat_values([y_true, y_pred], sample_weight)
    y_pred, y_true = losses_utils.squeeze_or_expand_dimensions(y_pred, y_true)
    ag_fn = autograph.tf_convert(self._fn, ag_ctx.control_status_ctx())
    matches = ag_fn(y_true, y_pred, **self._fn_kwargs)
    return super(MeanMetricWrapper, self).update_state(matches, sample_weight=sample_weight)

Accumulates metric statistics.

    `y_true` and `y_pred` should have the same shape.

Args:
        y_true: Ground truth values. shape = `[batch_size, d0, .. dN]`.
        y_pred: The predicted values. shape = `[batch_size, d0, .. dN]`.
        sample_weight: Optional `sample_weight` acts as a
        coefficient for the metric. If a scalar is provided, then the metric is
        simply scaled by the given value. If `sample_weight` is a tensor of size
        `[batch_size]`, then the metric for each sample of the batch is rescaled
        by the corresponding element in the `sample_weight` vector. If the shape
        of `sample_weight` is `[batch_size, d0, .. dN-1]` (or can be broadcasted
        to this shape), then each metric element of `y_pred` is scaled by the
        corresponding value of `sample_weight`. (Note on `dN-1`: all metric
        functions reduce by 1 dimension, usually the last axis (-1)).

Returns:
        Update op.


def _NonEagerInputs(op: ops.Operation, xs_set):
    return [t for t in _Inputs(op, xs_set) if not isinstance(t, ops.EagerTensor)]

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

    Does not return any captured EagerTensors, i.e., the number of tensors
    returned may be less than the actual number of inputs.

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

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


def __init__(self, x: int, *args, y: str, **kwargs) -> float:
    del x, y, args, kwargs

Constructor.

Args:
        x: Input 1.
        *args: Variable positional args.
        y: Input 2.
        **kwargs: Variable keyword args.

Returns:
        The result.


def get_special_tokens_mask(self, token_ids_0: List[int], token_ids_1: Optional[List[int]]=None, already_has_special_tokens: Optional[bool]=False) -> List[int]:
    if already_has_special_tokens:
        return super().get_special_tokens_mask(token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True)
    if token_ids_1 is None:
        return [0] * len(token_ids_0) + [1]
    return [0] * len(token_ids_0) + [1] + [0] * len(token_ids_1) + [1]

Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
    special tokens using the tokenizer `prepare_for_model` method.

Args:
        token_ids_0 (`List[int]`):
            List of IDs.
        token_ids_1 (`List[int]`, *optional*):
            Optional second list of IDs for sequence pairs.
        already_has_special_tokens (`bool`, *optional*, defaults to `False`):
            Whether or not the token list is already formatted with special tokens for the model.

Returns:
        `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.


def _output_ast(ast: pytd.TypeDeclUnit, options: config.Options | None=None) -> str:
    options = options or config.Options.create()
    with config.verbosity_from(options):
        result = pytd_utils.Print(ast)
        log.info('=========== pyi optimized =============')
        log.info('\n%s', result)
        log.info('========================================')
    result += '\n'
    if options.quick:
        result = '# (generated with --quick)\n\n' + result
    return result

Transforms the given analysis result into a pyi representation.

Args:
        ast: pytd.TypeDeclUnit to output in pyi format.
        options: config.Options object.

Returns:
        A pyi representation of the given AST as a string.


def forward(self, pixel_values: Optional[torch.FloatTensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[Tuple, BaseModelOutputWithPooling]:
    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
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    return self.vision_model(pixel_values=pixel_values, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)

Example:

    ```python
    >>> import requests
    >>> from PIL import Image
    >>> from transformers import AutoProcessor, MLCDVisionModel
    >>> model = MLCDVisionModel.from_pretrained("DeepGlint-AI/mlcd-vit-bigG-patch14-448")
    >>> processor = AutoProcessor.from_pretrained("DeepGlint-AI/mlcd-vit-bigG-patch14-448")

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

    >>> with torch.no_grad():
    ...     outputs = model(**inputs, output_attentions=True)

    >>> features = outputs.last_hidden_state
    >>> print(f"Extracted features shape: {features.shape}")
    >>> print(f"Number of attention layers: {len(outputs.attentions)}")
    >>> print(f"Attention shape: {outputs.attentions[0].shape}")
    ```


def VerifyMaps(self, conf):
    retval = 0
    for map_name in conf.maps:
        self.log.info('Verifying map: %s.', map_name)
        if map_name == config.MAP_NETGROUP:
            self.log.info('The netgroup map does not support enumeration, skipping.')
            continue
        if map_name == config.MAP_AUTOMOUNT:
            self.log.info('The automount map does not support enumeration, skipping.')
            continue
        try:
            nss_map = nss.GetMap(map_name)
        except error.UnsupportedMap:
            self.log.warning('Verification of %s map is unsupported!', map_name)
            continue
        self.log.debug('built NSS map of %d entries', len(nss_map))
        cache_options = conf.options[map_name].cache
        cache = cache_factory.Create(cache_options, map_name)
        try:
            cache_map = cache.GetMap()
        except error.CacheNotFound:
            self.log.error('Cache missing!')
            retval += 1
            continue
        self.log.debug('built cache map of %d entries', len(cache_map))
        missing_entries = 0
        for map_entry in cache_map:
            if map_entry not in nss_map:
                self.log.info('The following entry is present in the cache but not availible via NSS! %s', map_entry.name)
                self.log.debug('missing entry data: %s', map_entry)
                missing_entries += 1
        if missing_entries > 0:
            self.log.warning('Missing %d entries in %s map', missing_entries, map_name)
            retval += 1
    return retval

Compare each configured map against data retrieved from NSS.

    For each configured map, build a Map object from NSS and compare
    it against a Map object retrieved directly from the cache.  We
    expect the cache Map to be a subset of the nss Map due to possible
    inclusion of other NSS map types (e.g. files, nis, ldap, etc).

    This could be done via series of get*nam calls, however at this
    time it appears to be more efficient to grab them in bulk and use
    the Map.__contains__() membership test.

Args:
        conf: nss_cache.config.Config object

Returns:
        count of failures when verifying


def __clean__(struct: Union[dict, list]) -> Union[dict, list]:
    if isinstance(struct, dict):
        for key, value in struct.items():
            if isinstance(value, bytes):
                struct[key] = base64.standard_b64encode(value).decode('ascii')
            elif isinstance(value, date):
                struct[key] = str(value)
            else:
                API.__clean__(value)
    elif isinstance(struct, list):
        for index, value in enumerate(struct):
            if isinstance(value, bytes):
                struct[index] = base64.standard_b64encode(value).decode('ascii')
            elif isinstance(value, date):
                struct[index] = str(value)
            else:
                API.__clean__(value)
    return struct

Helper to recursively clean up JSON data for API call.

    Converts bytes -> base64.
    Converts date -> str (yyyy-mm-dd).
    TODO: Add Converts datetime, time -> string.

Args:
        struct: The kwargs being cleaned up.

Returns:
        struct: The kwargs with replacments.


def fftshift(x, axes=None, name=None):
    with _ops.name_scope(name, 'fftshift') as name:
        x = _ops.convert_to_tensor(x)
        if axes is None:
            axes = tuple(range(x.shape.ndims))
            shift = _array_ops.shape(x) // 2
        elif isinstance(axes, int):
            shift = _array_ops.shape(x)[axes] // 2
        else:
            rank = _array_ops.rank(x)
            axes = _array_ops.where(_math_ops.less(axes, 0), axes + rank, axes)
            shift = _array_ops.gather(_array_ops.shape(x), axes) // 2
        return manip_ops.roll(x, shift, axes, name)

Shift the zero-frequency component to the center of the spectrum.

    This function swaps half-spaces for all axes listed (defaults to all).
    Note that ``y[0]`` is the Nyquist component only if ``len(x)`` is even.

    @compatibility(numpy)
    Equivalent to numpy.fft.fftshift.
    https://docs.scipy.org/doc/numpy/reference/generated/numpy.fft.fftshift.html
    @end_compatibility

    For example:

    ```python
    x = tf.signal.fftshift([ 0.,  1.,  2.,  3.,  4., -5., -4., -3., -2., -1.])
    x.numpy() # array([-5., -4., -3., -2., -1.,  0.,  1.,  2.,  3.,  4.])
    ```

Args:
        x: `Tensor`, input tensor.
        axes: `int` or shape `tuple`, optional Axes over which to shift.  Default is
        None, which shifts all axes.
        name: An optional name for the operation.

Returns:
        A `Tensor`, The shifted tensor.


def UpdateFromSource(self, source, incremental=True, force_write=False):
    cache = cache_factory.Create(self.cache_options, self.map_name)
    return self.UpdateCacheFromSource(cache, source, incremental, force_write, location=None)

Update this map's cache from the source provided.

    The FileMapUpdater expects to fetch as single map from the source
    and write/merge it to disk.  We create a cache to write to, and then call
    UpdateCacheFromSource() with that cache.

    Note that AutomountUpdater also calls UpdateCacheFromSource() for each
    cache it is writing, hence the distinct seperation.

Args:
        source: A nss_cache.sources.Source object.
        incremental: A boolean flag indicating that an incremental update should
        be performed, defaults to True.
        force_write: A boolean flag forcing empty map updates, defaults to False.

Returns:
        An int indicating success of update (0 == good, fail otherwise).


def _add_variable_with_custom_getter(self, name, shape=None, dtype=dtypes.float32, initializer=None, getter=None, overwrite=False, **kwargs_for_getter):
    self._maybe_initialize_trackable()
    with ops.init_scope():
        if context.executing_eagerly():
            checkpoint_initializer = self._preload_simple_restoration(name=name)
        else:
            checkpoint_initializer = None
        if checkpoint_initializer is not None and (not (isinstance(initializer, CheckpointInitialValueCallable) and initializer.restore_uid > checkpoint_initializer.restore_uid)):
            initializer = checkpoint_initializer
    new_variable = getter(name=name, shape=shape, dtype=dtype, initializer=initializer, **kwargs_for_getter)
    if not overwrite or isinstance(new_variable, Trackable):
        return self._track_trackable(new_variable, name=name, overwrite=overwrite)
    else:
        return new_variable

Restore-on-create for a variable be saved with this `Trackable`.

    If the user has requested that this object or another `Trackable` which
    depends on this object be restored from a checkpoint (deferred loading
    before variable object creation), `initializer` may be ignored and the value
    from the checkpoint used instead.

Args:
        name: A name for the variable. Must be unique within this object.
        shape: The shape of the variable.
        dtype: The data type of the variable.
        initializer: The initializer to use. Ignored if there is a deferred
        restoration stored in the Trackable.
        getter: The getter to wrap which actually fetches the variable.
        overwrite: If True, disables unique name and type checks.
        **kwargs_for_getter: Passed to the getter.

Returns:
        The new variable object.

Raises:
        ValueError: If the variable name is not unique.


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

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 _print_extension(self, extension: message.Message) -> None:
    if not fhir_types.is_type_or_profile_of_extension(extension):
        raise ValueError(f'Message of type: {extension.DESCRIPTOR.full_name} is not a FHIR Extension.')
    if self.json_format == _FhirJsonFormat.ANALYTIC:
        self.generator.push(f'"{cast(Any, extension).url.value}"')
    else:
        self._print_message(extension)

Pushes the Extension into the JSON text generator.

    If the _FhirJsonFormat is set to ANALYTIC, this method only prints the url.

Args:
        extension: The Extension to print.


def _maybe_colocate_with(self, value):
    if not self._colocate_with_first_write_call:
        yield
    else:
        if not self._colocate_with:
            self._colocate_with.append(value)
        with ops.colocate_with(self._colocate_with[0]):
            yield

Colocate operations with an internal colocation group or `value`.

Args:
        value: `Tensor`, the tensor to try to colocate with.

Yields:
        Does not yield anything, but the new context is a colocation context.

    If no internal colocation group is set, colocate with `value` and set
    the internal colocation group to be value.


def _expand_variables(input_str, cmake_vars):

    def replace(match):
        if match.group(1) in cmake_vars:
            return cmake_vars[match.group(1)]
        return ''
    return _CMAKE_ATVAR_REGEX.sub(replace, _CMAKE_VAR_REGEX.sub(replace, input_str))

Expands ${VARIABLE}s and @VARIABLE@s in 'input_str', using dictionary 'cmake_vars'.

Args:
        input_str: the string containing ${VARIABLE} or @VARIABLE@ expressions to expand.
        cmake_vars: a dictionary mapping variable names to their values.

Returns:
        The expanded string.


def __init__(self, name=None):
    rr = gen_io_ops.identity_reader_v2(name=name)
    super(IdentityReader, self).__init__(rr, supports_serialize=True)

Create a IdentityReader.

Args:
        name: A name for the operation (optional).


def get_adafactor_schedule(optimizer, initial_lr=0.0):
    return AdafactorSchedule(optimizer, initial_lr)

Get a proxy schedule for [`~optimization.Adafactor`]

Args:
        optimizer ([`~torch.optim.Optimizer`]):
            The optimizer for which to schedule the learning rate.
        initial_lr (`float`, *optional*, defaults to 0.0):
            Initial lr

Returns:
        [`~optimization.Adafactor`] proxy schedule object.


def from_proto(context_def, import_scope=None):
    ret = WhileContext(context_def=context_def, import_scope=import_scope)
    ret.Enter()
    for nested_def in context_def.nested_contexts:
        from_control_flow_context_def(nested_def, import_scope=import_scope)
    ret.Exit()
    return ret

Returns a `WhileContext` object created from `context_def`.

Args:
        context_def: A `WhileContextDef` protocol buffer.
        import_scope: Optional `string`. Name scope to add.

Returns:
        A `WhileContext` Python object.


def append_value_at_field(msg: message.Message, field: Union[descriptor.FieldDescriptor, str], value: Any):
    if isinstance(field, str):
        field = _field_descriptor_for_name(msg, field)
    if not field_is_repeated(field):
        raise ValueError(f'{field.name} is not repeated. Unable to append value.')
    getattr(msg, field.name).append(value)

Appends the provided value to the repeated field.

Args:
        msg: The message whose field to append to.
        field: The FieldDescriptor or name of the field to append to.
        value: The value to append.

Raises:
        ValueError: In the event that the field is not repeated.


def dispatch_command(self, prefix, argv, screen_info=None):
    if not prefix:
        raise ValueError('Prefix is empty')
    resolved_prefix = self._resolve_prefix(prefix)
    if not resolved_prefix:
        raise ValueError('No handler is registered for command prefix "%s"' % prefix)
    handler = self._handlers[resolved_prefix]
    try:
        output = handler(argv, screen_info=screen_info)
    except CommandLineExit as e:
        raise e
    except SystemExit as e:
        lines = ['Syntax error for command: %s' % prefix, 'For help, do "help %s"' % prefix]
        output = RichTextLines(lines)
    except BaseException as e:
        lines = ['Error occurred during handling of command: %s %s:' % (resolved_prefix, ' '.join(argv)), '%s: %s' % (type(e), str(e))]
        lines.append('')
        lines.extend(traceback.format_exc().split('\n'))
        output = RichTextLines(lines)
    if not isinstance(output, RichTextLines) and output is not None:
        raise ValueError('Return value from command handler %s is not None or a RichTextLines instance' % str(handler))
    return output

Handles a command by dispatching it to a registered command handler.

Args:
        prefix: Command prefix, as a str, e.g., "print".
        argv: Command argument vector, excluding the command prefix, represented
        as a list of str, e.g.,
        ["tensor_1"]
        screen_info: A dictionary containing screen info, e.g., {"cols": 100}.

Returns:
        An instance of RichTextLines or None. If any exception is caught during
        the invocation of the command handler, the RichTextLines will wrap the
        error type and message.

Raises:
        ValueError: If
        1) prefix is empty, or
        2) no command handler is registered for the command prefix, or
        3) the handler is found for the prefix, but it fails to return a
        RichTextLines or raise any exception.
        CommandLineExit:
        If the command handler raises this type of exception, this method will
        simply pass it along.


def get_coder_from_spec(coder_spec):
    assert coder_spec is not None
    ignored_wrappers = 'com.google.cloud.dataflow.sdk.util.TimerOrElement$TimerOrElementCoder'
    if coder_spec['@type'] in ignored_wrappers:
        assert len(coder_spec['component_encodings']) == 1
        coder_spec = coder_spec['component_encodings'][0]
        return get_coder_from_spec(coder_spec)
    if coder_spec['@type'] == 'kind:pair':
        assert len(coder_spec['component_encodings']) == 2
        component_coders = [get_coder_from_spec(c) for c in coder_spec['component_encodings']]
        return coders.TupleCoder(component_coders)
    elif coder_spec['@type'] == 'kind:stream':
        assert len(coder_spec['component_encodings']) == 1
        return coders.IterableCoder(get_coder_from_spec(coder_spec['component_encodings'][0]))
    elif coder_spec['@type'] == 'kind:windowed_value':
        assert len(coder_spec['component_encodings']) == 2
        value_coder, window_coder = [get_coder_from_spec(c) for c in coder_spec['component_encodings']]
        return coders.coders.WindowedValueCoder(value_coder, window_coder=window_coder)
    elif coder_spec['@type'] == 'kind:interval_window':
        assert 'component_encodings' not in coder_spec or not coder_spec['component_encodings']
        return coders.coders.IntervalWindowCoder()
    elif coder_spec['@type'] == 'kind:global_window':
        assert 'component_encodings' not in coder_spec or not coder_spec['component_encodings']
        return coders.coders.GlobalWindowCoder()
    elif coder_spec['@type'] == 'kind:varint':
        assert 'component_encodings' not in coder_spec or len(coder_spec['component_encodings'] == 0)
        return coders.coders.VarIntCoder()
    elif coder_spec['@type'] == 'kind:length_prefix':
        assert len(coder_spec['component_encodings']) == 1
        return coders.coders.LengthPrefixCoder(get_coder_from_spec(coder_spec['component_encodings'][0]))
    elif coder_spec['@type'] == 'kind:bytes':
        assert 'component_encodings' not in coder_spec or len(coder_spec['component_encodings'] == 0)
        return coders.BytesCoder()
    return coders.coders.deserialize_coder(coder_spec['@type'].encode('ascii'))

Return a coder instance from a coder spec.

Args:
        coder_spec: A dict where the value of the '@type' key is a pickled instance
        of a Coder instance.

Returns:
        A coder instance (has encode/decode methods).


def _process_v1_graph_mode_tensor(self, op_type, tensor, debug_tensor, tensor_debug_mode):
    if op_type in ('Placeholder', 'PlaceholderWithDefault'):
        self._placeholder_to_debug_tensor[tensor] = debug_tensor
        return tensor
    elif tensor_debug_mode == debug_event_pb2.TensorDebugMode.FULL_TENSOR and op_type != 'Const':
        self._tensor_aliases[debug_tensor.name] = tensor.name
        return debug_tensor
    else:
        with self._symbolic_tensor_counter_lock:
            identity_name = 'tfdbg_identity_%d' % self._symbolic_tensor_counter
        identity = array_ops.identity(tensor, name=identity_name)
        identity.op._add_control_input(debug_tensor.op)
        self._tensor_aliases[identity.name] = tensor.name
        return identity

For V1 graph mode, determine what tensor to output from callback.

Args:
        op_type: Type of the op that outputs the original symbolic tensor.
        tensor: The original output symbolic tensor.
        debug_tensor: The debugger-instrumented tensor.
        tensor_debug_mode: Debug mode used, a tfdbg TensorDebugMode enum.

Returns:
        A symbolic tensor to be returned by the dumping op_callback.


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

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

Args:
        image (`np.ndarray`):
            Image to resize.
        size (`Dict[str, int]`):
            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_generic_distributions(generic_dists, metric_id):
    return sum((get_all_distributions_by_type(dist, metric_id) for dist in generic_dists), [])

Creates flatten list of distributions per its value type.
    A generic distribution is the one which is not processed but saved in
    the most raw version.

Args:
        generic_dists: list of distributions to be saved
        metric_id(uuid): id of the current test run

Returns:
        list of dictionaries made from :class:`DistributionMetric`


def _parallel_part_processors(part_processors: Sequence[PartProcessorWithMatchFn]) -> PartProcessorFn:

    async def part_processor(content: ProcessorPart) -> AsyncIterable[ProcessorPart]:
        output_queue = asyncio.Queue()
        processors = []
        match_fns = []
        passthrough_fallback = False
        passthrough_always = False
        for p in part_processors:
            if p is PASSTHROUGH_FALLBACK:
                passthrough_fallback = True
                continue
            if p is PASSTHROUGH_ALWAYS:
                passthrough_always = True
                continue
            processors.append(_CaptureReservedSubstreams(output_queue, p))
            match_fns.append(p.match)
        parallel_processor = _CaptureReservedSubstreams(output_queue, map_processor.parallel_part_functions(processors, match_fns, with_default_output=passthrough_fallback, with_always_output=passthrough_always))
        content = parallel_processor(content)
        create_task(_enqueue_content(content, output_queue))
        while (part := (await output_queue.get())) is not None:
            yield part
            output_queue.task_done()
    return part_processor

Combine **part processors** in parallel.

    Adds debug and status streams to the output.

    NOTE: Substreams debug and status are yielded immediately instead of passing
        them to the next processor.

Args:
        part_processors: sequence of part processors to compute concurrently.

Returns:
        Part processor that computes the output of the provided sequence of part
        processors concurrently.


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

The set of ops that terminate the gradient computation.

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

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

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

Returns:
        The set of operations.


def _ref(self):
    return self._variable

Returns a reference to this variable.

    You usually do not need to call this method as all ops that need a reference
    to the variable call it automatically.

    Returns is a `Tensor` which holds a reference to the variable.  You can
    assign a new value to the variable by passing the tensor to an assign op.
    See `tf.Variable.value` if you want to get the value of the
    variable.

Returns:
        A `Tensor` that is a reference to the variable.


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

Args:
    hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
    attention_mask (`torch.FloatTensor`, *optional*):
        attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,
        query_sequence_length, key_sequence_length)` if default attention is used.
    output_attentions (`bool`, *optional*):
        Whether or not to return the attentions tensors of all attention layers. See `attentions` under
        returned tensors for more detail.
    use_cache (`bool`, *optional*):
        If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
        (see `past_key_values`).
    past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
    cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
        Indices depicting the position of the input sequence tokens in the sequence
    kwargs (`dict`, *optional*):
        Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
        into the model


def _is_guaranteed_const(tensor):
    if isinstance(tensor, ops.EagerTensor):
        return False

    class Work(object):

        def __init__(self, op: ops.Operation, leaving):
            self.op = op
            self.leaving = leaving
    is_guaranteed_const = lambda op: op.node_def.op == 'GuaranteeConst'
    constants = set([])

    def all_inputs_const(op: ops.Operation):
        return op.inputs and all((inp.op in constants for inp in op.inputs))
    visited = set([])
    stack = [Work(tensor.op, leaving=False)]
    while stack:
        work = stack.pop()
        if work.leaving:
            if all_inputs_const(work.op):
                constants.add(work.op)
            continue
        visited.add(work.op)
        if is_guaranteed_const(work.op):
            constants.add(work.op)
            continue
        stack.append(Work(work.op, leaving=True))
        for inp in work.op.inputs:
            if inp.op not in visited:
                stack.append(Work(inp.op, leaving=False))
    return tensor.op in constants

Determines whether `tensor` is guaranteed to be a constant.

    A tensor is guaranteed to be a constant if either it was produced by
    a `GuaranteeConst` op or if all of its children are guaranteed to be
    constants.

Args:
        tensor: The tensor for which to determine const-ness.

Returns:
        True if `tensor` is guaranteed to be a constant, False otherwise.


def encode(self, table: 'pd.DataFrame', query: Optional[Union[TextInput, PreTokenizedInput, EncodedInput]]=None, add_special_tokens: bool=True, padding: Union[bool, str, PaddingStrategy]=False, truncation: Union[bool, str, TapasTruncationStrategy]=False, max_length: Optional[int]=None, return_tensors: Optional[Union[str, TensorType]]=None, **kwargs) -> List[int]:
    encoded_inputs = self.encode_plus(table, query=query, add_special_tokens=add_special_tokens, padding=padding, truncation=truncation, max_length=max_length, return_tensors=return_tensors, **kwargs)
    return encoded_inputs['input_ids']

Prepare a table and a string for the model. This method does not return token type IDs, attention masks, etc.
    which are necessary for the model to work correctly. Use that method if you want to build your processing on
    your own, otherwise refer to `__call__`.

Args:
        table (`pd.DataFrame`):
            Table containing tabular data. Note that all cell values must be text. Use *.astype(str)* on a Pandas
            dataframe to convert it to string.
        query (`str` or `List[str]`):
            Question related to a table to be encoded.


def num_parameters(self, only_trainable: bool=False) -> int:
    if only_trainable:
        return int(sum((np.prod(w.shape.as_list()) for w in self.trainable_variables)))
    else:
        return self.count_params()

Get the number of (optionally, trainable) parameters in the model.

Args:
        only_trainable (`bool`, *optional*, defaults to `False`):
            Whether or not to return only the number of trainable parameters

Returns:
        `int`: The number of parameters.


def __init__(self, saved_model_dir, saved_model_tags, saved_model_exported_names, experimental_debug_info_func=None):
    """Constructor for TFLiteConverter.

Args:
  saved_model_dir: Directory of the SavedModel.
  saved_model_tags: Set of tags identifying the MetaGraphDef within the
    SavedModel to analyze. All tags in the tag set must be present. (default
    {tf.saved_model.SERVING}).
  saved_model_exported_names: Names to be exported when the saved model
    import path is on.
  experimental_debug_info_func: An experimental function to retrieve the
    graph debug info for a set of nodes from the `graph_def`.

Raises:
  ValueError: Invalid arguments.
"""
    super(TFLiteSavedModelConverter, self).__init__(experimental_debug_info_func)
    self.saved_model_dir = saved_model_dir
    self._saved_model_tags = saved_model_tags
    self._saved_model_exported_names = saved_model_exported_names
    if len(self._saved_model_exported_names) != 1:
        raise ValueError('Only supports a single signature key.')
    signature_key = self._saved_model_exported_names[0]
    result = _freeze_saved_model(self.saved_model_dir, None, None, None, self._saved_model_tags, signature_key)
    self._graph_def = result[0]
    self._input_tensors = result[1]
    self._output_tensors = result[2]
    self._parse_saved_model_args()

Converts the given SavedModel into TensorFlow Lite model.

Attributes:
        saved_model_dir: Directory of the SavedModel.


def _shape_invariant_to_type_spec(self, shape):
    raise NotImplementedError(f'{type(self).__name__}._shape_invariant_to_type_spec')

Returns a TypeSpec given a shape invariant (used by `tf.while_loop`).

Args:
        shape: A `tf.TensorShape` object.  The shape invariant for this
        `CompositeTensor`, or `None` if a default shape invariant should be used
        (based on the value of this `CompositeTensor`).

Returns:
        A nested structure whose values are `tf.TensorShape` objects, specifying
        the shape invariants for the tensors that comprise this `CompositeTensor`.


def reconstruct_non_debug_graph_def(debug_graph_def):
    return DebugGraph(debug_graph_def).non_debug_graph_def

Reconstruct original (non-debugger-decorated) partition GraphDef.

    This method strips the input `tf.compat.v1.GraphDef` of the Copy* and
    Debug*-type nodes inserted by the debugger.

    The reconstructed partition graph is identical to the original (i.e.,
        non-debugger-decorated) partition graph except in the following respects:
        1) The exact names of the runtime-inserted internal nodes may differ.
        These include _Send, _Recv, _HostSend, _HostRecv, _Retval ops.
        2) As a consequence of 1, the nodes that receive input directly from such
        send- and recv-type ops will have different input names.
        3) The parallel_iteration attribute of while-loop Enter ops are set to 1.

Args:
        debug_graph_def: The debugger-decorated `tf.compat.v1.GraphDef`, with the
        debugger-inserted Copy* and Debug* nodes.

Returns:
        The reconstructed `tf.compat.v1.GraphDef` stripped of the debugger-inserted
        nodes.


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

Gets the absolute value of an [`EventSet`][temporian.EventSet]'s
    features.

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

        ```

Returns:
        EventSet with positive valued features.


def patch_retry(testcase, module):
    from mock import Mock
    from mock import patch
    real_retry_with_exponential_backoff = retry.with_exponential_backoff

    def patched_retry_with_exponential_backoff(**kwargs):
        """A patch for retry decorator to use a mock dummy clock and logger."""
        kwargs.update(logger=Mock(), clock=Mock())
        return real_retry_with_exponential_backoff(**kwargs)
    patch.object(retry, 'with_exponential_backoff', side_effect=patched_retry_with_exponential_backoff).start()
    importlib.reload(module)

    def remove_patches():
        patch.stopall()
        importlib.reload(module)
    testcase.addCleanup(remove_patches)

A function to patch retry module to use mock clock and logger.

    Clock and logger that defined in retry decorator will be replaced in test
    in order to skip sleep phase when retry happens.

Args:
        testcase: An instance of unittest.TestCase that calls this function to
        patch retry module.
        module: The module that uses retry and need to be replaced with mock
        clock and logger in test.


def tokenize(self, text):
    output_tokens = []
    for token in whitespace_tokenize(text):
        chars = list(token)
        if len(chars) > self.max_input_chars_per_word:
            output_tokens.append(self.unk_token)
            continue
        is_bad = False
        start = 0
        sub_tokens = []
        while start < len(chars):
            end = len(chars)
            cur_substr = None
            while start < end:
                substr = ''.join(chars[start:end])
                if start > 0:
                    substr = '##' + substr
                if substr in self.vocab:
                    cur_substr = substr
                    break
                end -= 1
            if cur_substr is None:
                is_bad = True
                break
            sub_tokens.append(cur_substr)
            start = end
        if is_bad:
            output_tokens.append(self.unk_token)
        else:
            output_tokens.extend(sub_tokens)
    return output_tokens

Tokenizes a piece of text into its word pieces. This uses a greedy longest-match-first algorithm to perform
    tokenization using the given vocabulary.

    For example, `input = "unaffable"` will return as output `["un", "##aff", "##able"]`.

Args:
        text: A single token or whitespace separated tokens. This should have
            already been passed through *BasicTokenizer*.

Returns:
        A list of wordpiece tokens.


def broadcast_tensor(self, tensor):
    return array_ops.gather(tensor, self.gather_index)

Broadcast from a dense tensor.

    It is assumed that the first axis of the dense tensor is indexed by the
    source shape, and at the end, the first axis of the dense tensor is
    indexed by the destination shape.

Args:
        tensor: a dense tensor.

Returns:
        A dense tensor.


def convert_segmentation_to_rle(segmentation):
    segment_ids = torch.unique(segmentation)
    run_length_encodings = []
    for idx in segment_ids:
        mask = torch.where(segmentation == idx, 1, 0)
        rle = binary_mask_to_rle(mask)
        run_length_encodings.append(rle)
    return run_length_encodings

Converts given segmentation map of shape `(height, width)` to the run-length encoding (RLE) format.

Args:
        segmentation (`torch.Tensor` or `numpy.array`):
            A segmentation map of shape `(height, width)` where each value denotes a segment or class id.

Returns:
        `List[List]`: A list of lists, where each list is the run-length encoding of a segment / class id.


def __init__(self, x_data: types.RealTensor, y_data: types.RealTensor, z_data: types.RealTensor, dtype: tf.DType=None, name: str=None):
    name = name or 'interpolation_2d'
    with tf.name_scope(name):
        self._xdata = tf.convert_to_tensor(x_data, dtype=dtype, name='x_data')
        self._dtype = dtype or self._xdata.dtype
        self._ydata = tf.convert_to_tensor(y_data, dtype=self._dtype, name='y_data')
        self._zdata = tf.convert_to_tensor(z_data, dtype=self._dtype, name='z_data')
        self._name = name
        self._spline_yz = cubic.build_spline(self._ydata, self._zdata, name='spline_y_direction')

Initialize the 2d-interpolation object.

Args:
        x_data: A `Tensor` of real `dtype` and shape
        `batch_shape + [num_x_data_points]`.
        Defines the x-coordinates of the input data. `num_x_data_points` should
        be >= 2. The elements of `x_data` should be in a non-decreasing order.
        y_data: A `Tensor` of the same `dtype` as `x_data` and shape
        `batch_shape + [num_x_data_points, num_y_data_points]`. Defines the
        y-coordinates of the input data. `num_y_data_points` should be >= 2.
        The elements of `y_data` should be in a non-decreasing order along last
        dimension.
        z_data: A `Tensor` of the same shape and `dtype` as `y_data`. Defines the
        z-coordinates of the input data (i.e., the function values).
        dtype: Optional dtype for the input `Tensor`s.
        Default value: `None` which maps to the default dtype inferred by
        TensorFlow.
        name: Python `str` name prefixed to ops created by this class.
        Default value: `None` which is mapped to the default name
        `interpolation_2d`.


def split(cls, tensor, split_dimension, num_devices, input_shape=None):
    if input_shape:
        shape = input_shape
    else:
        shape = tensor.shape.as_list()
    if shape[split_dimension] is not None and shape[split_dimension] < num_devices:
        raise ValueError('Split dimension was smaller than the required number of splits: shape=%r, dimension=%r, num_devices=%r' % (shape, split_dimension, num_devices))
    tile_assignment_dims = [1] * len(shape)
    tile_assignment_dims[split_dimension] = num_devices
    return Sharding(proto=xla_data_pb2.OpSharding(type=xla_data_pb2.OpSharding.OTHER, tile_assignment_dimensions=tile_assignment_dims, tile_assignment_devices=range(num_devices)))

Returns a Sharding that splits a tensor across a dimension.

    This creates a Tiled attribute, similar to tile(), but easier to use for the
    common case of tiling a tensor N ways in one dimension.

Args:
        tensor: A tf.Tensor to split.
        split_dimension: The dimension number to split.
        num_devices: The number of cores to split `tensor` over.
        input_shape: The shape of the original tensor.

Raises:
        ValueError: The tensor to split was smaller in the split dimension than
        the number of devices to split over.


def node_op_type(self, node_name, device_name=None):
    if not self._debug_graphs:
        raise LookupError('Node op types are not loaded from partition graphs yet.')
    device_name = self._infer_device_name(device_name, node_name)
    return self._debug_graphs[device_name].node_op_types[node_name]

Get the op type of given node.

Args:
        node_name: (`str`) name of the node.
        device_name: (`str`) name of the device. If there is only one device or if
        node_name exists on only one device, this argument is optional.

Returns:
        (`str`) op type of the node.

Raises:
        LookupError: If node op types have not been loaded
        from partition graphs yet.


def VisitNamedType(self, node):
    return pytd.ClassType(node.name)

Converts a named type to a class type, to be filled in later.

Args:
        node: The NamedType. This type only has a name.

Returns:
        A ClassType. This ClassType will (temporarily) only have a name.


def __init__(self, dtensor_components: Tuple[tensor.Tensor], global_element_spec: tensor_spec.TensorSpec, layouts: Any):
    [self._iterator_resource_dtensor] = dtensor_components
    self._global_element_spec = global_element_spec
    self._layouts = layouts
    self._layouts_str = nest.map_structure(lambda layout: layout.to_string(), layouts)
    super().__init__(components=dtensor_components, element_spec=global_element_spec)

Initializes a distributed iterator for DTensor datasets.

    This iterator encapsulates tf.data iterators for the underlying devices, and
    treats it as a packed DTensor of iterator resource tensors.

Args:
        dtensor_components: a tuple containing the underlying iterator resources
        packed into a DTensor. This is expected to be a tuple with a single
        element.
        global_element_spec: the underlying dataset's element spec from a global
        view.
        layouts: a structure of DTensor layouts to be applied to the elements
        returned by the underlying iterators. This can be a single layout or
        (possibly nested) tuples or dictionaries of layouts, and the structure
        must match the structure of the iterator elements.


def _calculate_scores(self, query, key):
    q_reshaped = array_ops.expand_dims(query, axis=-2)
    k_reshaped = array_ops.expand_dims(key, axis=-3)
    if self.use_scale:
        scale = self.scale
    else:
        scale = 1.0
    return math_ops.reduce_sum(scale * math_ops.tanh(q_reshaped + k_reshaped), axis=-1)

Calculates attention scores as a nonlinear sum of query and key.

Args:
        query: Query tensor of shape `[batch_size, Tq, dim]`.
        key: Key tensor of shape `[batch_size, Tv, dim]`.

Returns:
        Tensor of shape `[batch_size, Tq, Tv]`.


def visit_Import(self, node):
    for import_alias in node.names:
        if import_alias.name == 'tensorflow.compat.v1' and import_alias.asname == 'tf':
            import_alias.name = 'tensorflow'
    self.generic_visit(node)

Handle visiting an import node in the AST.

Args:
        node: Current Node


def to_yaml(self, **kwargs):
    raise RuntimeError('Method `model.to_yaml()` has been removed due to security risk of arbitrary code execution. Please use `model.to_json()` instead.')

Returns a yaml string containing the network configuration.

    Note: Since TF 2.6, this method is no longer supported and will raise a
    RuntimeError.

    To load a network from a yaml save file, use
    `keras.models.model_from_yaml(yaml_string, custom_objects={})`.

    `custom_objects` should be a dictionary mapping
    the names of custom losses / layers / etc to the corresponding
    functions / classes.

Args:
        **kwargs: Additional keyword arguments
            to be passed to `yaml.dump()`.

Returns:
        A YAML string.

Raises:
        RuntimeError: announces that the method poses a security risk


def apply_gradients(self, grads_and_vars, global_step=None, name=None):
    if distribute_lib.in_cross_replica_context():
        raise ValueError('apply_gradients() must be called in a replica context.')
    if not self._doing_dynamic_loss_scaling():
        return self._optimizer.apply_gradients(grads_and_vars, global_step, name)
    replica_context = distribute_lib.get_replica_context()
    grads_and_vars = tuple(grads_and_vars)
    return replica_context.merge_call(self._distributed_apply, args=(grads_and_vars, global_step, name))

Apply gradients to variables.

    This is the second part of `minimize()`. It returns an `Operation` that
    conditionally applies gradients if all gradient values are finite.
    Otherwise no update is performed (nor is `global_step` incremented).

Args:
        grads_and_vars: List of (gradient, variable) pairs as returned by
        `compute_gradients()`.
        global_step: Optional `Variable` to increment by one after the variables
        have been updated.
        name: Optional name for the returned operation.  Default to the name
        passed to the `Optimizer` constructor.

Returns:
        An `Operation` that conditionally applies the specified gradients. If
        `global_step` was not None, that operation also increments `global_step`.

Raises:
        RuntimeError: If you should use `_distributed_apply()` instead.


def compute_gradient(f, x, delta=None):
    if not isinstance(x, (list, tuple)):
        raise ValueError('`x` must be a list or tuple of values convertible to a Tensor (arguments to `f`), not a %s' % type(x))
    if delta is None:
        delta = 1.0 / 1024
    return _compute_gradient_list(f, x, delta)

Computes the theoretical and numeric Jacobian of `f`.

    With y = f(x), computes the theoretical and numeric Jacobian dy/dx.

Args:
        f: the function.
        x: the arguments for the function as a list or tuple of values convertible
        to a Tensor.
        delta: (optional) perturbation used to compute numeric Jacobian.

Returns:
        A pair of lists, where the first is a list of 2-d numpy arrays representing
        the theoretical Jacobians for each argument, and the second list is the
        numerical ones. Each 2-d array has "y_size" rows
        and "x_size" columns where "x_size" is the number of elements in the
        corresponding argument and "y_size" is the number of elements in f(x).

Raises:
        ValueError: If result is empty but the gradient is nonzero.
        ValueError: If x is not list, but any other type.

Example:
    >>> @tf.function
    ... def test_func(x):
    ...   return x*x
    ...
    >>>
    >>> class MyTest(tf.test.TestCase):
    ...
    ...   def test_gradient_of_test_func(self):
    ...     theoretical, numerical = tf.test.compute_gradient(test_func, [1.0])
    ...     # ((array([[2.]], dtype=float32),),
    ...     #  (array([[2.000004]], dtype=float32),))
    ...     self.assertAllClose(theoretical, numerical)


def _get_index_proposal_probs(self, x):
    x_float = tf.cast(x, tf.float32)
    with tf.GradientTape(watch_accessed_variables=False) as tape:
        tape.watch(x_float)
        current_energy = tf.squeeze(self._energy(tf.expand_dims(x_float, 0)))
    f_grad = -1.0 * tape.gradient(current_energy, x_float)
    approx_energy_diff = -(2 * tf.cast(x, tf.float32) - 1) * f_grad
    return tf.nn.softmax(approx_energy_diff / 2.0)

Returns the value of equation 6 of the paper.

    Given current state x, the locally-informed proposal evaluates f(x') - f(x)
    for each x' in a Hamming ball of radius 1 around x, written H(x).  Let
    d(x) = [f(x') - f(x) for x' in H(x)].  The locally-informed conditional
    distribution is then q(i | x) = C * e^(d(x) / T) for some temperature T,
    where C is the normalization constant (T=2 is optimal under a standard
    criterion).

    The insight of the paper is that d(x) can be approximated using a Taylor
    series, yielding d(x) ~ -(2x - 1) * df(x)/dx.  This function computes the
    locally-informed conditional probabilities q(i | x) using the approximation.

    Note that in terms of the unnormalized log probability of the state, f(x),
    used in the paper, the energy of the state is: E(x) = -f(x).  So we have
    d(x) ~ (2x - 1) * dE(x)/dx

Args:
        x: 1D tensor which is the current state of the chain.

Returns:
        The conditional probabilities q(i | x) given in equation 6.


def join(self, basepath, *paths):
    if not basepath.startswith(GCSFileSystem.GCS_PREFIX):
        raise ValueError('Basepath %r must be GCS path.' % basepath)
    path = basepath
    for p in paths:
        path = path.rstrip('/') + '/' + p.lstrip('/')
    return path

Join two or more pathname components for the filesystem

Args:
        basepath: string path of the first component of the path
        paths: path components to be added

    Returns: full path after combining all the passed components


Test function

Args:
        x: The first input

        y: The second input. This is a longer description
        that spans multiple lines with indentation and stuff.

Returns:
        God knows what


def categorical_crossentropy(y_true, y_pred, from_logits=False, label_smoothing=0, axis=-1):
    y_pred = tensor_conversion.convert_to_tensor_v2_with_dispatch(y_pred)
    y_true = math_ops.cast(y_true, y_pred.dtype)
    label_smoothing = tensor_conversion.convert_to_tensor_v2_with_dispatch(label_smoothing, dtype=backend.floatx())

    def _smooth_labels():
        num_classes = math_ops.cast(array_ops.shape(y_true)[-1], y_pred.dtype)
        return y_true * (1.0 - label_smoothing) + label_smoothing / num_classes
    y_true = smart_cond.smart_cond(label_smoothing, _smooth_labels, lambda: y_true)
    return backend.categorical_crossentropy(y_true, y_pred, from_logits=from_logits, axis=axis)

Computes the categorical crossentropy loss.

    Standalone usage:

    >>> y_true = [[0, 1, 0], [0, 0, 1]]
    >>> y_pred = [[0.05, 0.95, 0], [0.1, 0.8, 0.1]]
    >>> loss = tf.keras.losses.categorical_crossentropy(y_true, y_pred)
    >>> assert loss.shape == (2,)
    >>> loss.numpy()
    array([0.0513, 2.303], dtype=float32)

Args:
        y_true: Tensor of one-hot true targets.
        y_pred: Tensor of predicted targets.
        from_logits: Whether `y_pred` is expected to be a logits tensor. By default,
        we assume that `y_pred` encodes a probability distribution.
        label_smoothing: Float in [0, 1]. If > `0` then smooth the labels. For
        example, if `0.1`, use `0.1 / num_classes` for non-target labels
        and `0.9 + 0.1 / num_classes` for target labels.
        axis: Defaults to -1. The dimension along which the entropy is
        computed.

Returns:
        Categorical crossentropy loss value.


def request(self, batch: Sequence[ExampleT], model: ModelT, inference_args: Optional[dict[str, Any]]=None) -> Iterable[PredictionT]:
    raise NotImplementedError(type(self))

Makes a request to a remote inference service and returns the response.
    Should raise an exception of some kind if there is an error to enable the
    retry and client-side throttling logic to work. Returns an iterable of the
    desired prediction type. This method should return the values directly, as
    handling return values as a generator can prevent the retry logic from
    functioning correctly.

Args:
        batch: A sequence of examples or features.
        model: The model used to make inferences.
        inference_args: Extra arguments for models whose inference call requires
        extra parameters.

Returns:
        An Iterable of Predictions.


def build_model_flags(change_concat_input_ranges=False, allow_nonexistent_arrays=False, saved_model_dir=None, saved_model_version=0, saved_model_tags=None, saved_model_exported_names=None, **_):
    model_flags = _model_flags_pb2.ModelFlags()
    model_flags.change_concat_input_ranges = change_concat_input_ranges
    model_flags.allow_nonexistent_arrays = allow_nonexistent_arrays
    if saved_model_dir:
        model_flags.saved_model_dir = saved_model_dir
    model_flags.saved_model_version = saved_model_version
    if saved_model_tags:
        model_flags.saved_model_tags.extend(saved_model_tags)
    if saved_model_exported_names:
        model_flags.saved_model_exported_names.extend(saved_model_exported_names)
    return model_flags

Builds the model flags object from params.

Args:
        change_concat_input_ranges: Boolean to change behavior of min/max ranges for
        inputs and outputs of the concat operator for quantized models. Changes
        the ranges of concat operator overlap when true. (default False)
        allow_nonexistent_arrays: Allow specifying array names that don't exist or
        are unused in the final graph. (default False)
        saved_model_dir: Filepath of the saved model to be converted. This value
        will be non-empty only when the saved model import path will be used.
        Otherwises, the graph def-based conversion will be processed.
        saved_model_version: SavedModel file format version of The saved model file
        to be converted. This value will be set only when the SavedModel import
        path will be used.
        saved_model_tags: Set of string saved model tags, formatted in the
        comma-separated value. This value will be set only when the SavedModel
        import path will be used.
        saved_model_exported_names: Names to be exported (default: export all) when
        the saved model import path is on. This value will be set only when the
        SavedModel import path will be used.

Returns:
        model_flags: protocol buffer describing the model.


def forward(self, data: torch.Tensor, observed_indicator: Optional[torch.Tensor]=None) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    scale = torch.ones_like(data, requires_grad=False).mean(dim=self.dim, keepdim=self.keepdim)
    loc = torch.zeros_like(data, requires_grad=False).mean(dim=self.dim, keepdim=self.keepdim)
    return (data, loc, scale)

Parameters:
    data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`):
        input for Batch norm calculation

Returns:
    tuple of `torch.Tensor` of shapes
        (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`,
        `(batch_size, 1, num_input_channels)`)


def convert_debug_info_func(saved_debug_info):

    def f(original_nodes):
        """Function to create `GraphDebugInfo` for the given `original_nodes`."""
        del original_nodes
        return saved_debug_info
    return f

Returns a method to retrieve the `GraphDebugInfo` from the original graph.

Args:
        saved_debug_info: The `GraphDebugInfo` containing all the debug info.

Returns:
        A function which retrieves the stack traces from the original graph and
        converts them to a `GraphDebugInfo` for a given set of nodes.


def update_candidate_strategy(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, num_matches: int):
    return

Updates the candidate generation strategy based on the outcomes.

Args:
        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
            Indices of input sequence tokens in the vocabulary. [What are input IDs?](../glossary#input-ids)
        scores (`torch.FloatTensor` of shape `(batch_size, candidate_length, config.vocab_size)`):
            Prediction scores of a language modeling head. These can be logits for each vocabulary when not using
            beam search or log softmax for each vocabulary token when using beam search
        num_matches (`int`):
            The number of matches between the candidate sequences and the model predictions.


def __init__(self, min_length: int, eos_token_id: Union[int, List[int], torch.Tensor], device: str='cpu'):
    if not isinstance(min_length, int) or min_length < 0:
        raise ValueError(f'`min_length` has to be a non-negative integer, but is {min_length}')
    if not isinstance(eos_token_id, torch.Tensor):
        if isinstance(eos_token_id, int):
            eos_token_id = [eos_token_id]
        eos_token_id = torch.tensor(eos_token_id, device=device)
    self.min_length = min_length
    self.eos_token_id = eos_token_id

[`LogitsProcessor`] enforcing a min-length by setting EOS probability to 0. Note that, for decoder-only models
    like most LLMs, the length includes the prompt.

Args:
        min_length (`int`):
            The minimum length below which the score of `eos_token_id` is set to `-float("Inf")`.
        eos_token_id (`Union[int, List[int], torch.Tensor]`):
            The id(s) of the *end-of-sequence* token.
        device (`str`, *optional*, defaults to `"cpu"`):
            The device to allocate the tensors.

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

    >>> tokenizer = AutoTokenizer.from_pretrained("bigscience/bloomz-560m")
    >>> model = AutoModelForCausalLM.from_pretrained("bigscience/bloomz-560m")

    >>> inputs = tokenizer("A number:", return_tensors="pt")
    >>> gen_out = model.generate(**inputs)
    >>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0])
    A number: one

    >>> # setting `min_length` to a value smaller than the uncontrolled output length has no impact
    >>> gen_out = model.generate(**inputs, min_length=3)
    >>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0])
    A number: one

    >>> # setting a larger `min_length` will force the model to generate beyond its natural ending point, which is not
    >>> # necessarily incorrect
    >>> gen_out = model.generate(**inputs, min_length=10)
    >>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0])
    A number: one thousand, nine hundred and ninety-four
    ```


def to_diff_dict(self) -> dict[str, Any]:
    config_dict = self.to_dict()
    default_config_dict = PretrainedConfig().to_dict()
    class_config_dict = self.__class__().to_dict() if not self.has_no_defaults_at_init else {}
    serializable_config_dict = {}
    for key, value in config_dict.items():
        if isinstance(getattr(self, key, None), PretrainedConfig) and key in class_config_dict and isinstance(class_config_dict[key], dict) or key in self.sub_configs:
            diff = recursive_diff_dict(value, default_config_dict, config_obj=getattr(self, key, None))
            if 'model_type' in value:
                diff['model_type'] = value['model_type']
            serializable_config_dict[key] = diff
        elif key not in default_config_dict or key == 'transformers_version' or key == 'vocab_file' or (value != default_config_dict[key]) or (key in default_config_dict and value != class_config_dict.get(key, value)):
            serializable_config_dict[key] = value
    self._remove_keys_not_serialized(serializable_config_dict)
    if '_name_or_path' in serializable_config_dict:
        del serializable_config_dict['_name_or_path']
    if hasattr(self, 'quantization_config'):
        serializable_config_dict['quantization_config'] = self.quantization_config.to_dict() if not isinstance(self.quantization_config, dict) else self.quantization_config
    self.dict_torch_dtype_to_str(serializable_config_dict)
    return serializable_config_dict

Removes all attributes from the configuration that correspond to the default config attributes for
    better readability, while always retaining the `config` attribute from the class. Serializes to a
    Python dictionary.

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


def _format_neighbors(self, neighbor_type, non_ctrls, ctrls):
    lines = []
    font_attr_segs = {}
    lines.append('')
    lines.append('  %d %s(s) + %d control %s(s):' % (len(non_ctrls), neighbor_type, len(ctrls), neighbor_type))
    lines.append('    %d %s(s):' % (len(non_ctrls), neighbor_type))
    for non_ctrl in non_ctrls:
        line = '      [%s] %s' % (self._debug_dump.node_op_type(non_ctrl), non_ctrl)
        lines.append(line)
        font_attr_segs[len(lines) - 1] = [(len(line) - len(non_ctrl), len(line), debugger_cli_common.MenuItem(None, 'ni -a -d -t %s' % non_ctrl))]
    if ctrls:
        lines.append('')
        lines.append('    %d control %s(s):' % (len(ctrls), neighbor_type))
        for ctrl in ctrls:
            line = '      [%s] %s' % (self._debug_dump.node_op_type(ctrl), ctrl)
            lines.append(line)
            font_attr_segs[len(lines) - 1] = [(len(line) - len(ctrl), len(line), debugger_cli_common.MenuItem(None, 'ni -a -d -t %s' % ctrl))]
    return debugger_cli_common.RichTextLines(lines, font_attr_segs=font_attr_segs)

List neighbors (inputs or recipients) of a node.

Args:
        neighbor_type: ("input" | "recipient")
        non_ctrls: Non-control neighbor node names, as a list of str.
        ctrls: Control neighbor node names, as a list of str.

Returns:
        A RichTextLines object.


def for_stmt(iter_, extra_test, body, get_state, set_state, symbol_names, opts):
    try:
        for_fn = for_loop_registry.lookup(iter_)
    except LookupError:
        for_fn = _py_for_stmt
        if tensor_util.is_tf_type(iter_):
            if tensors.is_range_tensor(iter_):
                for_fn = _tf_range_for_stmt
            else:
                for_fn = _known_len_tf_for_stmt
        elif isinstance(iter_, distribute.Iterator):
            for_fn = _tf_iterator_for_stmt
        elif isinstance(iter_, distribute.Iterable):
            for_fn = _tf_distributed_iterable_for_stmt
    for_fn(iter_, extra_test, body, get_state, set_state, symbol_names, opts)

Functional form of a for statement.

    The loop operates on a state, which includes all symbols that are
    variant across loop iterations, excluding the variables local to the loop.

    For example, given the loop below that calculates the geometric and
    arithmetic means or some numbers:

    ```
        geo_mean = 1
        arith_mean = 0
        for i in range(n):
        a = numbers[i]
        geo_mean *= a
        arith_mean += a
    ```

    The state is represented by the variables named geo_mean and arith_mean. The
    `extra_test`, `body`, `get_state` and `set_state` functions must bind to the
    original `geo_mean` and `arith_mean` symbols, using `nonlocal`.

    The inputs and outputs of the callables representing the loop blocks are not
    explicit - instead, these functions must use nonlocal/global for side effects.
    The inputs and outputs are instead controlled by the set_state/get_state
    functions.

Args:
        iter_: The entity being iterated over.
        extra_test: Callable with boolean return type. An additional loop condition.
        body: Callable representing the actual loop body.
        get_state: Additional callable which can capture additional state (such as
        the values of composite symbols). This is only useful when staging the
        loop.
        set_state: Additional callable which save values captured by get_state back
        into the Python environment. This is only useful when staging the loop.
        symbol_names: Tuple containing names of the loop variables returned by
        get_state.
        opts: Optional dict of extra loop parameters.


def apply(self, inputs, *args, **kwargs):
    warnings.warn('`layer.apply` is deprecated and will be removed in a future version. Please use `layer.__call__` method instead.')
    return self.__call__(inputs, *args, **kwargs)

Deprecated, do NOT use!

    This is an alias of `self.__call__`.

Args:
        inputs: Input tensor(s).
        *args: additional positional arguments to be passed to `self.call`.
        **kwargs: additional keyword arguments to be passed to `self.call`.

Returns:
        Output tensor(s).


def from_dim_sizes(dim_sizes):
    with ops.name_scope(None, 'RaggedTensorDynamicShapeFromDimensionSizes', [dim_sizes]):
        dim_sizes = tuple((ops.convert_to_tensor(size, preferred_dtype=dtypes.int64, name='dim_sizes') for size in dim_sizes))
        inner_split = 0
        for dim, dim_size in enumerate(dim_sizes):
            if dim_size.shape.ndims == 1:
                inner_split = dim + 1
            elif dim_size.shape.ndims != 0:
                raise ValueError('Each dim_size must be a scalar or a vector')
        return RaggedTensorDynamicShape(dim_sizes[:inner_split], dim_sizes[inner_split:])

Constructs a ragged shape from a list of dimension sizes.

    This list contains a single tensor for each dimension, where the tensor
    is a scalar if the dimension is uniform, or a vector if the dimension is
    ragged.

Args:
        dim_sizes: List of int32 or int64 scalars or vectors.

Returns:
        A RaggedTensorDynamicShape.


def forward(self, pixel_values: torch.FloatTensor, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> BackboneOutput:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    outputs = self.encoder(pixel_values=pixel_values, output_attentions=output_attentions, output_hidden_states=True, return_dict=return_dict)
    hidden_states = outputs.hidden_states
    feature_maps = ()
    for idx, stage in enumerate(self.stage_names):
        if stage in self.out_features:
            feature_maps += (hidden_states[idx],)
    if not return_dict:
        output = (feature_maps,)
        if output_hidden_states:
            output += (outputs.hidden_states,)
        return output
    return BackboneOutput(feature_maps=feature_maps, hidden_states=outputs.hidden_states if output_hidden_states else None, attentions=None)

Examples:

    ```python
    >>> from transformers import AutoImageProcessor, AutoBackbone
    >>> import torch
    >>> from PIL import Image
    >>> import requests

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

    >>> processor = AutoImageProcessor.from_pretrained("OpenGVLab/pvt_v2_b0")
    >>> model = AutoBackbone.from_pretrained(
    ...     "OpenGVLab/pvt_v2_b0", out_features=["stage1", "stage2", "stage3", "stage4"]
    ... )

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

    >>> outputs = model(**inputs)
    >>> feature_maps = outputs.feature_maps
    >>> list(feature_maps[-1].shape)
    [1, 256, 7, 7]
    ```


def fixed_field_for_type_code(type_code: str) -> str:
    url_type_code = _TYPE_CODE_URI_RE.search(type_code)
    if url_type_code is not None:
        type_code = url_type_code[1]
    fixed_field = stringcase.snakecase(type_code[:1].lower() + type_code[1:])
    return 'string_value' if fixed_field == 'string' else fixed_field

Retrieves the `ElementDefinition.fixed.choice` oneof field for `type_code`.

Args:
        type_code: The FHIR type code to look up. Could be a value like 'boolean' or
        URL like 'http://hl7.org/fhirpath/System.Boolean'

Returns:
        The attribute corresponding to this type code on the
        ElementDefinition.FixedX.choice oneof.


def _test_end(self, result, e):
    if self.begin_time is not None:
        self.end_time = utils.get_current_epoch_time()
    self.result = result
    if e:
        self.termination_signal = ExceptionRecord(e)

Marks the end of the test logic.

Args:
        result: One of the TEST_RESULT enums in TestResultEnums.
        e: A test termination signal (usually an exception object). It can
        be any exception instance or of any subclass of
        mobly.signals.TestSignal.


def configTestMesh(device_type_mesh_map: typing.Dict[typing.Text, layout_lib.Mesh]) -> layout_lib.Mesh:
    reset_context()

    def get_mesh(device_type):
        mesh = device_type_mesh_map.get(device_type, None)
        if mesh is None:
            raise ValueError('Requires a %s mesh to run test on %s.' % (device_type, device_type))
        return mesh
    mesh = None
    if is_tpu_present():
        mesh = get_mesh('TPU')
        reset_context()
        accelerator_util.initialize_accelerator_system('TPU')
    elif tf_config.list_physical_devices('GPU'):
        mesh = get_mesh('GPU')
        reset_logical_devices('GPU', np.prod(mesh.shape()))
        accelerator_util.initialize_accelerator_system('GPU')
    else:
        mesh = get_mesh('CPU')
        reset_logical_devices('CPU', np.prod(mesh.shape()))
        accelerator_util.initialize_accelerator_system('CPU')
    return mesh

Configs corresponding mesh given test context.

    If runs on a CPU mesh, set virtual device on CPU.
    If runs on a GPU mesh, sets virtual device on GPU with proper memory limits.
    if runs on a TPU mesh, initializes TPU system.

Args:
        device_type_mesh_map: A dictionary containing device_type -> mesh mapping.

Returns:
        A properly configured mesh for use in test.


def _process_parameters_section(func_documentation, sig, func, class_name, model_name_lowercase, parent_class, indent_level):
    docstring = set_min_indent('Args:\n', indent_level + 4)
    undocumented_parameters = []
    documented_params = {}
    documented_kwargs = {}
    if func_documentation is not None:
        documented_params, func_documentation = parse_docstring(func_documentation)
        if model_name_lowercase is not None:
            documented_params = format_args_docstring(documented_params, model_name_lowercase)
    param_docstring, missing_args = _process_regular_parameters(sig, func, class_name, documented_params, indent_level, undocumented_parameters)
    docstring += param_docstring
    kwargs_docstring = _process_kwargs_parameters(sig, func, parent_class, model_name_lowercase, documented_kwargs, indent_level, undocumented_parameters)
    docstring += kwargs_docstring
    if len(undocumented_parameters) > 0:
        print('\n'.join(undocumented_parameters))
    return docstring

Process the parameters section of the docstring.

Args:
        func_documentation (`str`): Existing function documentation (manually specified in the docstring)
        sig (`inspect.Signature`): Function signature
        func (`function`): Function the parameters belong to
        class_name (`str`): Name of the class the function belongs to
        model_name_lowercase (`str`): Lowercase model name
        parent_class (`class`): Parent class of the function (if any)
        indent_level (`int`): Indentation level


def check_collective_ops_peer_health(self, task, timeout_in_ms):
    self.ensure_initialized()
    pywrap_tfe.TFE_CollectiveOpsCheckPeerHealth(self._handle, task, timeout_in_ms)

Check collective peer health.

    This probes each task to see if they're still alive. Note that restarted
    tasks are considered a different one, and they're considered not healthy.

    This should only be used in multi client multi worker training.

Args:
        task: a task string, must be in the format of /job:xxx/replica:0/task:N.
        timeout_in_ms: an integer, the timeout. If zero, there's no timeout.

Raises:
        tf.errors.UnavailableError: when a peer is down.
        tf.errors.FailedPreconditionError: when a peer is a different one from the
        one this task has talked to, e.g. the peer has restarted.
        tf.errors.InvalidArgumentError: when the task string is invalid.


def parallel(processor_list: Sequence[PartProcessor]) -> PartProcessor:
    if not processor_list:
        raise ValueError('processor_list is empty')
    return _ParallelPartProcessor(processor_list)

Create a sequence of part processors to be run in parallel.

Args:
        processor_list: list of part processors.

Returns:
        A processor consisting of the parallel run of all the processors in the
        list. The execution is sequential from the first processor to the last but
        parts are processed concurrently overall.


def _html_tree_view(self, *, view: 'HtmlTreeView', name: Optional[str]=None, parent: Any=None, root_path: Optional[KeyPath]=None, **kwargs) -> Html:
    return view.render(self, name=name, parent=parent, root_path=root_path, **kwargs)

Returns the topmost HTML representation of this extension.

Args:
        view: The view to render the object.
        name: The name of the object.
        parent: The parent of the object.
        root_path: The key path of the object relative to the root.
        **kwargs: kwargs to pass to the view. See `_html_tree_view_config` for
        the builtin arguments.

Returns:
        The rendered HTML.


def _get_kind_name(item):
    if isinstance(item, (str, bytes)):
        kind = 'bytes_list'
    elif isinstance(item, int):
        kind = 'int64_list'
    elif isinstance(item, float):
        kind = 'float_list'
    elif isinstance(item, Any):
        kind = 'any_list'
    else:
        kind = 'node_list'
    return kind

Returns the kind name in CollectionDef.

Args:
        item: A data item.

Returns:
        The string representation of the kind in CollectionDef.


def restore_op(self, filename_tensor, saveable, preferred_shard):
    tensors = []
    for spec in saveable.specs:
        tensors.append(io_ops.restore_v2(filename_tensor, [spec.name], [spec.slice_spec], [spec.dtype])[0])
    return tensors

Create ops to restore 'saveable'.

    This is intended to be overridden by subclasses that want to generate
    different Ops.

Args:
        filename_tensor: String Tensor.
        saveable: A BaseSaverBuilder.SaveableObject object.
        preferred_shard: Int.  Shard to open first when loading a sharded file.

Returns:
        A list of Tensors resulting from reading 'saveable' from
        'filename'.


def _pyval_empty_list_depth(pyval):
    if isinstance(pyval, list):
        if not pyval:
            return 1
        depths = [_pyval_empty_list_depth(v) for v in pyval]
        if any((depth is None for depth in depths)):
            return None
        else:
            return max(depths) + 1
    else:
        return None

Find the max depth for nested empty lists.

Args:
        pyval: A nested python list.

Returns:
        The maximum depth of empty lists in `pyval`, or None if `pyval` contains
        anything other than nested empty lists.


def __init__(self, elements: Iterable[Any]):
    self._elements = elements

Asserts that the input contains exactly the elements provided.

    This is primarily used for testing; it will cause the entire pipeline to
    fail if the input to this transform is not exactly the set of `elements`
    given in the config parameter.

    As with Create, YAML/JSON-style mappings are interpreted as Beam rows,
    e.g.::

        type: AssertEqual
        input: SomeTransform
        config:
        elements:
            - {a: 0, b: "foo"}
            - {a: 1, b: "bar"}

    would ensure that `SomeTransform` produced exactly two elements with values
    `(a=0, b="foo")` and `(a=1, b="bar")` respectively.

Args:
        elements: The set of elements that should belong to the PCollection.
            YAML/JSON-style mappings will be interpreted as Beam rows.


def _grappler_enabled_session_config():
    rewriter_config = rewriter_config_pb2.RewriterConfig(disable_model_pruning=False, arithmetic_optimization=rewriter_config_pb2.RewriterConfig.ON)
    graph_options = config_pb2.GraphOptions(rewrite_options=rewriter_config)
    return config_pb2.ConfigProto(graph_options=graph_options)

Constructs a Session config proto that explicitly enables Grappler.

Returns:
        A config proto that obtains extra safety for the unit tests in this
        file by ensuring that the relevant Grappler rewrites are always enabled.


def view_as_real(x):
    if any_symbolic_tensors((x,)):
        return ViewAsReal().symbolic_call(x)
    x = backend.convert_to_tensor(x)
    real_part = backend.numpy.real(x)
    imag_part = backend.numpy.imag(x)
    return backend.numpy.stack((real_part, imag_part), axis=-1)

Converts a complex tensor to a real tensor with shape `(..., 2)`,
    where the last dimension represents the real and imaginary components.

Args:
        x: A complex tensor.

Returns:
        A real tensor where the last dimension contains the
        real and imaginary parts.

Example:
    ```
    >>> import numpy as np
    >>> from keras import ops

    >>> complex_tensor = np.array([1 + 2j, 3 + 4j])
    >>> real = ops.view_as_real(complex_tensor)
    >>> real
    array([[1., 2.],
        [3., 4.]])
    ```


def __init__(self, file_pattern, batch_size=1, buffer_size=1, parallelism=1, shift_ratio=0, seed=0, name=None, batches=None, compression_type=None):
    self._batch_size = batch_size
    if batches is not None:
        self._batch_size *= batches
    self._batches = batches
    self._file_pattern = file_pattern
    self._buffer_size = buffer_size
    self._parallelism = parallelism
    self._shift_ratio = shift_ratio
    self._seed = seed
    self._name = name
    self._compression_type = python_io.TFRecordCompressionType.NONE
    if compression_type is not None:
        self._compression_type = compression_type

Constructs a RecordInput Op.

Args:
        file_pattern: File path to the dataset, possibly containing wildcards.
        All matching files will be iterated over each epoch.
        batch_size: How many records to return at a time.
        buffer_size: The maximum number of records the buffer will contain.
        parallelism: How many reader threads to use for reading from files.
        shift_ratio: What percentage of the total number files to move the start
        file forward by each epoch.
        seed: Specify the random number seed used by generator that randomizes
        records.
        name: Optional name for the operation.
        batches: None by default, creating a single batch op. Otherwise specifies
        how many batches to create, which are returned as a list when
        `get_yield_op()` is called. An example use case is to split processing
        between devices on one computer.
        compression_type: The type of compression for the file. Currently ZLIB and
        GZIP are supported. Defaults to none.

Raises:
        ValueError: If one of the arguments is invalid.


def mean_absolute_error(y_true, y_pred):
    y_pred = tensor_conversion.convert_to_tensor_v2_with_dispatch(y_pred)
    y_true = math_ops.cast(y_true, y_pred.dtype)
    return backend.mean(math_ops.abs(y_pred - y_true), axis=-1)

Computes the mean absolute error between labels and predictions.

    `loss = mean(abs(y_true - y_pred), axis=-1)`

    Standalone usage:

    >>> y_true = np.random.randint(0, 2, size=(2, 3))
    >>> y_pred = np.random.random(size=(2, 3))
    >>> loss = tf.keras.losses.mean_absolute_error(y_true, y_pred)
    >>> assert loss.shape == (2,)
    >>> assert np.array_equal(
    ...     loss.numpy(), np.mean(np.abs(y_true - y_pred), axis=-1))

Args:
        y_true: Ground truth values. shape = `[batch_size, d0, .. dN]`.
        y_pred: The predicted values. shape = `[batch_size, d0, .. dN]`.

Returns:
        Mean absolute error values. shape = `[batch_size, d0, .. dN-1]`.


def forward(self, pixel_values: Tensor, task_inputs: Tensor, text_inputs: Optional[Tensor]=None, pixel_mask: Optional[Tensor]=None, output_hidden_states: Optional[bool]=None, output_attentions: Optional[bool]=None, return_dict: Optional[bool]=None) -> OneFormerModelOutput:
    if pixel_values is None:
        raise ValueError('You have to specify pixel_values')
    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
    batch_size, _, height, width = pixel_values.shape
    if pixel_mask is None:
        pixel_mask = torch.ones((batch_size, height, width), device=pixel_values.device)
    pixel_level_module_output = self.pixel_level_module(pixel_values, output_hidden_states)
    multi_scale_features = pixel_level_module_output.decoder_features
    mask_features = pixel_level_module_output.decoder_last_feature
    task_token = self.task_encoder(task_inputs.to(self.dtype))
    if self.is_training:
        text_queries = self.text_mapper(text_inputs)
    else:
        text_queries = None
    transformer_module_output = self.transformer_module(multi_scale_features=multi_scale_features, mask_features=mask_features, task_token=task_token, output_attentions=output_attentions)
    queries = transformer_module_output.object_queries
    encoder_hidden_states = None
    pixel_decoder_hidden_states = None
    transformer_decoder_hidden_states = None
    if output_hidden_states:
        encoder_hidden_states = pixel_level_module_output.encoder_features
        pixel_decoder_hidden_states = (pixel_level_module_output.decoder_last_feature,)
        for f in pixel_level_module_output.decoder_features:
            pixel_decoder_hidden_states += (f,)
        transformer_decoder_hidden_states = transformer_module_output.auxiliary_predictions
    output = OneFormerModelOutput(encoder_hidden_states=encoder_hidden_states, pixel_decoder_hidden_states=pixel_decoder_hidden_states, transformer_decoder_hidden_states=transformer_decoder_hidden_states, transformer_decoder_object_queries=queries, transformer_decoder_contrastive_queries=transformer_module_output.contrastive_logits, transformer_decoder_mask_predictions=transformer_module_output.prediction_masks, transformer_decoder_class_predictions=transformer_module_output.prediction_class, transformer_decoder_auxiliary_predictions=transformer_module_output.auxiliary_predictions, text_queries=text_queries, task_token=task_token, attentions=transformer_module_output.attentions)
    if not return_dict:
        output = tuple((v for v in output.values()))
    return output

task_inputs (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):
    Task inputs. Task inputs can be obtained using [`AutoImageProcessor`]. See [`OneFormerProcessor.__call__`]
    for details.
    text_inputs (`List[torch.Tensor]`, *optional*):
    Tensor fof shape `(num_queries, sequence_length)` to be fed to a model

Example:
    ```python
    >>> import torch
    >>> from PIL import Image
    >>> import requests
    >>> from transformers import OneFormerProcessor, OneFormerModel

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

    >>> # load processor for preprocessing the inputs
    >>> processor = OneFormerProcessor.from_pretrained("shi-labs/oneformer_ade20k_swin_tiny")
    >>> model = OneFormerModel.from_pretrained("shi-labs/oneformer_ade20k_swin_tiny")
    >>> inputs = processor(image, ["semantic"], return_tensors="pt")

    >>> with torch.no_grad():
    ...     outputs = model(**inputs)

    >>> mask_predictions = outputs.transformer_decoder_mask_predictions
    >>> class_predictions = outputs.transformer_decoder_class_predictions

    >>> f"👉 Mask Predictions Shape: {list(mask_predictions.shape)}, Class Predictions Shape: {list(class_predictions.shape)}"
    '👉 Mask Predictions Shape: [1, 150, 128, 171], Class Predictions Shape: [1, 150, 151]'
    ```