def unbatch(self, spec):
    raise NotImplementedError(f'{type(self).__name__}.unbatch')

Returns the TypeSpec for a single unbatched element in `spec`.

Args:
        spec: The `TypeSpec` for a batch of values.

Returns:
        A `TypeSpec` for an individual value.


def forward(self, hidden_states):
    hidden_states = hidden_states.view(self.num_experts, -1, self.hidden_size)
    num_tokens = None
    next_states = torch.empty_like(hidden_states)
    for i in range(self.num_experts):
        expert_hidden = hidden_states[i]
        expert_hidden_reshaped = expert_hidden.reshape(-1, self.hidden_size)
        expert_quantized, expert_scale = torch.ops.fbgemm.quantize_fp8_per_row(expert_hidden_reshaped, num_tokens, self.input_scale_ub)
        sharded_expert_dim = self.gate_up_proj.shape[-1] // 2
        gate_up_proj_scale_float32 = self.gate_up_proj_scale.to(torch.float32)
        gate = torch.ops.fbgemm.f8f8bf16_rowwise(expert_quantized, self.gate_up_proj[i].transpose(0, 1)[:sharded_expert_dim].contiguous(), expert_scale, gate_up_proj_scale_float32[i][0][:sharded_expert_dim].view(-1, 1).contiguous(), use_fast_accum=True)
        up = torch.ops.fbgemm.f8f8bf16_rowwise(expert_quantized, self.gate_up_proj[i].transpose(0, 1)[sharded_expert_dim:].contiguous(), expert_scale, gate_up_proj_scale_float32[i][0][sharded_expert_dim:].view(-1, 1).contiguous(), use_fast_accum=True)
        activated = up * self.act_fn(gate)
        activated_quantized, activated_scale = torch.ops.fbgemm.quantize_fp8_per_row(activated, num_tokens, self.input_scale_ub)
        down_proj_scale_float32 = self.down_proj_scale.to(torch.float32)
        expert_output = torch.ops.fbgemm.f8f8bf16_rowwise(activated_quantized, self.down_proj[i].transpose(0, 1).contiguous(), activated_scale, down_proj_scale_float32[i].view(-1, 1).contiguous(), use_fast_accum=True)
        next_states[i] = expert_output
    next_states = next_states.to(hidden_states.device)
    return next_states.view(-1, self.hidden_size)

Args:
    hidden_states (torch.Tensor): (batch_size * token_num, hidden_size)

Returns:
    torch.Tensor: (batch_size * token_num, hidden_size)


def add_input(self, mutable_accumulator, element, *args, **kwargs):
    raise NotImplementedError(str(self))

Return result of folding element into accumulator.

    CombineFn implementors must override add_input.

Args:
        mutable_accumulator: the current accumulator,
        may be modified and returned for efficiency
        element: the element to add, should not be mutated
        *args: Additional arguments and side inputs.
        **kwargs: Additional arguments and side inputs.


def is_compatible_with(self, other):
    other = as_dtype(other)
    return self._type_enum in (other.as_datatype_enum, other.base_dtype.as_datatype_enum)

Returns True if the `other` DType will be converted to this DType (TF1).

    Programs written for TensorFlow 2.x do not need this function.
    Instead, they can do equality comparison on `DType` objects directly:
    `tf.as_dtype(this) == tf.as_dtype(other)`.

    This function exists only for compatibility with TensorFlow 1.x, where it
    additionally allows conversion from a reference type (used by
    `tf.compat.v1.Variable`) to its base type.

Args:
        other: A `DType` (or object that may be converted to a `DType`).

Returns:
        True if a Tensor of the `other` `DType` will be implicitly converted to
        this `DType`.


def call(self, pixel_values: np.ndarray | tf.Tensor | None=None, decoder_input_ids: np.ndarray | tf.Tensor | None=None, decoder_attention_mask: np.ndarray | tf.Tensor | None=None, encoder_outputs: Optional[Union[Tuple, TFBaseModelOutput]]=None, past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]]=None, decoder_inputs_embeds: np.ndarray | tf.Tensor | None=None, labels: np.ndarray | tf.Tensor | None=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, training: bool=False, **kwargs) -> Union[TFSeq2SeqLMOutput, Tuple[tf.Tensor]]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    kwargs_encoder = {argument: value for argument, value in kwargs.items() if not argument.startswith('decoder_')}
    kwargs_decoder = {argument[len('decoder_'):]: value for argument, value in kwargs.items() if argument.startswith('decoder_')}
    if encoder_outputs is not None:
        if return_dict and (not isinstance(encoder_outputs, ModelOutput)):
            raise ValueError(f'If `return_dict=True` and `encoder_outputs` is provided, it should be an instance of `ModelOutput`. Got an instance {type(encoder_outputs)} for `encoder_outputs`.')
    if encoder_outputs is None:
        encoder_inputs = {'input_ids': pixel_values, 'output_attentions': output_attentions, 'output_hidden_states': output_hidden_states, 'return_dict': return_dict, 'training': training}
        encoder_inputs.update(kwargs_encoder)
        if 'input_ids' in encoder_inputs:
            encoder_inputs['pixel_values'] = encoder_inputs.pop('input_ids')
        if encoder_inputs['pixel_values'] is None:
            raise ValueError('You have to specify pixel_values')
        if 'labels' in encoder_inputs:
            labels = encoder_inputs.pop('labels')
        if 'decoder_input_ids' in encoder_inputs:
            decoder_input_ids = encoder_inputs.pop('decoder_input_ids')
        if 'decoder_attention_mask' in encoder_inputs:
            decoder_attention_mask = encoder_inputs.pop('decoder_attention_mask')
        encoder_outputs = self.encoder(**encoder_inputs)
    encoder_hidden_states = encoder_outputs[0]
    if self.encoder.config.hidden_size != self.decoder.config.hidden_size and self.decoder.config.cross_attention_hidden_size is None:
        encoder_hidden_states = self.enc_to_dec_proj(encoder_hidden_states)
    if labels is not None and (decoder_input_ids is None and decoder_inputs_embeds is None):
        decoder_input_ids = shift_tokens_right(labels, self.config.pad_token_id, self.config.decoder_start_token_id)
    batch_size, sequence_length = shape_list(encoder_hidden_states)[:2]
    encoder_attention_mask = tf.ones(shape=(batch_size, sequence_length), dtype=tf.int32)
    decoder_inputs = {'input_ids': decoder_input_ids, 'attention_mask': decoder_attention_mask, 'encoder_hidden_states': encoder_hidden_states, 'encoder_attention_mask': encoder_attention_mask, 'inputs_embeds': decoder_inputs_embeds, 'output_attentions': output_attentions, 'output_hidden_states': output_hidden_states, 'use_cache': use_cache, 'past_key_values': past_key_values, 'return_dict': return_dict, 'training': training}
    decoder_inputs.update(kwargs_decoder)
    decoder_outputs = self.decoder(**decoder_inputs)
    logits = decoder_outputs[0]
    loss = None
    if labels is not None:
        warnings.warn(DEPRECATION_WARNING, FutureWarning)
        loss = self.hf_compute_loss(labels, logits)
    if not return_dict:
        past_key_values = None
        if use_cache:
            past_key_values = decoder_outputs[1]
        start_index = sum([1 if x is not None else 0 for x in (loss, logits, past_key_values)])
        if not isinstance(encoder_outputs, tuple):
            encoder_outputs = encoder_outputs.to_tuple()
        output = (loss, logits, past_key_values) + decoder_outputs[start_index:] + encoder_outputs
        output = tuple([x for x in output if x is not None])
        return output
    return TFSeq2SeqLMOutput(loss=loss, logits=decoder_outputs.logits, past_key_values=decoder_outputs.past_key_values, decoder_hidden_states=decoder_outputs.hidden_states, decoder_attentions=decoder_outputs.attentions, cross_attentions=decoder_outputs.cross_attentions, encoder_last_hidden_state=encoder_outputs.last_hidden_state, encoder_hidden_states=encoder_outputs.hidden_states, encoder_attentions=encoder_outputs.attentions)

Returns:

Example:
    ```python
    >>> from transformers import AutoImageProcessor, AutoTokenizer, TFVisionEncoderDecoderModel
    >>> from PIL import Image
    >>> import requests

    >>> image_processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
    >>> decoder_tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")

    >>> # initialize a bert2gpt2 from a pretrained BERT and GPT2 models. Note that the cross-attention layers will be randomly initialized
    >>> model = TFVisionEncoderDecoderModel.from_encoder_decoder_pretrained(
    ...     "google/vit-base-patch16-224-in21k", "openai-community/gpt2"
    ... )

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

    >>> # forward
    >>> pixel_values = image_processor(images=img, return_tensors="tf").pixel_values  # Batch size 1
    >>> decoder_input_ids = decoder_tokenizer("Linda Davis", return_tensors="tf").input_ids  # Batch size 1
    >>> outputs = model(pixel_values=pixel_values, decoder_input_ids=decoder_input_ids)

    >>> # training
    >>> outputs = model(pixel_values=pixel_values, decoder_input_ids=decoder_input_ids, labels=decoder_input_ids)
    >>> loss, logits = outputs.loss, outputs.logits

    >>> # save and load from pretrained
    >>> model.save_pretrained("vit-gpt2")
    >>> model = TFVisionEncoderDecoderModel.from_pretrained("vit-gpt2")

    >>> # generation
    >>> generated = model.generate(pixel_values, decoder_start_token_id=model.config.decoder.bos_token_id)
    ```


def reduce_sum(values, index, name='segmented_reduce_sum'):
    return _segment_reduce(values, index, tf.math.unsorted_segment_sum, name)

Sums a tensor over its segments. Outputs 0 for empty segments. This operations computes the sum over segments, with
    support for:

        - Batching using the first dimensions [B1, B2, ..., Bn]. Each element in a batch can have different indices.
        - Vectorization using the last dimension [V1, V2, ...]. If they are present the output will be a sum of vectors
        rather than scalars.
    Only the middle dimensions [I1, ..., Ik] are reduced by the operation.

Args:
        values: [B1, B2, ..., Bn, I1, .., Ik, V1, V2, ..] tensor of values to be
        averaged.
        index: IndexMap [B1, B2, ..., Bn, I1, .., Ik] index defining the segments.
        name: Name for the TensorFlow ops.

Returns:
        A pair (output_values, output_index) where `output_values` is a tensor of shape [B1, B2, ..., Bn, num_segments,
        V1, V2, ..] and `index` is an IndexMap with shape [B1, B2, ..., Bn, num_segments].


def __init__(self, sess_creator):
    self._sess_creator = sess_creator
    _WrappedSession.__init__(self, self._create_session())

Create a new `_RecoverableSession`.

    The value returned by calling `sess_creator.create_session()` will be the
    session wrapped by this recoverable session.

Args:
        sess_creator: A 'SessionCreator' to be wrapped by recoverable.


def regex_full_match(input, pattern, name=None):
    if isinstance(pattern, util_compat.bytes_or_text_types):
        return gen_string_ops.static_regex_full_match(input=input, pattern=pattern, name=name)
    return gen_string_ops.regex_full_match(input=input, pattern=pattern, name=name)

Match elements of `input` with regex `pattern`.

Args:
        input: string `Tensor`, the source strings to process.
        pattern: string or scalar string `Tensor`, regular expression to use,
        see more details at https://github.com/google/re2/wiki/Syntax
        name: Name of the op.

Returns:
        bool `Tensor` of the same shape as `input` with match results.


def request_stop(self, ex=None):
    self._coord.request_stop(ex=ex)

Request that the coordinator stop the threads.

    See `Coordinator.request_stop()`.

Args:
        ex: Optional `Exception`, or Python `exc_info` tuple as returned by
        `sys.exc_info()`.  If this is the first call to `request_stop()` the
        corresponding exception is recorded and re-raised from `join()`.


def _infer_type(str_val, na_value, prev_type):
    if str_val in ('', na_value):
        return prev_type
    type_list = [dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64, dtypes.string]
    type_functions = [_is_valid_int32, _is_valid_int64, lambda str_val: _is_valid_float(str_val, dtypes.float32), lambda str_val: _is_valid_float(str_val, dtypes.float64), lambda str_val: True]
    for i in range(len(type_list)):
        validation_fn = type_functions[i]
        if validation_fn(str_val) and (prev_type is None or prev_type in type_list[:i + 1]):
            return type_list[i]

Given a string, infers its tensor type.

    Infers the type of a value by picking the least 'permissive' type possible,
    while still allowing the previous type inference for this column to be valid.

Args:
        str_val: String value to infer the type of.
        na_value: Additional string to recognize as a NA/NaN CSV value.
        prev_type: Type previously inferred based on values of this column that
        we've seen up till now.

Returns:
        Inferred dtype.


def __init__(self, config: PerceiverConfig, in_channels: int, postproc_type: str='patches') -> None:
    super().__init__()
    if postproc_type not in ('patches',):
        raise ValueError('Invalid postproc_type!')
    self.classifier = nn.Linear(in_channels, config.samples_per_patch)

Audio postprocessing for Perceiver. Can be used to convert the decoder output to audio features.

Args:
        config ([*PerceiverConfig*]):
            Model configuration.
        in_channels (`int`):
            Number of channels in the input.
        postproc_type (`str`, *optional*, defaults to `"patches"`):
            Postprocessor type to use. Currently, only "patches" is supported.


def annotate_test_case(self, test_case: Iterator[str]) -> tuple[Iterator[str], Iterator[str]]:
    split_input = HloStreamSplitter(test_case, record_directives={DirectiveComment.RUN})
    optimizer_input, passthrough_input = itertools.tee(_fix_whitespace(split_input.non_directive_lines()))
    test_case_output = optimize_hlo(optimizer_input, self._optimizer_path, self._optimizer_args, expand_to_input=self._expand_to_input)
    test_case_checks = HloFileCheckLines(test_case_output)
    transformed_test_case = itertools.chain(test_case_checks, ['\n'], passthrough_input)
    run_directives = split_input.directive_history(DirectiveComment.RUN)
    return (transformed_test_case, run_directives)

Transforms a test case, inserting FileCheck directives above it.

Args:
        test_case: An iterator over the lines of a single test case.

Returns:
        A tuple of two iterators. The first iterator yields the annotated test
        case containing FileCheck directives followed by the original HLO IR. The
        second iterator yields any RUN directives that were stripped from the
        original test case.


def report_fhir_path_error(self, element_path: str, fhir_path_constraint: str, msg: str) -> None:

Reports a FHIRPath constraint error during validation and/or encoding.

    The base implementation logs to the `error` context and raises `e` by
    default. Subclasses should override this behavior as necessary.

Args:
        element_path: The path to the field that the constraint is defined on.
        fhir_path_constraint: The FHIRPath constraint expression.
        msg: The error message produced.


def create_proxy_api_files(output_files, proxy_module_root, output_dir):
    for file_path in output_files:
        module = get_module(os.path.dirname(file_path), output_dir)
        if not os.path.isdir(os.path.dirname(file_path)):
            os.makedirs(os.path.dirname(file_path))
        contents = f'from {proxy_module_root}.{module} import *'
        with open(file_path, 'w') as fp:
            fp.write(contents)

Creates __init__.py files in proxy format for the Python API.

Args:
        output_files: List of __init__.py file paths to create.
        proxy_module_root: Module root for proxy-import format. If specified, proxy
        files with content like `from proxy_module_root.proxy_module import *`
        will be created to enable import resolution under TensorFlow.
        output_dir: output API root directory.


def __init__(self, expr_token):
    self.expr_token = expr_token
    self.for_token = None
    self.has_split_at_for = False
    self.has_interior_split = False

Maintains the state of list comprehension formatting decisions.

    A stack of ComprehensionState objects are kept to ensure that list
    comprehensions are wrapped with well-defined rules.

Attributes:
        expr_token: The first token in the comprehension.
        for_token: The first 'for' token of the comprehension.
        opening_bracket: The opening bracket of the list comprehension.
        closing_bracket: The closing bracket of the list comprehension.
        has_split_at_for: Whether there is a newline immediately before the
        for_token.
        has_interior_split: Whether there is a newline within the comprehension.
        That is, a split somewhere after expr_token or before closing_bracket.


def __init__(self, config: Llama4VisionConfig):
    super().__init__()
    self.config = config
    self.layers = nn.ModuleList([Llama4VisionEncoderLayer(config) for _ in range(config.num_hidden_layers)])
    self.gradient_checkpointing = False
    self.config = config

Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
    [`Llama4VisionEncoderLayer`].

Args:
        config: Llama4VisionConfig


def __init__(self, serialized=None, mesh_shape=None, device_coordinates=None):
    self._serialized = serialized
    if serialized:
        self._parse_topology(serialized)
    else:
        self._mesh_shape = numpy_compat.np_asarray(mesh_shape, dtype=np.int32)
        self._device_coordinates = numpy_compat.np_asarray(device_coordinates, dtype=np.int32)
        if len(self._mesh_shape) != 4 or any(self._mesh_shape < 1):
            raise ValueError(f'`mesh_shape` must be a sequence of 4 positive entries; got `mesh_shape={self._mesh_shape}`')
        if len(self._device_coordinates.shape) != 3 or self._device_coordinates.shape[2] != len(self._mesh_shape):
            raise ValueError('`device_coordinates` must be a rank 3 int32 array with minor dimension equal to the `mesh_shape` rankgot device_coordinates.shape={} len(device_coordinates.shape)={} device_coordinates.shape[2]={} mesh_shape={}, len(mesh_shape)={}'.format(self._device_coordinates.shape, len(self._device_coordinates.shape), self._device_coordinates.shape[2], self._mesh_shape, len(self._mesh_shape)))
    self._topology_tasks, self._topology_devices = self._invert_topology()
    self._missing_devices = np.argwhere(self._topology_tasks < 0)

Builds a Topology object.

    If `serialized` is not `None`, the topology is parsed from `serialized` and
    the other arguments are ignored. Otherwise, the topology is computed from
    `mesh_shape` and `device_coordinates`.

Args:
        serialized: A serialized `TopologyProto`, or `None`. If not `None`, the
        serialized proto is parsed to discover the topology.
        mesh_shape: A sequence of 4 positive integers, or `None`. If not `None`,
        the shape of the TPU topology, in number of cores. Ignored if
        `serialized` is not `None`.
        device_coordinates: A rank 3 numpy array that describes the mapping from
        TensorFlow TPU devices to TPU fabric coordinates, or `None`. If
        specified, array is a rank 3 int32 array with shape
        `[tasks, devices, axis]`.  `tasks` is the number of tasks in the TPU
        cluster, `devices` is the number of TPU devices per task, and `axis` is
        the number of axes in the TPU cluster topology. Each entry gives the
        `axis`-th coordinate in the topology of a task/device pair. TPU
        topologies are 4-dimensional, with dimensions `(x, y, z, core number)`.
        This arg is ignored if `serialized is not `None`.

Raises:
        ValueError: If `serialized` does not describe a well-formed topology.
        ValueError: If `serialized` is `None` and `mesh_shape` is not a sequence
        of 4 positive integers.
        ValueError: If `serialized` is `None` and `device_coordinates` is not a
        rank 3 numpy int32 array that describes a valid coordinate mapping.


def make_one_shot_iterator(self) -> Union[iterator_ops.Iterator, iterator_ops.OwnedIterator]:
    return self._make_one_shot_iterator()

Creates an iterator for elements of this dataset.

    Note: The returned iterator will be initialized automatically.
    A "one-shot" iterator does not currently support re-initialization. For
    that see `make_initializable_iterator`.

Example:
    ```python
    # Building graph ...
    dataset = ...
    next_value = dataset.make_one_shot_iterator().get_next()

    # ... from within a session ...
    try:
        while True:
        value = sess.run(next_value)
        ...
    except tf.errors.OutOfRangeError:
        pass
    ```

Returns:
        An `tf.data.Iterator` for elements of this dataset.


def retrieve_collected_errors():
    serialized_message_list = wrap_converter.wrapped_retrieve_collected_errors()
    return list(map(converter_error_data_pb2.ConverterErrorData.FromString, serialized_message_list))

Returns and clears the list of collected errors in ErrorCollector.

    The RetrieveCollectedErrors function in C++ returns a list of serialized proto
    messages. This function will convert them to ConverterErrorData instances.

Returns:
        A list of ConverterErrorData.


def send_eager_tracebacks(destinations, origin_stack, send_source=True):
    _send_call_tracebacks(destinations, origin_stack, is_eager_execution=True, send_source=send_source)

Send the tracebacks of an eager execution call to debug server(s).

Args:
        destinations: gRPC destination addresses, a `str` or a `list` of `str`s,
        e.g., "localhost:4242". If a `list`, gRPC requests containing the same
        origin_stack: The traceback of the eager operation invocation.
        send_source: Whether the source files involved in the op tracebacks but
        outside the TensorFlow library are to be sent.


def forward(self, input_ids: Optional[torch.LongTensor]=None, pixel_values: Optional[torch.FloatTensor]=None, attention_mask: Optional[torch.Tensor]=None, return_loss: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, interpolate_pos_encoding: bool=False, return_base_image_embeds: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[Tuple, OwlViTOutput]:
    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
    vision_outputs = self.vision_model(pixel_values=pixel_values, output_attentions=output_attentions, output_hidden_states=output_hidden_states, interpolate_pos_encoding=interpolate_pos_encoding, return_dict=return_dict)
    text_outputs = self.text_model(input_ids=input_ids, attention_mask=attention_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    text_embeds = text_outputs[1]
    text_embeds = self.text_projection(text_embeds)
    image_embeds = vision_outputs[1]
    image_embeds = self.visual_projection(image_embeds)
    image_embeds = image_embeds / torch.linalg.norm(image_embeds, ord=2, dim=-1, keepdim=True)
    text_embeds_norm = text_embeds / torch.linalg.norm(text_embeds, ord=2, dim=-1, keepdim=True)
    logit_scale = self.logit_scale.exp().to(image_embeds.device)
    logits_per_text = torch.matmul(text_embeds_norm, image_embeds.t()) * logit_scale
    logits_per_image = logits_per_text.t()
    loss = None
    if return_loss:
        loss = owlvit_loss(logits_per_text)
    text_embeds = text_embeds_norm
    if not return_dict:
        output = (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)
        return (loss,) + output if loss is not None else output
    return OwlViTOutput(loss=loss, logits_per_image=logits_per_image, logits_per_text=logits_per_text, text_embeds=text_embeds, image_embeds=image_embeds, text_model_output=text_outputs, vision_model_output=vision_outputs)

return_loss (`bool`, *optional*):
    Whether or not to return the contrastive loss.
    return_base_image_embeds (`bool`, *optional*):
    Whether or not to return the base image embeddings.

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

    >>> model = OwlViTModel.from_pretrained("google/owlvit-base-patch32")
    >>> processor = AutoProcessor.from_pretrained("google/owlvit-base-patch32")
    >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
    >>> image = Image.open(requests.get(url, stream=True).raw)
    >>> inputs = processor(text=[["a photo of a cat", "a photo of a dog"]], images=image, return_tensors="pt")
    >>> outputs = model(**inputs)
    >>> logits_per_image = outputs.logits_per_image  # this is the image-text similarity score
    >>> probs = logits_per_image.softmax(dim=1)  # we can take the softmax to get the label probabilities
    ```


def zeros(shape, dtype=None, name=None):
    with ops.init_scope():
        if dtype is None:
            dtype = floatx()
        tf_dtype = dtypes_module.as_dtype(dtype)
        v = array_ops.zeros(shape=shape, dtype=tf_dtype, name=name)
        if py_all(v.shape.as_list()):
            return variable(v, dtype=dtype, name=name)
        return v

Instantiates an all-zeros variable and returns it.

Args:
        shape: Tuple or list of integers, shape of returned Keras variable
        dtype: data type of returned Keras variable
        name: name of returned Keras variable

Returns:
        A variable (including Keras metadata), filled with `0.0`.
        Note that if `shape` was symbolic, we cannot return a variable,
        and will return a dynamically-shaped tensor instead.

Example:
    >>> kvar = tf.keras.backend.zeros((3,4))
    >>> tf.keras.backend.eval(kvar)
    array([[0.,  0.,  0.,  0.],
        [0.,  0.,  0.,  0.],
        [0.,  0.,  0.,  0.]], dtype=float32)
    >>> A = tf.constant([1,2,3])
    >>> kvar2 = tf.keras.backend.zeros(A.shape) # [0., 0., 0.]
    >>> tf.keras.backend.eval(kvar2)
    array([0., 0., 0.], dtype=float32)
    >>> kvar3 = tf.keras.backend.zeros(A.shape,dtype=tf.int32)
    >>> tf.keras.backend.eval(kvar3)
    array([0, 0, 0], dtype=int32)
    >>> kvar4 = tf.keras.backend.zeros([2,3])
    >>> tf.keras.backend.eval(kvar4)
    array([[0., 0., 0.],
        [0., 0., 0.]], dtype=float32)


def _step(time, output_ta_t, prev_output, *states):
    current_input = tuple((ta[time] for ta in input_ta))
    current_input = tree.pack_sequence_as(inputs, current_input)
    mask_t = masking_fn(time)
    output, new_states = step_function(current_input, tuple(states) + tuple(constants))
    flat_output = tree.flatten(output)
    flat_mask_output = flat_zero_output if zero_output_for_mask else tree.flatten(prev_output)
    flat_new_output = compute_masked_output(mask_t, flat_output, flat_mask_output)
    flat_state = tree.flatten(states)
    flat_new_state = tree.flatten(new_states)
    flat_final_state = compute_masked_output(mask_t, flat_new_state, flat_state)
    new_states = tree.pack_sequence_as(new_states, flat_final_state)
    ta_index_to_write = time if return_all_outputs else 0
    for ta, out in zip(output_ta_t, flat_new_output):
        ta[ta_index_to_write] = out
    return (time + 1, output_ta_t, tuple(flat_new_output)) + tuple(new_states)

RNN step function.

Args:
        time: Current timestep value.
        output_ta_t: TensorArray.
        prev_output: tuple of outputs from time - 1.
        *states: List of states.

Returns:
        Tuple: `(time + 1, output_ta_t, output) + tuple(new_states)`


def _get_dir_size(self, path: str='.'):
    total = 0
    for root, _, files in os.walk(path):
        for filename in files:
            total += os.path.getsize(os.path.join(root, filename))
    return total

Get the total size of files and sub-directories under the path.

Args:
        path: Path of a directory or a file to calculate the total size.

Returns:
        Total size of the directory or a file.


def AddCalledComponent(self, component, target, args, filename, lineno, capacity, action=CALLED_CALLABLE):
    element = FireTraceElement(component=component, action=action, target=target, args=args, filename=filename, lineno=lineno, capacity=capacity)
    self.elements.append(element)

Adds an element to the trace indicating that a component was called.

    Also applies to instantiating a class.

Args:
        component: The result of calling the callable.
        target: The name of the callable.
        args: The args consumed in order to call this callable.
        filename: The file in which the callable is defined, or None if N/A.
        lineno: The line number on which the callable is defined, or None if N/A.
        capacity: (bool) Whether the callable could have accepted additional args.
        action: The value to include as the action in the FireTraceElement.


def _google_section(line_info):
    colon_index = line_info.remaining.find(':')
    possible_title = line_info.remaining[:colon_index]
    return _section_from_possible_title(possible_title)

Checks whether the current line is the start of a new Google-style section.

    This docstring is a Google-style docstring. Google-style sections look like
    this:

        Section Name:
        section body goes here

Args:
        line_info: Information about the current line.

Returns:
        A Section type if one matches, or None if no section type matches.


def _stop_profiler(self, save=True):
    if not self._profiler_started:
        return
    try:
        backend.tensorboard.stop_trace(save=save)
    except Exception as e:
        logging.error('Failed to stop profiler: %s', e)
    finally:
        self._profiler_started = False

Stops the profiler if currently active.

Args:
        save: Whether to save the profiler results to TensorBoard.


def hertz_to_mel(frequencies_hertz):
    return _MEL_HIGH_FREQUENCY_Q * np.log(1.0 + frequencies_hertz / _MEL_BREAK_FREQUENCY_HERTZ)

Convert frequencies to mel scale using HTK formula.

    Copied from
    https://github.com/tensorflow/models/blob/master/research/audioset/mel_features.py.

Args:
        frequencies_hertz: Scalar or np.array of frequencies in hertz.

Returns:
        Object of same size as frequencies_hertz containing corresponding values
        on the mel scale.


def __init__(self, x, y=None, **kwargs):
    if not self.can_handle(x, y):
        raise ValueError('{} Cannot handle input {}, {}'.format(self.__class__, x, y))

Create a DataAdapter based on data inputs.

    The caller must make sure to call `can_handle()` first before invoking this
    method. Provide unsupported data type will result into unexpected behavior.

Args:
        x: input features.
        y: target labels. Note that y could be None in the case of prediction.
        **kwargs: Other keyword arguments for DataAdapter during the construction
        of the tf.dataset.Dataset. For example:
        - Numpy data might have `sample_weights` which will be used for
        weighting the loss function during training.
        - Numpy data might need to have `batch_size` parameter when constructing
        the dataset and iterator.
        - Certain input might need to be distribution strategy aware. When
        `distribution_strategy` is passed, the created dataset need to respect
        the strategy.
        DataAdapter might choose to ignore any keyword argument if it doesn't
        use it, or raise exception if any required argument is not provide.


def __init__(self, session_creator, hooks, should_recover, stop_grace_period_secs=120):
    self._graph_was_finalized = ops.get_default_graph().finalized
    self._hooks = hooks or []
    for h in self._hooks:
        h.begin()
    worker_context = distribute_coordinator_context.get_current_worker_context()
    if not session_creator and worker_context:
        session_creator = worker_context.session_creator()
    self._coordinated_creator = self._CoordinatedSessionCreator(session_creator=session_creator or ChiefSessionCreator(), hooks=self._hooks, stop_grace_period_secs=stop_grace_period_secs)
    if should_recover:
        self._sess = _RecoverableSession(self._coordinated_creator)
    else:
        self._sess = self._coordinated_creator.create_session()

Sets up a Monitored or Hooked Session.

Args:
        session_creator: A factory object to create session. Typically a
        `ChiefSessionCreator` or a `WorkerSessionCreator`.
        hooks: An iterable of `SessionRunHook' objects.
        should_recover: A bool. Indicates whether to recover from `AbortedError`
        and `UnavailableError` or not.
        stop_grace_period_secs: Number of seconds given to threads to stop after
        `close()` has been called.


def _init_values_from_proto(self, values_def, import_scope=None):
    assert isinstance(values_def, control_flow_pb2.ValuesDef)
    self._values = set((ops.prepend_name_scope(value, import_scope) for value in values_def.values))
    g = ops.get_default_graph()
    self._external_values = {}
    for k, v in values_def.external_values.items():
        k = ops.prepend_name_scope(k, import_scope)
        self._external_values[k] = g.as_graph_element(ops.prepend_name_scope(v, import_scope))
    op_names = set([op.split(':')[0] for op in self._values - set(self._external_values.keys())])
    for op in op_names:
        g.as_graph_element(op)._set_control_flow_context(self)

Initializes values and external_values from `ValuesDef` protocol buffer.

Args:
        values_def: `ValuesDef` protocol buffer.
        import_scope: Optional `string`. Name scope to add.


def assert_sources_equal_reference_source(reference_source_info, sources_info):
    if not (isinstance(reference_source_info, tuple) and len(reference_source_info) == 3 and isinstance(reference_source_info[0], iobase.BoundedSource)):
        raise ValueError('reference_source_info must a three-tuple where firstitem of the tuple gives a iobase.BoundedSource. Received: %r' % reference_source_info)
    reference_records = read_from_source(*reference_source_info)
    source_records = []
    for source_info in sources_info:
        assert isinstance(source_info, tuple)
        assert len(source_info) == 3
        if not (isinstance(source_info, tuple) and len(source_info) == 3 and isinstance(source_info[0], iobase.BoundedSource)):
            raise ValueError('source_info must a three tuple where firstitem of the tuple gives a iobase.BoundedSource. Received: %r' % source_info)
        if type(reference_source_info[0].default_output_coder()) != type(source_info[0].default_output_coder()):
            raise ValueError('Reference source %r and the source %r must use the same coder. They are using %r and %r respectively instead.' % (reference_source_info[0], source_info[0], type(reference_source_info[0].default_output_coder()), type(source_info[0].default_output_coder())))
        source_records.extend(read_from_source(*source_info))
    if len(reference_records) != len(source_records):
        raise ValueError('Reference source must produce the same number of records as the list of sources. Number of records were %d and %d instead.' % (len(reference_records), len(source_records)))
    if equal_to(reference_records)(source_records):
        raise ValueError('Reference source and provided list of sources must produce the same set of records.')

Tests if a reference source is equal to a given set of sources.

    Given a reference source (a :class:`~apache_beam.io.iobase.BoundedSource`
    and a position range) and a list of sources, assert that the union of the
    records read from the list of sources is equal to the records read from the
    reference source.

Args:
        reference_source_info        (Tuple[~apache_beam.io.iobase.BoundedSource, int, int]):
        a three-tuple that gives the reference
        :class:`~apache_beam.io.iobase.BoundedSource`, position to start
        reading at, and position to stop reading at.
        sources_info        (Iterable[Tuple[~apache_beam.io.iobase.BoundedSource, int, int]]):
        a set of sources. Each source is a three-tuple that is of the same
        format described above.

Raises:
        ValueError: if the set of data produced by the reference source
        and the given set of sources are not equivalent.


def __init__(self, input_energy: energy.BitstringEnergy, num_expectation_samples: int, initial_seed: Union[None, tf.Tensor]=None, name: Union[None, str]=None):
    super().__init__(input_energy, initial_seed, name)
    self.num_expectation_samples = num_expectation_samples

Initializes an EnergyInference.

Args:
        input_energy: The parameterized energy function which defines this
        distribution via the equations of an energy based model.  This class
        assumes that all parameters of `energy` are `tf.Variable`s and that
        they are all returned by `energy.variables`.
        num_expectation_samples: Number of samples to draw and use for estimating
        the expectation value.
        initial_seed: PRNG seed; see tfp.random.sanitize_seed for details. This
        seed will be used in the `sample` method.  If None, the seed is updated
        after every inference call.  Otherwise, the seed is fixed.
        name: Optional name for the model.


def get_special_tokens_mask(self, token_ids_0: List, token_ids_1: Optional[List]=None, already_has_special_tokens: bool=False) -> List[int]:
    if already_has_special_tokens:
        if token_ids_1 is not None:
            raise ValueError('You should not supply a second sequence if the provided sequence of ids is already formatted with special tokens for the model.')
        return [1 if token in self.all_special_ids else 0 for token in token_ids_0]
    mask = [1] + [0] * len(token_ids_0) + [1]
    if token_ids_1 is not None:
        mask += [0] * len(token_ids_1) + [1]
    return mask

Retrieves 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` or `encode_plus` methods.

Args:
        token_ids_0 (`List[int]`):
            List of ids of the first sequence.
        token_ids_1 (`List[int]`, *optional*):
            List of ids of the second sequence.
        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:
        A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.


def make_batched_images(images) -> List[List[ImageInput]]:
    if isinstance(images, (list, tuple)) and isinstance(images[0], (list, tuple)) and is_valid_image(images[0][0]):
        return [img for img_list in images for img in img_list]
    elif isinstance(images, (list, tuple)) and is_valid_image(images[0]):
        return images
    elif is_valid_image(images):
        return [images]
    raise ValueError(f'Could not make batched images from {images}')

Accepts images in list or nested list format, and makes a list of images for preprocessing.

Args:
        images (`Union[List[List[ImageInput]], List[ImageInput], ImageInput]`):
            The input image.

Returns:
        list: A list of images.


def load_raw_type(self, typ: type[Any]) -> abstract.BaseValue:
    if typ is type(None):
        return self.consts[None]
    pytd_node = self._pytd_loader.lookup_pytd(typ.__module__, typ.__name__)
    return self._load_pytd_node(pytd_node)

Converts a raw type to an abstract value.

    For convenience, this method can also be called via ctx.types[typ].

Args:
        typ: The type.

Returns:
        The abstract representation of the type. For example, when passed `int`,
        this function returns `abstract.SimpleClass(int)`.


def recursive_create_dir_v2(path):
    _pywrap_file_io.RecursivelyCreateDir(compat.path_to_bytes(path))

Creates a directory and all parent/intermediate directories.

    It succeeds if path already exists and is writable.

Args:
        path: string, name of the directory to be created

Raises:
        errors.OpError: If the operation fails.


def conv(self, input_tensor: core.Tensor) -> Mapping[str, core.Tensor]:
    scale = [1.0] * self.out_channel_size
    offset = [0.5] * self.out_channel_size
    mean, variance = (scale, offset)
    out = nn_ops.conv2d(input_tensor, self.filters, strides=strides, dilations=dilations, padding=padding, data_format='NHWC')
    if has_bias:
        out = nn_ops.bias_add(out, self.bias, data_format='NHWC')
    if has_batch_norm:
        out, _, _, _, _, _ = nn_ops.fused_batch_norm_v3(out, scale, offset, mean, variance, is_training=False)
    if activation_fn is not None:
        out = activation_fn(out)
    return {'output': out}

Performs a 2D convolution operation.

Args:
        input_tensor: Input tensor to perform convolution on.

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


def _make_class_weight_map_fn(class_weight):
    class_ids = list(sorted(class_weight.keys()))
    expected_class_ids = list(range(len(class_ids)))
    if class_ids != expected_class_ids:
        error_msg = 'Expected `class_weight` to be a dict with keys from 0 to one less than the number of classes, found {}'.format(class_weight)
        raise ValueError(error_msg)
    class_weight_tensor = tensor_conversion.convert_to_tensor_v2_with_dispatch([class_weight[int(c)] for c in class_ids])

    def _class_weights_map_fn(*data):
        """Convert `class_weight` to `sample_weight`."""
        x, y, sw = unpack_x_y_sample_weight(data)
        if nest.is_nested(y):
            raise ValueError('`class_weight` is only supported for Models with a single output.')
        if y.shape.rank > 2:
            raise ValueError('`class_weight` not supported for 3+ dimensional targets.')
        y_classes = smart_cond.smart_cond(y.shape.rank == 2 and backend.shape(y)[1] > 1, lambda: backend.argmax(y, axis=1), lambda: math_ops.cast(backend.reshape(y, (-1,)), dtypes.int64))
        cw = array_ops.gather_v2(class_weight_tensor, y_classes)
        if sw is not None:
            cw = math_ops.cast(cw, sw.dtype)
            sw, cw = expand_1d((sw, cw))
            sw = sw * cw
        else:
            sw = cw
        return (x, y, sw)
    return _class_weights_map_fn

Applies class weighting to a `Dataset`.

    The `Dataset` is assumed to be in format `(x, y)` or `(x, y, sw)`, where
    `y` must be a single `Tensor`.

Args:
        class_weight: A map where the keys are integer class ids and values are
        the class weights, e.g. `{0: 0.2, 1: 0.6, 2: 0.3}`

Returns:
        A function that can be used with `tf.data.Dataset.map` to apply class
        weighting.


def rgb_to_yuv(images):
    images = ops.convert_to_tensor(images, name='images')
    kernel = ops.convert_to_tensor(_rgb_to_yuv_kernel, dtype=images.dtype, name='kernel')
    ndims = images.get_shape().ndims
    return math_ops.tensordot(images, kernel, axes=[[ndims - 1], [0]])

Converts one or more images from RGB to YUV.

    Outputs a tensor of the same shape as the `images` tensor, containing the YUV
    value of the pixels.
    The output is only well defined if the value in images are in [0, 1].
    There are two ways of representing an image: [0, 255] pixel values range or
    [0, 1] (as float) pixel values range. Users need to convert the input image
    into a float [0, 1] range.

Args:
        images: 2-D or higher rank. Image data to convert. Last dimension must be
        size 3.

Returns:
        images: tensor with the same shape as `images`.


def _ImageDimensions(image, rank):
    if image.get_shape().is_fully_defined():
        return image.get_shape().as_list()
    else:
        static_shape = image.get_shape().with_rank(rank).as_list()
        dynamic_shape = array_ops_stack.unstack(array_ops.shape(image), rank)
        return [s if s is not None else d for s, d in zip(static_shape, dynamic_shape)]

Returns the dimensions of an image tensor.

Args:
        image: A rank-D Tensor. For 3-D  of shape: `[height, width, channels]`.
        rank: The expected rank of the image

Returns:
        A list of corresponding to the dimensions of the
        input image.  Dimensions that are statically known are python integers,
        otherwise, they are integer scalar tensors.


def get_help(self, cmd_prefix=None):
    if not cmd_prefix:
        help_info = RichTextLines([])
        if self._help_intro:
            help_info.extend(self._help_intro)
        sorted_prefixes = sorted(self._handlers)
        for cmd_prefix in sorted_prefixes:
            lines = self._get_help_for_command_prefix(cmd_prefix)
            lines.append('')
            lines.append('')
            help_info.extend(RichTextLines(lines))
        return help_info
    else:
        return RichTextLines(self._get_help_for_command_prefix(cmd_prefix))

Compile help information into a RichTextLines object.

Args:
        cmd_prefix: Optional command prefix. As the prefix itself or one of its
        aliases.

Returns:
        A RichTextLines object containing the help information. If cmd_prefix
        is None, the return value will be the full command-line help. Otherwise,
        it will be the help information for the specified command.


def get_image_features(self, pixel_values: Optional[torch.FloatTensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, interpolate_pos_encoding: bool=False) -> torch.FloatTensor:
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    vision_outputs: BaseModelOutputWithPooling = self.vision_model(pixel_values=pixel_values, output_attentions=output_attentions, output_hidden_states=output_hidden_states, interpolate_pos_encoding=interpolate_pos_encoding)
    pooled_output = vision_outputs.pooler_output
    return pooled_output

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

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

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

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

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

    >>> with torch.no_grad():
    ...     image_features = model.get_image_features(**inputs)
    ```


def forward(self, embedding: torch.Tensor) -> torch.Tensor:
    embedding = self.linear(self.dropout(embedding))
    if self.use_cls_token:
        embedding = embedding[:, :, 1:, :]
    return embedding

Parameters:
    embedding (`torch.Tensor` of shape `(bs, num_channels, num_patches, d_model)` or
            `(bs, num_channels, num_patches+1, d_model)` if `cls_token` is set to True, *required*):
        Embedding from the model

Returns:
    `torch.Tensor` of shape `(bs, num_channels, num_patches, d_model)` or
                    `(bs, num_channels, num_patches+1, d_model)` if `cls_token` is set to True


def get_extended_attention_mask(self, word_attention_mask: torch.LongTensor, entity_attention_mask: Optional[torch.LongTensor]):
    attention_mask = word_attention_mask
    if entity_attention_mask is not None:
        attention_mask = torch.cat([attention_mask, entity_attention_mask], dim=-1)
    if attention_mask.dim() == 3:
        extended_attention_mask = attention_mask[:, None, :, :]
    elif attention_mask.dim() == 2:
        extended_attention_mask = attention_mask[:, None, None, :]
    else:
        raise ValueError(f'Wrong shape for attention_mask (shape {attention_mask.shape})')
    extended_attention_mask = extended_attention_mask.to(dtype=self.dtype)
    extended_attention_mask = (1.0 - extended_attention_mask) * torch.finfo(self.dtype).min
    return extended_attention_mask

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

Args:
        word_attention_mask (`torch.LongTensor`):
            Attention mask for word tokens with ones indicating tokens to attend to, zeros for tokens to ignore.
        entity_attention_mask (`torch.LongTensor`, *optional*):
            Attention mask for entity tokens with ones indicating tokens to attend to, zeros for tokens to ignore.

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


def _retrieve_variables_impl(config: Text, hosts: List[Tuple[int, Text]], variables: Dict[Text, Dict[Text, tf_variables.Variable]], table_config: tpu_embedding_v2_utils.TableConfig):
    for host_id, host in enumerate(hosts):
        with ops.device(host):
            for table in table_config:
                retrieved = table.optimizer._retrieve()(table_name=table.name, num_shards=len(hosts), shard_id=host_id, config=config)
                if not isinstance(retrieved, tuple):
                    retrieved = (retrieved,)
                for i, slot in enumerate(['parameters'] + table.optimizer._slot_names()):
                    sharded_var = variables[table.name][slot]
                    if host_id < len(sharded_var.variables):
                        sharded_var.variables[host_id].assign(retrieved[i])
                config = None

Retrieve embedding tables from TPU to host memory.

Args:
        config: A serialized TPUEmbeddingConfiguration proto.
        hosts: A list of all the host CPU devices.
        variables: A dictionary of dictionaries of TPUEmbeddingVariables. First key
        is the table name, second key is 'parameters' or the optimizer slot name.
        table_config: A list of tf.tpu.experimental.embedding.TableConfig objects.


def __init__(self, name, func=None):
    NamedObject.__init__(self, name)
    self.func = func
    self.args = []

External function call.

Attributes:
        func: The callable function this wraps and runs.
        args: List of arguments for given function with |name|.


def path_of_module(self, mod: nn.Module) -> str:
    try:
        return super().path_of_module(mod)
    except NameError as e:
        if self.allow_insert_stateless_mods and len(list(mod.parameters())) == 0 and (len(list(mod.buffers())) == 0):
            path = self._insert_module_as_submodule(mod)
            return path
        raise e

Helper method to find the qualified name of `mod` in the Module hierarchy of `root`. For example, if `root` has
    a submodule named `foo`, which has a submodule named `bar`, passing `bar` into this function will return the
    string "foo.bar".

Args:
        mod (str): The `Module` to retrieve the qualified name for.


def interpolate_graph(message, graph):
    parsed_messaged, _, node_tags = parse_message(message)
    error_message = ['Graph execution error:', '']
    for tag in node_tags:
        try:
            op = graph.get_operation_by_name(tag.name)
        except KeyError:
            continue
        else:
            error_message.append(_build_node_error_message(op))
    error_message.append(parsed_messaged.strip())
    return '\n'.join(error_message)

Interpolates an error message.

    The error message can contain tags of form `{{node_type node_name}}`
    which will be parsed to identify the tf.Graph and op. If the op contains
    traceback, the traceback will be attached to the error message.

Args:
        message: A string to interpolate.
        graph: ops.Graph object containing all nodes referenced in the error
        message.

Returns:
        The error message string with node definition traceback.


def get_input_shape_at(self, node_index):
    return self._get_node_attribute_at_index(node_index, 'input_shapes', 'input shape')

Retrieves the input shape(s) of a layer at a given node.

Args:
        node_index: Integer, index of the node
            from which to retrieve the attribute.
            E.g. `node_index=0` will correspond to the
            first time the layer was called.

Returns:
        A shape tuple
        (or list of shape tuples if the layer has multiple inputs).

Raises:
        RuntimeError: If called in Eager mode.


def wait_on_failure(self, on_failure_fn=None, on_transient_failure_fn=None, on_recovery_fn=None, worker_device_name='(unknown)'):
    assert self._should_preemption_thread_run
    try:
        yield
    except (errors.OpError, ClosureInputError, ClosureAbortedError) as e:
        with self._next_task_state_cond:
            self._next_task_state_cond.wait(_POLL_FREQ_IN_SEC * 1.25)
        with self._next_task_state_cond:
            self._next_task_state_cond.wait(_POLL_FREQ_IN_SEC * 1.25)
        if not self._task_states:
            self._log_ps_failure_and_raise(e, 0)
        worker_states = self._task_states[:self._num_workers]
        ps_states = self._task_states[self._num_workers:]
        if any(ps_states):
            failed_ps_index = [ix for ix, ps_state in enumerate(ps_states) if ps_state]
            self._log_ps_failure_and_raise(e, failed_ps_index[0])
        worker_ix = int(worker_device_name.split(':')[-1])
        if worker_states[worker_ix]:
            if self._cluster.closure_queue._cancellation_mgr.is_cancelled:
                if isinstance(e, errors.CancelledError):
                    raise e
                else:
                    raise errors.CancelledError(None, None, 'The corresponding function was cancelled while attempting to recover from worker failure.')
            self._handle_failure_and_recovery(e, on_failure_fn, on_transient_failure_fn, on_recovery_fn, worker_device_name)
            return
        if self._cluster._record_and_ignore_transient_timeouts(e):
            logging.error('Remote function on worker %s failed with %r:%s\nThis derived error is ignored and not reported to users.', worker_device_name, e, e)
            if on_transient_failure_fn:
                on_transient_failure_fn()
            return
        raise e

Catches errors during closure execution and handles them.

Args:
        on_failure_fn: an optional function to run if preemption happens.
        on_transient_failure_fn: an optional function to run if transient failure
        happens.
        on_recovery_fn: an optional function to run when a worker is recovered
        from preemption.
        worker_device_name: the device name of the worker instance that is passing
        through the failure.

Yields:
        None.


def filter_def_file(def_file: str, filter_file: str, filtered_file: str) -> None:
    with open(filter_file, 'r', encoding='utf-8') as filter_file_handle:
        filter_json: Dict[str, Any] = json.load(filter_file_handle)
        inclusion_patterns: List[str] = filter_json['global'] + ['EXPORTS', '*;*']
        incl_patterns: List[Pattern[str]] = [re.compile(re.escape(p).replace('\\*', '.*')) for p in inclusion_patterns]
        exclusion_patterns: List[str] = filter_json['local']
        excl_patterns: List[Pattern[str]] = [re.compile(re.escape(p).replace('\\*', '.*')) for p in exclusion_patterns]
    with open(def_file, 'r') as orig_file, open(filtered_file, 'w') as filt_file:
        for l in orig_file:
            if not matches_any(excl_patterns, l) or matches_any(incl_patterns, l):
                filt_file.write(l)

Filters a windows .def file based on a filter .json.

Args:
        def_file: The path to the input windows .def file.
        filter_file: The path to the filter file (JSON format).
        filtered_file: The path to the output filtered windows .def file.


def convert_tensor_tf_type_to_tflite_type(tf_type: dtypes.DType, usage: str='') -> _types_pb2.IODataType:
    mapping = {dtypes.float16: _types_pb2.FLOAT16, dtypes.float32: _types_pb2.FLOAT, dtypes.float64: _types_pb2.FLOAT64, dtypes.int8: _types_pb2.INT8, dtypes.int16: _types_pb2.INT16, dtypes.uint16: _types_pb2.UINT16, dtypes.int32: _types_pb2.INT32, dtypes.int64: _types_pb2.INT64, dtypes.uint8: _types_pb2.UINT8, dtypes.uint32: _types_pb2.UINT32, dtypes.uint64: _types_pb2.UINT64, dtypes.string: _types_pb2.STRING, dtypes.bool: _types_pb2.BOOL, dtypes.complex64: _types_pb2.COMPLEX64, dtypes.complex128: _types_pb2.COMPLEX128}
    tflite_type = mapping.get(tf_type)
    if tflite_type is None:
        raise ValueError('Unsupported TensorFlow type `{0}` provided for the {1}'.format(tf_type, usage))
    return tflite_type

Convert tensor type from tf type to tflite type.

Args:
        tf_type: TensorFlow type.
        usage: Text describing the reason for invoking this function.

Raises:
        ValueError: If `tf_type` is unsupported.

Returns:
        tflite_type: TFLite type. Refer to compiler/mlir/lite/types.proto.


def load(path, compile=True, options=None):
    metadata = saved_metadata_pb2.SavedMetadata()
    meta_graph_def = loader_impl.parse_saved_model(path).meta_graphs[0]
    object_graph_def = meta_graph_def.object_graph_def
    path_to_metadata_pb = os.path.join(path, constants.SAVED_METADATA_PATH)
    if gfile.Exists(path_to_metadata_pb):
        try:
            with gfile.GFile(path_to_metadata_pb, 'rb') as f:
                file_content = f.read()
            metadata.ParseFromString(file_content)
        except message.DecodeError as e:
            raise IOError('Cannot parse keras metadata {}: {}.'.format(path_to_metadata_pb, str(e)))
    else:
        logging.warning('SavedModel saved prior to TF 2.5 detected when loading Keras model. Please ensure that you are saving the model with model.save() or tf.keras.models.save_model(), *NOT* tf.saved_model.save(). To confirm, there should be a file named "keras_metadata.pb" in the SavedModel directory.')
        _read_legacy_metadata(object_graph_def, metadata)
    if not metadata.nodes:
        return tf_load.load(path, options=options)
    keras_loader = KerasObjectLoader(metadata, object_graph_def)
    keras_loader.load_layers(compile=compile)
    nodes_to_load = {'root': None}
    for node_id, loaded_node in keras_loader.loaded_nodes.items():
        nodes_to_load[keras_loader.get_path(node_id)] = loaded_node
    loaded = tf_load.load_partial(path, nodes_to_load, options=options)
    keras_loader.finalize_objects()
    keras_loader.del_tracking()
    model = loaded['root']
    if isinstance(model, training_lib.Model) and compile:
        training_config = model._serialized_attributes['metadata'].get('training_config', None)
        if training_config is not None:
            model.compile(**saving_utils.compile_args_from_training_config(training_config), from_serialized=True)
            saving_utils.try_build_compiled_arguments(model)
            if isinstance(model.optimizer, optimizer_v2.OptimizerV2):
                if model.optimizer.get_slot_names():
                    logging.warning('Your optimizer uses slots. Slots cannot be restored from saved_model, as a result, your model is starting with  a new initialized optimizer.')
        else:
            logging.warning('No training configuration found in save file, so the model was *not* compiled. Compile it manually.')
    if not context.executing_eagerly():
        sess = backend.get_session()
        sess.run(ops.get_collection(ops.GraphKeys.TABLE_INITIALIZERS))
    return model

Loads Keras objects from a SavedModel.

    Any Keras layer or model saved to the SavedModel will be loaded back
    as Keras objects. Other objects are loaded as regular trackable objects (same
    as `tf.saved_model.load`).

    Currently, Keras saving/loading only retains the Keras object's weights,
    losses, and call function.

    The loaded model can be re-compiled, but the original optimizer, compiled loss
    functions, and metrics are not retained. This is temporary, and `model.save`
    will soon be able to serialize compiled models.

Args:
        path: Path to SavedModel.
        compile: If true, compile the model after loading it.
        options: Optional `tf.saved_model.LoadOptions` object that specifies
        options for loading from SavedModel.

Returns:
        Object loaded from SavedModel.


def __init__(self, config: PatchTSMixerConfig):
    super().__init__()
    self.norm = PatchTSMixerNormLayer(config)
    self.self_attn = config.self_attn
    self.gated_attn = config.gated_attn
    self.mlp = PatchTSMixerMLP(in_features=config.num_patches, out_features=config.num_patches, config=config)
    if config.gated_attn:
        self.gating_block = PatchTSMixerGatedAttention(in_size=config.num_patches, out_size=config.num_patches)
    if config.self_attn:
        self.self_attn_layer = PatchTSMixerAttention(embed_dim=config.d_model, num_heads=config.self_attn_heads, dropout=config.dropout, config=config)
        self.norm_attn = PatchTSMixerNormLayer(config)

This module mixes the patch dimension.

Args:
        config (`PatchTSMixerConfig`):
            Configuration.


def split(self, path: str) -> Tuple[str, str]:
    raise NotImplementedError

Splits the given path into two parts.

    Splits the path into a pair (head, tail) such that tail contains the last
    component of the path and head contains everything up to that.

    For file-systems other than the local file-system, head should include the
    prefix.

Args:
        path: path as a string

Returns:
        a pair of path components as strings.


def data_gen_sig2() -> repr_dataset.RepresentativeDataset:
    for _ in range(4):
        yield {'conv_input': random_ops.random_uniform(shape=(1, 3, 4, 3))}

Generates tuple-style samples for signature 'sig2'.

    The first element of the tuple identifies the signature key the input data
    is for.

Yields:
        Representative sample for 'sig2'.


def add_to_collection(self, name, value) -> None:
    self._check_not_finalized()
    with self._lock:
        if name not in self._collections:
            self._collections[name] = [value]
        else:
            self._collections[name].append(value)

Stores `value` in the collection with the given `name`.

    Note that collections are not sets, so it is possible to add a value to
    a collection several times.

Args:
        name: The key for the collection. The `GraphKeys` class contains many
        standard names for collections.
        value: The value to add to the collection.


def __init__(self):
    self._last_step_outputs = {}
    self._last_step_outputs_reduce_ops = {}
    self._non_tensor_outputs = {}

Initialize an output context.

Returns:
        A context object.


def select_best_resolution(original_size: tuple, possible_resolutions: list) -> tuple:
    original_height, original_width = original_size
    best_fit = None
    max_effective_resolution = 0
    min_wasted_resolution = float('inf')
    for height, width in possible_resolutions:
        scale = min(width / original_width, height / original_height)
        downscaled_width, downscaled_height = (int(original_width * scale), int(original_height * scale))
        effective_resolution = min(downscaled_width * downscaled_height, original_width * original_height)
        wasted_resolution = width * height - effective_resolution
        if effective_resolution > max_effective_resolution or (effective_resolution == max_effective_resolution and wasted_resolution < min_wasted_resolution):
            max_effective_resolution = effective_resolution
            min_wasted_resolution = wasted_resolution
            best_fit = (height, width)
    return best_fit

Selects the best resolution from a list of possible resolutions based on the original size.

    This is done by calculating the effective and wasted resolution for each possible resolution.

    The best fit resolution is the one that maximizes the effective resolution and minimizes the wasted resolution.

Args:
        original_size (tuple):
            The original size of the image in the format (height, width).
        possible_resolutions (list):
            A list of possible resolutions in the format [(height1, width1), (height2, width2), ...].

Returns:
        tuple: The best fit resolution in the format (height, width).


def get_gradients(self, loss, params):
    params = nest.flatten(params)
    with backend.get_graph().as_default(), backend.name_scope(self._name + '/gradients'):
        grads = gradients.gradients(loss, params)
        for grad, param in zip(grads, params):
            if grad is None:
                raise ValueError('Variable {} has `None` for gradient. Please make sure that all of your ops have a gradient defined (i.e. are differentiable). Common ops without gradient: K.argmax, K.round, K.eval.'.format(param))
    return grads

Returns gradients of `loss` with respect to `params`.

    Should be used only in legacy v1 graph mode.

Args:
        loss: Loss tensor.
        params: List of variables.

Returns:
        List of gradient tensors.

Raises:
        ValueError: In case any gradient cannot be computed (e.g. if gradient
        function not implemented).


def __init__(self, input_ids: torch.LongTensor, assistant_model: 'PreTrainedModel', generation_config: 'GenerationConfig', model_kwargs: Dict, inputs_tensor: Optional[torch.Tensor]=None, logits_processor: 'LogitsProcessorList'=None):
    super().__init__(input_ids=input_ids, assistant_model=assistant_model, generation_config=generation_config, model_kwargs=model_kwargs, inputs_tensor=inputs_tensor, logits_processor=logits_processor)
    self.assistant_early_exit = self.generation_config.assistant_early_exit
    self.generation_config.assistant_early_exit = None

`CandidateGenerator` class to be used for assisted generation and speculative decoding. This class generates
    candidates through the use of **the model itself**, exiting early. Can only be used with models that support early
    exit, e.g., `facebook/layerskip-llama3.2-1B`.

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)
        assistant_model (`PreTrainedModel`):
            The original model. This model must support early exit (i.e. is trained to compute logits in earlier
            layers).
        generation_config (`~generation.GenerationConfig`, *optional*):
            The generation configuration to be used as base parametrization for the generation call.
        logits_processor (`LogitsProcessorList`):
            An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`]
            used to modify the prediction scores of the language modeling head applied at each generation step.
        model_kwargs (`Dict`):
            The keyword arguments that will be passed to the main model, and are used as base inputs for the assistant
            model as well.
        inputs_tensor (`torch.Tensor`, *optional*):
            The model input tensor. In encoder-decoder models, this is the encoder input.


def build_results(self, session, tensor_values):
    full_values = []
    assert len(self._final_fetches) == len(tensor_values)
    i = 0
    j = 0
    for is_op in self._ops:
        if is_op:
            full_values.append(None)
        else:
            if self._fetches[i].ref() in self._feed_handles:
                value = self._feed_handles[self._fetches[i].ref()].eval()
            else:
                value = self._feeds.get(self._fetches[i].ref())
            if value is None:
                value = tensor_values[j]
                j += 1
            dtype = self._fetch_handles.get(self._fetches[i].ref())
            if dtype:
                full_values.append(session_ops.TensorHandle(value, dtype, session))
            else:
                full_values.append(value)
            i += 1
    assert j == len(tensor_values)
    return self._fetch_mapper.build_results(full_values)

Build results matching the original fetch shape.

    `tensor_values` must be a list of the same length as
    the one returned by `fetches()`, and holding the requested
    fetch values.

    This method builds a struct with the same shape as the original `fetches`
    passed to the constructor, in which the fetches are replaced by their
    fetched value.

Args:
        session: The enclosing session.  Used for tensor handles.
        tensor_values: List of values matching the list returned by fetches().

Returns:
        A structure of the same shape as the original `fetches` argument but
        containing tensors or None (for fetched ops).


def __init__(self, original_lm_head, assistant_overlap_token_ids):
    super().__init__()
    self.pruned_lm_head = prune_linear_layer(original_lm_head, assistant_overlap_token_ids).to(original_lm_head.weight.dtype)

A class to prune and reindex the language model head.

    This class prunes the language model head to only include the specified token IDs and reindexes the logits
    to map back to the original vocabulary.

Args:
        original_lm_head (nn.Module): The original language model head.
        token_ids (list[int]): The list of token IDs to keep.


def _is_executing(self) -> bool:
    executors = self._executor.transform_executor_services.executors
    if not executors:
        return False
    for transform_executor in executors:
        if not transform_executor.blocked:
            return True
    return False

Checks whether the job is still executing.

Returns:
        True if there is at least one non-blocked TransformExecutor active.


def forward(self, inputs: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None, head_mask: Optional[torch.Tensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, labels: Optional[torch.Tensor]=None, return_dict: Optional[bool]=None) -> Union[Tuple, PerceiverClassifierOutput]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    loss = None
    if labels is not None:
        raise NotImplementedError('Optical flow training is not yet supported')
    outputs = self.perceiver(inputs=inputs, attention_mask=attention_mask, head_mask=head_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    logits = outputs.logits if return_dict else outputs[0]
    if not return_dict:
        output = (logits,) + outputs[2:]
        return (loss,) + output if loss is not None else output
    return PerceiverClassifierOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions, cross_attentions=outputs.cross_attentions)

inputs (`torch.FloatTensor`):
    Inputs to the perceiver. Can be anything: images, text, audio, video, etc.
    labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
    Labels for computing the optical flow loss. Indices should be in `[0, ..., config.num_labels - 1]`.

Example:
    ```python
    >>> from transformers import PerceiverForOpticalFlow
    >>> import torch

    >>> model = PerceiverForOpticalFlow.from_pretrained("deepmind/optical-flow-perceiver")

    >>> # in the Perceiver IO paper, the authors extract a 3 x 3 patch around each pixel,
    >>> # leading to 3 x 3 x 3 = 27 values for each pixel (as each pixel also has 3 color channels)
    >>> # patches have shape (batch_size, num_frames, num_channels, height, width)
    >>> # the authors train on resolutions of 368 x 496
    >>> patches = torch.randn(1, 2, 27, 368, 496)
    >>> outputs = model(inputs=patches)
    >>> logits = outputs.logits
    >>> list(logits.shape)
    [1, 368, 496, 2]
    ```


def fn(x: list[list[Union[str, int]]]):
    return x

Test function

Args:
        x: The input


def __init__(self, total: int, prefix: Optional[str]=None, leave: bool=True, parent: Optional['NotebookTrainingTracker']=None, width: int=300):
    self.total = total
    self.prefix = '' if prefix is None else prefix
    self.leave = leave
    self.parent = parent
    self.width = width
    self.last_value = None
    self.comment = None
    self.output = None
    self.value = None
    self.label = None
    if 'VSCODE_PID' in os.environ:
        self.update_every = 0.5

A progress par for display in a notebook.

    Class attributes (overridden by derived classes)

        - **warmup** (`int`) -- The number of iterations to do at the beginning while ignoring `update_every`.
        - **update_every** (`float`) -- Since calling the time takes some time, we only do it every presumed
        `update_every` seconds. The progress bar uses the average time passed up until now to guess the next value
        for which it will call the update.

Args:
        total (`int`):
            The total number of iterations to reach.
        prefix (`str`, *optional*):
            A prefix to add before the progress bar.
        leave (`bool`, *optional*, defaults to `True`):
            Whether or not to leave the progress bar once it's completed. You can always call the
            [`~utils.notebook.NotebookProgressBar.close`] method to make the bar disappear.
        parent ([`~notebook.NotebookTrainingTracker`], *optional*):
            A parent object (like [`~utils.notebook.NotebookTrainingTracker`]) that spawns progress bars and handle
            their display. If set, the object passed must have a `display()` method.
        width (`int`, *optional*, defaults to 300):
            The width (in pixels) that the bar will take.

Example:
    ```python
    import time

    pbar = NotebookProgressBar(100)
    for val in range(100):
        pbar.update(val)
        time.sleep(0.07)
    pbar.update(100)
    ```


def create_and_fill_np_array(start_or_end_logits, dataset, max_len):
    step = 0
    logits_concat = np.full((len(dataset), max_len), -100, dtype=np.float32)
    for i, output_logit in enumerate(start_or_end_logits):
        batch_size = output_logit.shape[0]
        cols = output_logit.shape[1]
        if step + batch_size < len(dataset):
            logits_concat[step:step + batch_size, :cols] = output_logit
        else:
            logits_concat[step:, :cols] = output_logit[:len(dataset) - step]
        step += batch_size
    return logits_concat

Create and fill numpy array of size len_of_validation_data * max_length_of_output_tensor

Args:
        start_or_end_logits(:obj:`tensor`):
            This is the output predictions of the model. We can only enter either start or end logits.
        eval_dataset: Evaluation dataset
        max_len(:obj:`int`):
            The maximum length of the output tensor. ( See the model.eval() part for more details )


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_pred, y_true], sample_weight = metrics_utils.ragged_assert_compatible_and_get_flat_values([y_pred, y_true], sample_weight)
    y_pred, y_true = losses_utils.squeeze_or_expand_dimensions(y_pred, y_true)
    y_pred, self.normalizer = losses_utils.remove_squeezable_dimensions(y_pred, self.normalizer)
    y_pred.shape.assert_is_compatible_with(y_true.shape)
    relative_errors = math_ops.div_no_nan(math_ops.abs(y_true - y_pred), self.normalizer)
    return super(MeanRelativeError, self).update_state(relative_errors, sample_weight=sample_weight)

Accumulates metric statistics.

Args:
        y_true: The ground truth values.
        y_pred: The predicted values.
        sample_weight: Optional weighting of each example. Defaults to 1. Can be a
        `Tensor` whose rank is either 0, or the same rank as `y_true`, and must
        be broadcastable to `y_true`.

Returns:
        Update op.


def gradient_tensor(self, x_tensor):
    x_tensor_name = self._get_tensor_name(x_tensor)
    if x_tensor_name not in self._gradient_tensors:
        raise LookupError('This GradientsDebugger has not received any gradient tensor for x-tensor %s' % x_tensor_name)
    return self._gradient_tensors[x_tensor_name]

Get the gradient tensor of an x-tensor.

Args:
        x_tensor: (`tf.Tensor`, `tf.Variable` or `str`) The x-tensor object or its
        name. x-tensor refers to the independent `tf.Tensor`, i.e., the tensor
        on the denominator of the differentiation.

Returns:
        If found, the gradient tensor.

Raises:
        TypeError: If `x_tensor` is not a `tf.Tensor`, `tf.Variable` or `str`.
        LookupError: If the `x_tensor` has not been registered with a gradient
        tensor.


def _is_variable(node_def: node_def_pb2.NodeDef) -> bool:
    return node_def.op == 'VarHandleOp'

Determines whether `node_def` is a variable node.

Args:
        node_def: `NodeDef` to test whether it is a variable or not.

Returns:
        Returns True if it is a variable.


def _collect_layer_statistics(self) -> Dict[str, Dict[str, float]]:
    layer_statistics = collections.defaultdict(lambda: collections.defaultdict(list))
    initialize = True
    for tensor_data in self._data_gen():
        self._set_input_tensors(self._quant_interpreter, tensor_data, initialize)
        initialize = False
        self._quant_interpreter.invoke()
        for tensor_detail in self._get_numeric_verify_tensor_details():
            tensor_name = tensor_detail['name']
            diffs = self._quant_interpreter.get_tensor(tensor_detail['index'])
            for metric_name, metric_fn in self._layer_debug_metrics.items():
                layer_statistics[tensor_name][metric_name].append(metric_fn(diffs))
        if self._debug_options.layer_direct_compare_metrics is not None:
            for tensor_detail in self._get_numeric_verify_tensor_details():
                tensor_name = tensor_detail['name']
                op_idx = self._defining_op[tensor_detail['index']]
                op_detail = self._quant_interpreter._get_op_details(op_idx)
                q_idx, f_idx = op_detail['inputs']
                quant_input_detail = self._quant_interpreter._get_tensor_details(q_idx, subgraph_index=0)
                for metric_name, metric_fn in self._debug_options.layer_direct_compare_metrics.items():
                    layer_statistics[tensor_name][metric_name].append(metric_fn(self._quant_interpreter.get_tensor(f_idx), self._quant_interpreter.get_tensor(q_idx), quant_input_detail['quantization_parameters']['scales'][0], quant_input_detail['quantization_parameters']['zero_points'][0]))
    for metrics in layer_statistics.values():
        for metric_name in metrics:
            metrics[metric_name] = np.nanmean(metrics[metric_name])
    return layer_statistics

Collects layer statistics by applying layer debug metrics.

    For all data from the given RepresentativeDataset, collect statistics per
    example by getting the NumericVerify op results in _quant_interpreter
    and calculating layer debug metrics on the results.

Returns:
        aggregated per-layer statistics of NumericVerify results.
        {layer_name: {metric_name: metric}}


def trace(self, name='trace'):
    with self._name_scope(name):
        return self._trace()

Trace of the linear operator, equal to sum of `self.diag_part()`.

    If the operator is square, this is also the sum of the eigenvalues.

Args:
        name:  A name for this `Op`.

Returns:
        Shape `[B1,...,Bb]` `Tensor` of same `dtype` as `self`.


def build_relative_position(query_size, key_size, bucket_size=-1, max_position=-1, device=None):
    q_ids = torch.arange(0, query_size, device=device)
    k_ids = torch.arange(0, key_size, device=device)
    rel_pos_ids = q_ids[:, None] - k_ids[None, :]
    if bucket_size > 0 and max_position > 0:
        rel_pos_ids = make_log_bucket_position(rel_pos_ids, bucket_size, max_position)
    rel_pos_ids = rel_pos_ids.to(torch.long)
    rel_pos_ids = rel_pos_ids[:query_size, :]
    rel_pos_ids = rel_pos_ids.unsqueeze(0)
    return rel_pos_ids

Build relative position according to the query and key

    We assume the absolute position of query \(P_q\) is range from (0, query_size) and the absolute position of key
    \(P_k\) is range from (0, key_size), The relative positions from query to key is \(R_{q \rightarrow k} = P_q -
    P_k\)

Args:
        query_size (int): the length of query
        key_size (int): the length of key
        bucket_size (int): the size of position bucket
        max_position (int): the maximum allowed absolute position
        device (`torch.device`): the device on which tensors will be created.

Returns:
        `torch.LongTensor`: A tensor with shape [1, query_size, key_size]


def _with_dependencies(self, dependencies):
    new_row_splits = control_flow_ops.with_dependencies(dependencies, self._row_splits)
    return RowPartition(row_splits=new_row_splits, row_lengths=self._row_lengths, value_rowids=self._value_rowids, nrows=self._nrows, uniform_row_length=self._uniform_row_length, internal=_row_partition_factory_key)

Returns a new RowPartition equal to self with control dependencies.

    Specifically, self._row_splits is gated by the given control dependencies.
    Used to add sanity checks to the constructors.

Args:
        dependencies: a list of tensors to use as dependencies.

Returns:
        A new RowPartition object.


def assert_true(expr, msg, extras=None):
    if not expr:
        fail(msg, extras)

Assert an expression evaluates to true, otherwise fail the test.

Args:
        expr: The expression that is evaluated.
        msg: A string explaining the details in case of failure.
        extras: An optional field for extra information to be included in
        test result.


def calendar_day_of_week(self: EventSetOrNode, tz: Union[str, float, int]=0) -> EventSetOrNode:
    from temporian.core.operators.calendar.day_of_week import calendar_day_of_week
    return calendar_day_of_week(self, tz)

Obtains the day of the week the timestamps in an
    [`EventSet`][temporian.EventSet]'s sampling are in.

    Features in the input are ignored, only the timestamps are used and
    they must be unix timestamps (`is_unix_timestamp=True`).

    Output feature contains numbers from 0 (Monday) to 6 (Sunday).

    By default, the timezone is UTC unless the `tz` argument is specified,
    as an offset in hours or a timezone name. See
    [`EventSet.calendar_hour()`][temporian.EventSet.calendar_hour] for an
    example using timezones.

    Usage example:
        ```python
        >>> a = tp.event_set(
        ...    timestamps=["2023-06-19", "2023-06-21", "2023-06-25", "2023-07-03"],
        ... )
        >>> b = a.calendar_day_of_week()
        >>> b
        indexes: ...
        features: [('calendar_day_of_week', int32)]
        events:
            (4 events):
                timestamps: [...]
                'calendar_day_of_week': [0  2  6  0]
        ...

        ```

Args:
        tz: timezone name (see `pytz.all_timezones`) or UTC offset in hours.

Returns:
        EventSet with a single feature with the day of the week each timestamp
            in `sampling` belongs to.


def __init__(self, token_ids: List[int]):
    super(Constraint, self).__init__()
    if not isinstance(token_ids, list) or len(token_ids) == 0:
        raise ValueError(f'`token_ids` has to be a non-empty list, but is {token_ids}.')
    if any((not isinstance(token_id, int) or token_id < 0 for token_id in token_ids)):
        raise ValueError(f'Each list in `token_ids` has to be a list of positive integers, but is {token_ids}.')
    self.token_ids = token_ids
    self.seqlen = len(self.token_ids)
    self.fulfilled_idx = -1
    self.completed = False

[`Constraint`] enforcing that an ordered sequence of tokens is included in the output.

Args:
        token_ids (`List[int]`):
            The id of the token that must be generated by the output.


def replace_flat_tensors_for_gradients(xs, flat_grads):
    xs_structure = [_get_tensors_for_gradient(x) for x in xs]
    grads = nest.pack_sequence_as(xs_structure, flat_grads)
    return [_replace_tensors_for_gradient(x, grad) for x, grad in zip(xs, grads)]

Replaces Tensors that should be differentiated in `xs` with `flat_grads`.

Args:
        xs: A list of `Tensor`s or `CompositeTensor`s.
        flat_grads: A list of `Tensor`.

Returns:
        A list of `Tensor` or `CompositeTensor`.


def _get_bounding_box(self, box: 'torch.Tensor') -> Dict[str, int]:
    if self.framework != 'pt':
        raise ValueError('The ObjectDetectionPipeline is only available in PyTorch.')
    xmin, ymin, xmax, ymax = box.int().tolist()
    bbox = {'xmin': xmin, 'ymin': ymin, 'xmax': xmax, 'ymax': ymax}
    return bbox

Turns list [xmin, xmax, ymin, ymax] into dict { "xmin": xmin, ... }

Args:
        box (`torch.Tensor`): Tensor containing the coordinates in corners format.

Returns:
        bbox (`Dict[str, int]`): Dict containing the coordinates in corners format.


def update_pcollection_stats(self, pcollection_stats):
    edge_dict = {}
    for pcoll_id, stats in pcollection_stats.items():
        attrs = {}
        pcoll_list = stats['sample']
        if pcoll_list:
            attrs['label'] = format_sample(pcoll_list, 1)
            attrs['labeltooltip'] = format_sample(pcoll_list, 10)
        else:
            attrs['label'] = '?'
        edge_dict[pcoll_id] = attrs
    self._update_graph(edge_dict=edge_dict)

Updates PCollection stats.

Args:
        pcollection_stats: (dict of dict) maps PCollection IDs to informations. In
        particular, we only care about the field 'sample' which should be a
        the PCollection result in as a list.


def before_starting_server(self):
    self._validate_snippet_app_on_device()
    self._disable_hidden_api_blocklist()

Performs the preparation steps before starting the remote server.

    This function performs following preparation steps:
    * Validate that the Mobly Snippet app is available on the device.
    * Disable hidden api blocklist if necessary and possible.

Raises:
        errors.ServerStartPreCheckError: if the server app is not installed
        for the current user.


def container_type_mismatch(self, stack, cls, mutations, name):
    details = f'Container: {self._pp.print_generic_type(cls)}\n'
    allowed_contained = ''
    new_contained = ''
    for formal in cls.formal_type_parameters.keys():
        if formal in mutations:
            params, values, _ = mutations[formal]
            allowed_content = self._pp.print_type_of_instance(cls.get_formal_type_parameter(formal))
            new_content = self._pp.join_printed_types(sorted((self._pp.print_type(v) for v in set(values.data) - set(params.data))))
            allowed_contained += f'  {formal}: {allowed_content}\n'
            new_contained += f'  {formal}: {new_content}\n'
    annotation = self._pp.print_type_of_instance(cls)
    details += 'Allowed contained types (from annotation %s):\n%sNew contained types:\n%s' % (annotation, allowed_contained, new_contained)
    suffix = '' if name is None else ' for ' + name
    err_msg = f'New container type{suffix} does not match type annotation'
    self.error(stack, err_msg, details=details)

Invalid combination of annotation and mutation.

Args:
        stack: the frame stack
        cls: the container type
        mutations: a dict of {parameter name: (annotated types, new types)}
        name: the variable name (or None)


def process(self, feed_item):
    item = self.get(feed_item)
    if item:
        self._process_update(item, feed_item)
        self._clean(item)
        self._update(item, feed_item)
    else:
        new_item = self._process_new(feed_item)
        self._clean(new_item)
        item = self._insert(new_item, feed_item)
        if self._id_field and feed_item.get(self._id_field, '').startswith('ext'):
            store.map(self._entity, feed_item.get(self._id_field), item['id'])
            store.set(self._entity, [feed_item[self._id_field]], item)
        if self._search_field and feed_item.get(self._search_field, ''):
            store.map(self._entity, feed_item.get(self._search_field), item['id'])
            store.set(self._entity, [feed_item[self._search_field]], item)
    if item:
        feed_item[self._id_field] = item['id']
        store.set(self._entity, [item['id']], item)
    self._post_process(feed_item, item)
    return item

Processes a Bulkdozer feed item.

    This method identifies if the item needs to be inserted or updated, cleans
    it, performs the CM operations required, and update the feed item with newly
    created ids and name lookups so that the feed can be updated.

Args:
        feed_item: Bulkdozer feed item to process.

Returns:
        Newly created or updated CM object.


def _execute_bundle(self, runner_execution_context: execution.FnApiRunnerExecutionContext, bundle_context_manager: execution.BundleContextManager, bundle_input: DataInput) -> beam_fn_api_pb2.InstructionResponse:
    worker_handler_manager = runner_execution_context.worker_handler_manager
    worker_handler_manager.register_process_bundle_descriptor(bundle_context_manager.process_bundle_descriptor)
    bundle_manager = self._get_bundle_manager(bundle_context_manager)
    last_result, deferred_inputs, newly_set_timers, watermark_updates = self._run_bundle(runner_execution_context, bundle_context_manager, bundle_input, bundle_context_manager.stage_data_outputs, bundle_context_manager.stage_timer_outputs, bundle_manager)
    for pc_name, watermark in watermark_updates.items():
        _BUNDLE_LOGGER.debug('Update: %s %s', pc_name, watermark)
        runner_execution_context.watermark_manager.set_pcoll_watermark(pc_name, watermark)
    if deferred_inputs:
        assert runner_execution_context.watermark_manager.get_stage_node(bundle_context_manager.stage.name).output_watermark() < timestamp.MAX_TIMESTAMP, 'wrong timestamp for %s. ' % runner_execution_context.watermark_manager.get_stage_node(bundle_context_manager.stage.name)
        runner_execution_context.queues.ready_inputs.enque((bundle_context_manager.stage.name, DataInput(deferred_inputs, {})))
    self._enqueue_set_timers(runner_execution_context, bundle_context_manager, newly_set_timers, bundle_input)
    if not deferred_inputs and (not newly_set_timers):
        for _, output_pc in bundle_context_manager.stage_data_outputs.items():
            _, update_output_pc = translations.split_buffer_id(output_pc)
            runner_execution_context.watermark_manager.get_pcoll_node(update_output_pc).set_produced_watermark(timestamp.MAX_TIMESTAMP)
    data_side_input = runner_execution_context.side_input_descriptors_by_stage.get(bundle_context_manager.stage.name, {})
    runner_execution_context.commit_side_inputs_to_state(data_side_input)
    buffers_to_clean = set()
    known_consumers = set()
    for transform_id, buffer_id in bundle_context_manager.stage_data_outputs.items():
        for consuming_stage_name, consuming_transform in runner_execution_context.buffer_id_to_consumer_pairs.get(buffer_id, []):
            buffer = runner_execution_context.pcoll_buffers.get(buffer_id, None)
            if buffer_id in runner_execution_context.pcoll_buffers and buffer_id not in buffers_to_clean:
                buffers_to_clean.add(buffer_id)
            elif buffer and buffer_id in buffers_to_clean:
                runner_execution_context.pcoll_buffers[buffer_id] = buffer.copy()
                buffer = runner_execution_context.pcoll_buffers[buffer_id]
            if buffer_id not in runner_execution_context.pcoll_buffers:
                buffer = bundle_context_manager.get_buffer(buffer_id, transform_id)
                runner_execution_context.pcoll_buffers[buffer_id] = buffer
                buffers_to_clean.add(buffer_id)
            if (consuming_stage_name, consuming_transform, buffer_id) in known_consumers:
                continue
            else:
                known_consumers.add((consuming_stage_name, consuming_transform, buffer_id))
            runner_execution_context.queues.watermark_pending_inputs.enque(((consuming_stage_name, timestamp.MAX_TIMESTAMP), DataInput({consuming_transform: buffer}, {})))
    for bid in buffers_to_clean:
        if bid in runner_execution_context.pcoll_buffers:
            del runner_execution_context.pcoll_buffers[bid]
    return last_result

Execute a bundle end-to-end.

Args:
        runner_execution_context (execution.FnApiRunnerExecutionContext): An
        object containing execution information for the pipeline.
        bundle_context_manager (execution.BundleContextManager): A description of
        the stage to execute, and its context.
        bundle_input: The set of buffers to input into this bundle


def __init__(self, resume_delay=0):
    self.resume_delay = resume_delay

Initializes a ProcessContinuation object.

Args:
        resume_delay: indicates the minimum time, in seconds, that should elapse
        before re-invoking process() method for resuming the invocation of the
        current element.


def parse_args(test: ArgList=None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(prog='budoux', formatter_class=lambda prog: BudouxHelpFormatter(prog, **{'width': shutil.get_terminal_size(fallback=(120, 50)).columns, 'max_help_position': 30}), description=textwrap.dedent('        BudouX is the successor to Budou,\n        the machine learning powered line break organizer tool.'), epilog='\n- '.join(['supported languages of `-l`, `--lang`:', *langs.keys()]))
    parser.add_argument('text', metavar='TXT', nargs='?', type=str, help='text')
    parser.add_argument('-H', '--html', action='store_true', help='HTML mode')
    model_select_group = parser.add_mutually_exclusive_group()
    model_select_group.add_argument('-m', '--model', metavar='JSON', type=check_file, default=check_lang('ja'), help='custom model file path')
    model_select_group.add_argument('-l', '--lang', metavar='LANG', type=check_lang, help='language of custom model')
    parser.add_argument('-s', '--sep', metavar='STR', type=str, default='\n', help='output phrase separator in TEXT mode')
    parser.add_argument('-d', '--delim', metavar='STR', type=str, default='---', help='output sentence delimiter in TEXT mode')
    parser.add_argument('-V', '--version', action='version', version='%(prog)s {}'.format(budoux.__version__))
    if test is not None:
        return parser.parse_args(test)
    else:
        return parser.parse_args()

Parse commandline arguments.

Args:
        test (typing.Optional[typing.List[str]], optional): Commandline args for testing. Defaults to None.

Returns:
        argparse.Namespace: Parsed data of args.


def parse_indices(indices_string):
    indices_string = re.sub('\\s+', '', indices_string)
    if indices_string.startswith('[') and indices_string.endswith(']'):
        indices_string = indices_string[1:-1]
    return [int(element) for element in indices_string.split(',')]

Parse a string representing indices.

    For example, if the input is "[1, 2, 3]", the return value will be a list of
    indices: [1, 2, 3]

Args:
        indices_string: (str) a string representing indices. Can optionally be
        surrounded by a pair of brackets.

Returns:
        (list of int): Parsed indices.


def deregister_context(self, context_words):
    for context_word in context_words:
        if context_word not in self._comp_dict:
            raise KeyError('Cannot deregister unregistered context word "%s"' % context_word)
    for context_word in context_words:
        del self._comp_dict[context_word]

Deregister a list of context words.

Args:
        context_words: A list of context words to deregister, as a list of str.

Raises:
        KeyError: if there are word(s) in context_words that do not correspond
        to any registered contexts.


def as_videos(content: ProcessorContentTypes, *, ignore_unsupported_types: bool=False) -> list[ProcessorPart]:
    return _as_format_helper(content, lambda mime: mime.startswith('video/'), ignore_unsupported_types)

Returns the video parts from the content.

Args:
        content: Input content.
        ignore_unsupported_types: By default if content contains non-video parts a
        ValueError would be raised. This argument allows ingoring such parts.

Returns:
        A list of video parts.


def _get_func_name():
    return tf_inspect.stack()[1][3]

Get the name of current function.

Returns:
        String that is the name of current function.


def call(self, input_ids: Optional[tf.Tensor]=None, position_ids: Optional[tf.Tensor]=None, token_type_ids: Optional[tf.Tensor]=None, inputs_embeds: Optional[tf.Tensor]=None, past_key_values_length=0, training: bool=False) -> tf.Tensor:
    if input_ids is None and inputs_embeds is None:
        raise ValueError('Need to provide either `input_ids` or `input_embeds`.')
    if input_ids is not None:
        check_embeddings_within_bounds(input_ids, self.config.vocab_size)
        inputs_embeds = tf.gather(params=self.weight, indices=input_ids)
    input_shape = shape_list(inputs_embeds)[:-1]
    if token_type_ids is None:
        token_type_ids = tf.fill(dims=input_shape, value=0)
    if position_ids is None:
        position_ids = tf.expand_dims(tf.range(start=past_key_values_length, limit=input_shape[1] + past_key_values_length), axis=0)
    position_embeds = tf.gather(params=self.position_embeddings, indices=position_ids)
    token_type_embeds = tf.gather(params=self.token_type_embeddings, indices=token_type_ids)
    final_embeddings = inputs_embeds + position_embeds + token_type_embeds
    final_embeddings = self.LayerNorm(inputs=final_embeddings)
    final_embeddings = self.dropout(inputs=final_embeddings, training=training)
    return final_embeddings

Applies embedding based on inputs tensor.

Returns:
        final_embeddings (`tf.Tensor`): output embedding tensor.


def create_detector(model_uri: str, **kwargs) -> OfflineDetector:
    model_handler = KeyedModelHandler(PyODModelHandler(model_uri=model_uri)).with_postprocess_fn(OfflineDetector.score_prediction_adapter)
    m = model_handler.load_model()
    assert isinstance(m, PyODBaseDetector)
    threshold = float(m.threshold_)
    detector = OfflineDetector(model_handler, threshold_criterion=FixedThreshold(threshold), **kwargs)
    return detector

A utility function to create OfflineDetector for a PyOD model.

    **NOTE:** This API and its implementation are currently under active
    development and may not be backward compatible.

Args:
        model_uri: The URI specifying the location of the pickled PyOD model.
        **kwargs: Additional keyword arguments.


def is_registered(self, prefix):
    return self._resolve_prefix(prefix) is not None

Test if a command prefix or its alias is has a registered handler.

Args:
        prefix: A prefix or its alias, as a str.

Returns:
        True iff a handler is registered for prefix.