def _convert_from_saved_model(self, graph_def):
    self._save_conversion_params_metric(graph_def)
    quant_mode = QuantizationMode(self.optimizations, self.target_spec, self.representative_dataset, graph_def, self._experimental_disable_per_channel, self.experimental_new_dynamic_range_quantizer, self._experimental_low_bit_qat, self._experimental_full_integer_quantization_bias_type, self._experimental_variable_quantization, self._experimental_strict_qdq)
    self._validate_inference_input_output_types(quant_mode)
    converter_kwargs = {'enable_tflite_resource_variables': self.experimental_enable_resource_variables}
    converter_kwargs.update(self._get_base_converter_args())
    converter_kwargs.update(quant_mode.converter_flags())
    result = _convert_saved_model(**converter_kwargs)
    return self._optimize_tflite_model(result, quant_mode, _build_conversion_flags(**converter_kwargs).debug_options, quant_io=self.experimental_new_quantizer)

Helper method that converts saved model.

Args:
  graph_def: GraphDef object for the model, used only for stats.

Returns:
  The converted TFLite model.


def __init__(self, error_name, error_id, error_msg, stack_patterns):
    self.error_name = error_name
    self.error_id = error_id
    self.error_msg = error_msg
    self._stack_patterns = stack_patterns

Create a ParserError that matches against any of the |stack_patterns|.

Args:
 error_name: A short, human readable name for the error,
     using lowercase-with-dashes-format.
 error_id: An integer to identify a specific error:
     100s: Lexer errors.
     200s: Low level parsing errors.
     300s: High level parsing errors.
 error_msg: A message to display with this error that describes
     clearly what caused the error.
 stack_patterns: A list of "stack patterns", where each stack pattern
     is a list of strings corresponding to symbols on the parser's symbol
     stack at the time it errored out. The string values for the symbols
     can match essentially any terminal or non-terminal symbol used in the
     grammar from parser.py.
     Examples: ['TRANSITION', 'NAME', 'params', '=']
     (or None to match against any symbol stack).

Returns:
  ParserError that matches against |stack_patterns|.


def compute(self, batch_values, accumulator=None):
    pass

Compute a step in this computation, returning a new accumulator.

This method computes a step of the computation described by this Combiner.
If an accumulator is passed, the data in that accumulator is also used; so
compute(batch_values) results in f(batch_values), while
compute(batch_values, accumulator) results in
merge(f(batch_values), accumulator).

Args:
  batch_values: A list of ndarrays representing the values of the inputs for
    this step of the computation.
  accumulator: the current accumulator. Can be None.

Returns:
  An accumulator that includes the passed batch of inputs.


def _parse_interval(interval_str):
    interval_str = interval_str.strip()
    if interval_str.startswith('<='):
        if _NUMBER_PATTERN.match(interval_str[2:].strip()):
            return Interval(start=None, start_included=False, end=interval_str[2:].strip(), end_included=True)
        else:
            raise ValueError("Invalid value string after <= in '%s'" % interval_str)
    if interval_str.startswith('<'):
        if _NUMBER_PATTERN.match(interval_str[1:].strip()):
            return Interval(start=None, start_included=False, end=interval_str[1:].strip(), end_included=False)
        else:
            raise ValueError("Invalid value string after < in '%s'" % interval_str)
    if interval_str.startswith('>='):
        if _NUMBER_PATTERN.match(interval_str[2:].strip()):
            return Interval(start=interval_str[2:].strip(), start_included=True, end=None, end_included=False)
        else:
            raise ValueError("Invalid value string after >= in '%s'" % interval_str)
    if interval_str.startswith('>'):
        if _NUMBER_PATTERN.match(interval_str[1:].strip()):
            return Interval(start=interval_str[1:].strip(), start_included=False, end=None, end_included=False)
        else:
            raise ValueError("Invalid value string after > in '%s'" % interval_str)
    if not interval_str.startswith(('[', '(')) or not interval_str.endswith((']', ')')):
        raise ValueError('Invalid interval format: %s. Valid formats are: [min, max], (min, max), <max, >min' % interval_str)
    interval = interval_str[1:-1].split(',')
    if len(interval) != 2:
        raise ValueError('Incorrect interval format: %s. Interval should specify two values: [min, max] or (min, max).' % interval_str)
    start_item = interval[0].strip()
    if not _NUMBER_PATTERN.match(start_item):
        raise ValueError("Invalid first item in interval: '%s'" % start_item)
    end_item = interval[1].strip()
    if not _NUMBER_PATTERN.match(end_item):
        raise ValueError("Invalid second item in interval: '%s'" % end_item)
    return Interval(start=start_item, start_included=interval_str[0] == '[', end=end_item, end_included=interval_str[-1] == ']')

Convert a human-readable interval to a tuple of start and end value.

Args:
  interval_str: (`str`) A human-readable str representing an interval
    (e.g., "[1M, 2M]", "<100k", ">100ms"). The items following the ">", "<",
    ">=" and "<=" signs have to start with a number (e.g., 3.0, -2, .98).
    The same requirement applies to the items in the parentheses or brackets.

Returns:
  Interval object where start or end can be None
  if the range is specified as "<N" or ">N" respectively.

Raises:
  ValueError: if the input is not valid.


def prediction_step(self, model: nn.Module, inputs: dict[str, Union[torch.Tensor, Any]], prediction_loss_only: bool, ignore_keys: Optional[list[str]]=None) -> tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
    inputs = self._prepare_inputs(inputs)
    gen_kwargs = {'max_length': self.data_args.val_max_target_length if self.data_args is not None else self.config.max_length, 'num_beams': self.data_args.eval_beams if self.data_args is not None else self.config.num_beams}
    if self.args.predict_with_generate and (not self.args.prediction_loss_only):
        generated_tokens = self.model.generate(inputs['input_ids'], attention_mask=inputs['attention_mask'], **gen_kwargs)
        if generated_tokens.shape[-1] < gen_kwargs['max_length']:
            generated_tokens = self._pad_tensors_to_max_len(generated_tokens, gen_kwargs['max_length'])
    labels = inputs.pop('labels')
    with torch.no_grad():
        loss, logits = self._compute_loss(model, inputs, labels)
    loss = loss.mean().detach()
    if self.args.prediction_loss_only:
        return (loss, None, None)
    logits = generated_tokens if self.args.predict_with_generate else logits
    if labels.shape[-1] < gen_kwargs['max_length']:
        labels = self._pad_tensors_to_max_len(labels, gen_kwargs['max_length'])
    return (loss, logits, labels)

Perform an evaluation step on :obj:`model` using obj:`inputs`.

Subclass and override to inject custom behavior.

Args:
    model (:obj:`nn.Module`):
        The model to evaluate.
    inputs (:obj:`Dict[str, Union[torch.Tensor, Any]]`):
        The inputs and targets of the model.

        The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
        argument :obj:`labels`. Check your model's documentation for all accepted arguments.
    prediction_loss_only (:obj:`bool`):
        Whether or not to return the loss only.

Return:
    Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
    A tuple with the loss, logits and labels (each being optional).


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

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

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

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

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

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


def forward(self, hidden_states: List[torch.Tensor], patch_height: Optional[int]=None, patch_width: Optional[int]=None, prompt_depth: Optional[torch.Tensor]=None) -> List[torch.Tensor]:
    if not isinstance(hidden_states, (tuple, list)):
        raise TypeError('hidden_states should be a tuple or list of tensors')
    if len(hidden_states) != len(self.config.neck_hidden_sizes):
        raise ValueError('The number of hidden states should be equal to the number of neck hidden sizes.')
    hidden_states = self.reassemble_stage(hidden_states, patch_height, patch_width)
    features = [self.convs[i](feature) for i, feature in enumerate(hidden_states)]
    output = self.fusion_stage(features, prompt_depth=prompt_depth)
    return output

Args:
    hidden_states (`List[torch.FloatTensor]`, each of shape `(batch_size, sequence_length, hidden_size)` or `(batch_size, hidden_size, height, width)`):
        List of hidden states from the backbone.


def constant(value, dtype=None, shape=None, name='Const') -> Union[ops.Operation, ops._EagerTensorBase]:
    return _constant_impl(value, dtype, shape, name, verify_shape=False, allow_broadcast=True)

Creates a constant tensor from a tensor-like object.

Note: All eager `tf.Tensor` values are immutable (in contrast to
`tf.Variable`). There is nothing especially _constant_ about the value
returned from `tf.constant`. This function is not fundamentally different from
`tf.convert_to_tensor`. The name `tf.constant` comes from the `value` being
embedded in a `Const` node in the `tf.Graph`. `tf.constant` is useful
for asserting that the value can be embedded that way.

If the argument `dtype` is not specified, then the type is inferred from
the type of `value`.

>>> # Constant 1-D Tensor from a python list.
>>> tf.constant([1, 2, 3, 4, 5, 6])
<tf.Tensor: shape=(6,), dtype=int32,
    numpy=array([1, 2, 3, 4, 5, 6], dtype=int32)>
>>> # Or a numpy array
>>> a = np.array([[1, 2, 3], [4, 5, 6]])
>>> tf.constant(a)
<tf.Tensor: shape=(2, 3), dtype=int64, numpy=
  array([[1, 2, 3],
         [4, 5, 6]])>

If `dtype` is specified, the resulting tensor values are cast to the requested
`dtype`.

>>> tf.constant([1, 2, 3, 4, 5, 6], dtype=tf.float64)
<tf.Tensor: shape=(6,), dtype=float64,
    numpy=array([1., 2., 3., 4., 5., 6.])>

If `shape` is set, the `value` is reshaped to match. Scalars are expanded to
fill the `shape`:

>>> tf.constant(0, shape=(2, 3))
  <tf.Tensor: shape=(2, 3), dtype=int32, numpy=
  array([[0, 0, 0],
         [0, 0, 0]], dtype=int32)>
>>> tf.constant([1, 2, 3, 4, 5, 6], shape=[2, 3])
<tf.Tensor: shape=(2, 3), dtype=int32, numpy=
  array([[1, 2, 3],
         [4, 5, 6]], dtype=int32)>

`tf.constant` has no effect if an eager Tensor is passed as the `value`, it
even transmits gradients:

>>> v = tf.Variable([0.0])
>>> with tf.GradientTape() as g:
...     loss = tf.constant(v + v)
>>> g.gradient(loss, v).numpy()
array([2.], dtype=float32)

But, since `tf.constant` embeds the value in the `tf.Graph` this fails for
symbolic tensors:

>>> with tf.compat.v1.Graph().as_default():
...   i = tf.compat.v1.placeholder(shape=[None, None], dtype=tf.float32)
...   t = tf.constant(i)
Traceback (most recent call last):
...
TypeError: ...

`tf.constant` will create tensors on the current device. Inputs which are
already tensors maintain their placements unchanged.

Related Ops:

* `tf.convert_to_tensor` is similar but:
  * It has no `shape` argument.
  * Symbolic tensors are allowed to pass through.

  >>> with tf.compat.v1.Graph().as_default():
  ...   i = tf.compat.v1.placeholder(shape=[None, None], dtype=tf.float32)
  ...   t = tf.convert_to_tensor(i)

* `tf.fill`: differs in a few ways:
  *   `tf.constant` supports arbitrary constants, not just uniform scalar
      Tensors like `tf.fill`.
  *   `tf.fill` creates an Op in the graph that is expanded at runtime, so it
      can efficiently represent large tensors.
  *   Since `tf.fill` does not embed the value, it can produce dynamically
      sized outputs.

Args:
  value: A constant value (or list) of output type `dtype`.
  dtype: The type of the elements of the resulting tensor.
  shape: Optional dimensions of resulting tensor.
  name: Optional name for the tensor.

Returns:
  A Constant Tensor.

Raises:
  TypeError: if shape is incorrectly specified or unsupported.
  ValueError: if called on a symbolic tensor.


def _remove_one_redundant_stack_unstack(in_graph_def):
    name_to_input_name, name_to_node, name_to_seq_num = _extract_graph_summary(in_graph_def)
    del name_to_seq_num
    do_generic_pack_unpack = True
    out = _graph_pb2.GraphDef()
    out.library.CopyFrom(in_graph_def.library)
    out.versions.CopyFrom(in_graph_def.versions)
    for n in in_graph_def.node:
        node_name = _tensor_name_base(n.name)
        if not node_name.startswith('OpHintStack') and (not n.op.startswith('Pack')):
            continue
        next_to_visit = [node_name]
        visited = set()
        unpack_nodes = set()
        pack_node = node_name
        matches_pattern = True
        is_hint_created_stack = False
        while next_to_visit:
            current_node_name = next_to_visit[0]
            visited.add(current_node_name)
            del next_to_visit[0]
            node = name_to_node[current_node_name]
            is_op_hint_stack = node.name.startswith('OpHintStack')
            is_op_hint_unstack = node.name.startswith('OpHintUnstack')
            if node.op == 'Identity' or is_op_hint_stack or (do_generic_pack_unpack and node.op == 'Pack'):
                is_hint_created_stack |= is_op_hint_stack
                next_to_visit += [input_node for input_node in name_to_input_name[current_node_name] if input_node not in visited]
            elif is_op_hint_unstack or (do_generic_pack_unpack and node.op == 'Unpack'):
                unpack_nodes.add(node.name)
                is_hint_created_stack &= is_op_hint_unstack
            else:
                matches_pattern = False
                break
            visited.add(node.name)
        if matches_pattern and len(unpack_nodes) == 1:
            pack_node = node_name
            no_external_dependency = True
            for other_n in in_graph_def.node:
                if other_n.name in visited:
                    continue
                for input_tensor in name_to_input_name[other_n.name]:
                    input_op = _tensor_name_base(input_tensor)
                    if input_op in visited and input_op != pack_node:
                        no_external_dependency = False
            if is_hint_created_stack or no_external_dependency:
                end = unpack_nodes.pop()
                end_input = name_to_node[end].input[0]
                for other_n in in_graph_def.node:
                    node_name = _tensor_name_base(other_n.name)
                    if node_name not in visited:
                        new_node = _copy.deepcopy(other_n)
                        new_node.input[:] = [end_input if stripped == pack_node else non_stripped for stripped, non_stripped in zip(name_to_input_name[node_name], new_node.input[:])]
                        out.node.extend([new_node])
                return (out, True)
    return (in_graph_def, False)

Removes a stack->unstack pattern from in_graph_def in a returned graph.

Args:
  in_graph_def: Graph def to use as input.

Returns:
  Simplified tuple (graph_def, changed_something) where changed_something
  is true if anything was done.


def __call__(self, images: ImageInput=None, text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]]=None, audio=None, videos=None, **kwargs: Unpack[ColQwen2ProcessorKwargs]) -> BatchFeature:
    output_kwargs = self._merge_kwargs(ColQwen2ProcessorKwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, **kwargs)
    suffix = output_kwargs['text_kwargs'].pop('suffix', None)
    return_token_type_ids = True if suffix is not None else False
    if text is None and images is None:
        raise ValueError('Either text or images must be provided')
    if text is not None and images is not None:
        raise ValueError('Only one of text or images can be processed at a time')
    if images is not None:
        if is_valid_image(images):
            images = [images]
        elif isinstance(images, list) and is_valid_image(images[0]):
            pass
        elif not (isinstance(images, list) and isinstance(images[0], list) and is_valid_image(images[0][0])):
            raise ValueError('images must be an image, list of images or list of list of images')
        texts_doc = [self.visual_prompt_prefix] * len(images)
        image_inputs = self.image_processor(images=images, **output_kwargs['images_kwargs'])
        image_grid_thw = image_inputs['image_grid_thw']
        if image_grid_thw is not None:
            merge_length = self.image_processor.merge_size ** 2
            index = 0
            for i in range(len(texts_doc)):
                while self.image_token in texts_doc[i]:
                    texts_doc[i] = texts_doc[i].replace(self.image_token, '<|placeholder|>' * (image_grid_thw[index].prod() // merge_length), 1)
                    index += 1
                texts_doc[i] = texts_doc[i].replace('<|placeholder|>', self.image_token)
        text_inputs = self.tokenizer(texts_doc, return_token_type_ids=False, **output_kwargs['text_kwargs'])
        return_data = BatchFeature(data={**text_inputs, **image_inputs})
        offsets = return_data['image_grid_thw'][:, 1] * return_data['image_grid_thw'][:, 2]
        pixel_values = list(torch.split(return_data['pixel_values'], offsets.tolist()))
        return_data['pixel_values'] = torch.nn.utils.rnn.pad_sequence(pixel_values, batch_first=True)
        if return_token_type_ids:
            labels = return_data['input_ids'].masked_fill(return_data['token_type_ids'] == 0, -100)
            return_data.update({'labels': labels})
        return return_data
    elif text is not None:
        if isinstance(text, str):
            text = [text]
        elif not (isinstance(text, list) and isinstance(text[0], str)):
            raise ValueError('Text must be a string or a list of strings')
        if suffix is None:
            suffix = self.query_augmentation_token * 10
        texts_query: List[str] = []
        for query in text:
            augmented_query = self.query_prefix + query + suffix
            texts_query.append(augmented_query)
        batch_query = self.tokenizer(texts_query, return_token_type_ids=False, **output_kwargs['text_kwargs'])
        return batch_query

Main method to prepare for the model either (1) one or several texts, either (2) one or several image(s). This method is a custom
wrapper around the Qwen2VLProcessor's [`~Qwen2VLProcessor.__call__`] method adapted for the ColQwen2 model. It cannot process
both text and images at the same time.

When preparing the the text(s), this method forwards the `text` and `kwargs` arguments to Qwen2TokenizerFast's
[`~Qwen2TokenizerFast.__call__`].
When preparing the the image(s), this method forwards the `images` and `kwargs` arguments to Qwen2VLImageProcessor's
[`~Qwen2VLImageProcessor.__call__`].
Please refer to the doctsring of the above two methods for more information.

Args:
    images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`):
        The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
        tensor. In case of a NumPy array/PyTorch tensor, each image should be of shape (C, H, W), where C is a
        number of channels, H and W are image height and width.
    text (`str`, `List[str]`, `List[List[str]]`):
        The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
        (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
        `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
    return_tensors (`str` or [`~utils.TensorType`], *optional*):
        If set, will return tensors of a particular framework. Acceptable values are:

        - `'tf'`: Return TensorFlow `tf.constant` objects.
        - `'pt'`: Return PyTorch `torch.Tensor` objects.
        - `'np'`: Return NumPy `np.ndarray` objects.
        - `'jax'`: Return JAX `jnp.ndarray` objects.

Returns:
    [`BatchFeature`]: A [`BatchFeature`] with the following fields:

    - **input_ids** -- List of token ids to be fed to a model.
    - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
      `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
      `None`).
    - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.


def _merge_with(self, other: 'DynamicRaggedShape') -> 'DynamicRaggedShape':
    max_num_row_partitions = max(self.num_row_partitions, other.num_row_partitions)
    a = self._with_num_row_partitions(max_num_row_partitions)
    b = other._with_num_row_partitions(max_num_row_partitions)
    new_row_partitions = [rp_a._merge_precomputed_encodings(rp_b) for rp_a, rp_b in zip(a._row_partitions, b._row_partitions)]
    new_dtype = b.dtype if a.dtype == dtypes.int32 else dtypes.int64
    new_static_inner_shape = a._static_inner_shape.merge_with(b._static_inner_shape)
    new_inner_shape = a._inner_shape
    return DynamicRaggedShape(new_row_partitions, new_inner_shape, new_dtype, True, new_static_inner_shape)

Merge two shapes that are equal modulo num_row_partitions.

The resulting num_row_partitions is the maximum of the two
num_row_partitions.

Args:
  other: a DynamicRaggedShape representing the same shape with a possibly
    different number of row partitions.

Returns:
  A DynamicRaggedShape with the same shape and the maximum of the
  num_row_partitions of the two shapes.


def get_diff(repo: Repo, base_commit: str, commits: List[str]) -> List[str]:
    print('\n### DIFF ###\n')
    code_diff = []
    for commit in commits:
        for diff_obj in commit.diff(base_commit):
            if diff_obj.change_type == 'A' and diff_obj.b_path.endswith('.py'):
                code_diff.append(diff_obj.b_path)
            elif diff_obj.change_type == 'D' and diff_obj.a_path.endswith('.py'):
                code_diff.append(diff_obj.a_path)
            elif diff_obj.change_type in ['M', 'R'] and diff_obj.b_path.endswith('.py'):
                if diff_obj.a_path != diff_obj.b_path:
                    code_diff.extend([diff_obj.a_path, diff_obj.b_path])
                elif diff_is_docstring_only(repo, commit, diff_obj.b_path):
                    print(f'Ignoring diff in {diff_obj.b_path} as it only concerns docstrings or comments.')
                else:
                    code_diff.append(diff_obj.a_path)
    return code_diff

Get the diff between a base commit and one or several commits.

Args:
    repo (`git.Repo`):
        A git repository (for instance the Transformers repo).
    base_commit (`str`):
        The commit reference of where to compare for the diff. This is the current commit, not the branching point!
    commits (`List[str]`):
        The list of commits with which to compare the repo at `base_commit` (so the branching point).

Returns:
    `List[str]`: The list of Python files with a diff (files added, renamed or deleted are always returned, files
    modified are returned if the diff in the file is not only in docstrings or comments, see
    `diff_is_docstring_only`).


def set_index(self: EventSetOrNode, indexes: Union[str, List[str]]) -> EventSetOrNode:
    from temporian.core.operators.add_index import set_index
    return set_index(self, indexes=indexes)

Replaces the index in an [`EventSet`][temporian.EventSet].

Usage example:
    ```python
    >>> a = tp.event_set(
    ...     timestamps=[1, 2, 1, 0, 1, 1],
    ...     features={
    ...         "f1": [1, 1, 1, 2, 2, 2],
    ...         "f2": [1, 1, 2, 1, 1, 2],
    ...         "f3": [1, 1, 1, 1, 1, 1]
    ...     },
    ...     indexes=["f1"],
    ... )

    >>> # "f1" is the current index
    >>> a
    indexes: [('f1', int64)]
    features: [('f2', int64), ('f3', int64)]
    events:
        f1=1 (3 events):
            timestamps: [1. 1. 2.]
            'f2': [1 2 1]
            'f3': [1 1 1]
        f1=2 (3 events):
            timestamps: [0. 1. 1.]
            'f2': [1 1 2]
            'f3': [1 1 1]
    ...

    >>> # Set "f2" as the only index, remove "f1"
    >>> b = a.set_index("f2")
    >>> b
    indexes: [('f2', int64)]
    features: [('f3', int64), ('f1', int64)]
    events:
        f2=1 (4 events):
            timestamps: [0. 1. 1. 2.]
            'f3': [1 1 1 1]
            'f1': [2 1 2 1]
        f2=2 (2 events):
            timestamps: [1. 1.]
            'f3': [1 1]
            'f1': [1 2]
    ...

    >>> # Set both "f1" and "f2" as indices
    >>> b = a.set_index(["f1", "f2"])
    >>> b
    indexes: [('f1', int64), ('f2', int64)]
    features: [('f3', int64)]
    events:
        f1=1 f2=1 (2 events):
            timestamps: [1. 2.]
            'f3': [1 1]
        f1=1 f2=2 (1 events):
            timestamps: [1.]
            'f3': [1]
        f1=2 f2=1 (2 events):
            timestamps: [0. 1.]
            'f3': [1 1]
        f1=2 f2=2 (1 events):
            timestamps: [1.]
            'f3': [1]
    ...

    ```

Args:
    indexes: List of index / feature names (strings) used as
        the new indexes. These names should be either indexes or
        features in the input.

Returns:
    EventSet with the updated indexes.

Raises:
    KeyError: If any of the specified `indexes` are not found in the
        input.


def _init_init_op(self, init_op=USE_DEFAULT, init_feed_dict=None):
    if init_op is Supervisor.USE_DEFAULT:
        init_op = self._get_first_op_from_collection(ops.GraphKeys.INIT_OP)
        if init_op is None:
            init_op = variables.global_variables_initializer()
            ops.add_to_collection(ops.GraphKeys.INIT_OP, init_op)
    self._init_op = init_op
    self._init_feed_dict = init_feed_dict

Initializes init_op.

Args:
  init_op: `Operation` to initialize the variables. If set to USE_DEFAULT,
    create an op that initializes all variables and tables.
  init_feed_dict: A dictionary that maps `Tensor` objects to feed values.
    This feed dictionary will be used when `init_op` is evaluated.


class BaseModelOutput(ModelOutput):
    last_hidden_state: Optional[torch.FloatTensor] = None
    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None

Base class for model's outputs, with potential hidden states and attentions.

Args:
    last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
        Sequence of hidden-states at the output of the last layer of the model.
    hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
        one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.

        Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
    attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
        Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
        sequence_length)`.

        Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
        heads.


class PromptDepthAnythingConfig(PretrainedConfig):
    model_type = 'prompt_depth_anything'

    def __init__(self, backbone_config=None, backbone=None, use_pretrained_backbone=False, use_timm_backbone=False, backbone_kwargs=None, patch_size=14, initializer_range=0.02, reassemble_hidden_size=384, reassemble_factors=[4, 2, 1, 0.5], neck_hidden_sizes=[48, 96, 192, 384], fusion_hidden_size=64, head_in_index=-1, head_hidden_size=32, depth_estimation_type='relative', max_depth=None, **kwargs):
        super().__init__(**kwargs)
        if backbone_config is None and backbone is None:
            logger.info('`backbone_config` is `None`. Initializing the config with the default `Dinov2` backbone.')
            backbone_config = CONFIG_MAPPING['dinov2'](image_size=518, hidden_size=384, num_attention_heads=6, out_indices=[9, 10, 11, 12], apply_layernorm=True, reshape_hidden_states=False)
        elif isinstance(backbone_config, dict):
            backbone_model_type = backbone_config.get('model_type')
            config_class = CONFIG_MAPPING[backbone_model_type]
            backbone_config = config_class.from_dict(backbone_config)
        verify_backbone_config_arguments(use_timm_backbone=use_timm_backbone, use_pretrained_backbone=use_pretrained_backbone, backbone=backbone, backbone_config=backbone_config, backbone_kwargs=backbone_kwargs)
        self.backbone_config = backbone_config
        self.backbone = backbone
        self.use_pretrained_backbone = use_pretrained_backbone
        self.use_timm_backbone = use_timm_backbone
        self.backbone_kwargs = backbone_kwargs
        self.reassemble_hidden_size = reassemble_hidden_size
        self.patch_size = patch_size
        self.initializer_range = initializer_range
        self.reassemble_factors = reassemble_factors
        self.neck_hidden_sizes = neck_hidden_sizes
        self.fusion_hidden_size = fusion_hidden_size
        self.head_in_index = head_in_index
        self.head_hidden_size = head_hidden_size
        if depth_estimation_type not in ['relative', 'metric']:
            raise ValueError("depth_estimation_type must be one of ['relative', 'metric']")
        self.depth_estimation_type = depth_estimation_type
        self.max_depth = max_depth if max_depth else 1

    def to_dict(self):
        """
        Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`]. Returns:
            `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
        """
        output = copy.deepcopy(self.__dict__)
        if output['backbone_config'] is not None:
            output['backbone_config'] = self.backbone_config.to_dict()
        output['model_type'] = self.__class__.model_type
        return output

This is the configuration class to store the configuration of a [`PromptDepthAnythingModel`]. It is used to instantiate a PromptDepthAnything
model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
defaults will yield a similar configuration to that of the PromptDepthAnything
[LiheYoung/depth-anything-small-hf](https://huggingface.co/LiheYoung/depth-anything-small-hf) architecture.

Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
documentation from [`PretrainedConfig`] for more information.

Args:
    backbone_config (`Union[Dict[str, Any], PretrainedConfig]`, *optional*):
        The configuration of the backbone model. Only used in case `is_hybrid` is `True` or in case you want to
        leverage the [`AutoBackbone`] API.
    backbone (`str`, *optional*):
        Name of backbone to use when `backbone_config` is `None`. If `use_pretrained_backbone` is `True`, this
        will load the corresponding pretrained weights from the timm or transformers library. If `use_pretrained_backbone`
        is `False`, this loads the backbone's config and uses that to initialize the backbone with random weights.
    use_pretrained_backbone (`bool`, *optional*, defaults to `False`):
        Whether to use pretrained weights for the backbone.
    use_timm_backbone (`bool`, *optional*, defaults to `False`):
        Whether or not to use the `timm` library for the backbone. If set to `False`, will use the [`AutoBackbone`]
        API.
    backbone_kwargs (`dict`, *optional*):
        Keyword arguments to be passed to AutoBackbone when loading from a checkpoint
        e.g. `{'out_indices': (0, 1, 2, 3)}`. Cannot be specified if `backbone_config` is set.
    patch_size (`int`, *optional*, defaults to 14):
        The size of the patches to extract from the backbone features.
    initializer_range (`float`, *optional*, defaults to 0.02):
        The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
    reassemble_hidden_size (`int`, *optional*, defaults to 384):
        The number of input channels of the reassemble layers.
    reassemble_factors (`List[int]`, *optional*, defaults to `[4, 2, 1, 0.5]`):
        The up/downsampling factors of the reassemble layers.
    neck_hidden_sizes (`List[str]`, *optional*, defaults to `[48, 96, 192, 384]`):
        The hidden sizes to project to for the feature maps of the backbone.
    fusion_hidden_size (`int`, *optional*, defaults to 64):
        The number of channels before fusion.
    head_in_index (`int`, *optional*, defaults to -1):
        The index of the features to use in the depth estimation head.
    head_hidden_size (`int`, *optional*, defaults to 32):
        The number of output channels in the second convolution of the depth estimation head.
    depth_estimation_type (`str`, *optional*, defaults to `"relative"`):
        The type of depth estimation to use. Can be one of `["relative", "metric"]`.
    max_depth (`float`, *optional*):
        The maximum depth to use for the "metric" depth estimation head. 20 should be used for indoor models
        and 80 for outdoor models. For "relative" depth estimation, this value is ignored.

Example:

```python
>>> from transformers import PromptDepthAnythingConfig, PromptDepthAnythingForDepthEstimation

>>> # Initializing a PromptDepthAnything small style configuration
>>> configuration = PromptDepthAnythingConfig()

>>> # Initializing a model from the PromptDepthAnything small style configuration
>>> model = PromptDepthAnythingForDepthEstimation(configuration)

>>> # Accessing the model configuration
>>> configuration = model.config
```


class VanConfig(PretrainedConfig):
    model_type = 'van'

    def __init__(self, image_size=224, num_channels=3, patch_sizes=[7, 3, 3, 3], strides=[4, 2, 2, 2], hidden_sizes=[64, 128, 320, 512], depths=[3, 3, 12, 3], mlp_ratios=[8, 8, 4, 4], hidden_act='gelu', initializer_range=0.02, layer_norm_eps=1e-06, layer_scale_init_value=0.01, drop_path_rate=0.0, dropout_rate=0.0, **kwargs):
        super().__init__(**kwargs)
        self.image_size = image_size
        self.num_channels = num_channels
        self.patch_sizes = patch_sizes
        self.strides = strides
        self.hidden_sizes = hidden_sizes
        self.depths = depths
        self.mlp_ratios = mlp_ratios
        self.hidden_act = hidden_act
        self.initializer_range = initializer_range
        self.layer_norm_eps = layer_norm_eps
        self.layer_scale_init_value = layer_scale_init_value
        self.drop_path_rate = drop_path_rate
        self.dropout_rate = dropout_rate

This is the configuration class to store the configuration of a [`VanModel`]. It is used to instantiate a VAN model
according to the specified arguments, defining the model architecture. Instantiating a configuration with the
defaults will yield a similar configuration to that of the VAN
[Visual-Attention-Network/van-base](https://huggingface.co/Visual-Attention-Network/van-base) architecture.

Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
documentation from [`PretrainedConfig`] for more information.

Args:
    image_size (`int`, *optional*, defaults to 224):
        The size (resolution) of each image.
    num_channels (`int`, *optional*, defaults to 3):
        The number of input channels.
    patch_sizes (`List[int]`, *optional*, defaults to `[7, 3, 3, 3]`):
        Patch size to use in each stage's embedding layer.
    strides (`List[int]`, *optional*, defaults to `[4, 2, 2, 2]`):
        Stride size to use in each stage's embedding layer to downsample the input.
    hidden_sizes (`List[int]`, *optional*, defaults to `[64, 128, 320, 512]`):
        Dimensionality (hidden size) at each stage.
    depths (`List[int]`, *optional*, defaults to `[3, 3, 12, 3]`):
        Depth (number of layers) for each stage.
    mlp_ratios (`List[int]`, *optional*, defaults to `[8, 8, 4, 4]`):
        The expansion ratio for mlp layer at each stage.
    hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
        The non-linear activation function (function or string) in each layer. If string, `"gelu"`, `"relu"`,
        `"selu"` and `"gelu_new"` are supported.
    initializer_range (`float`, *optional*, defaults to 0.02):
        The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
    layer_norm_eps (`float`, *optional*, defaults to 1e-06):
        The epsilon used by the layer normalization layers.
    layer_scale_init_value (`float`, *optional*, defaults to 0.01):
        The initial value for layer scaling.
    drop_path_rate (`float`, *optional*, defaults to 0.0):
        The dropout probability for stochastic depth.
    dropout_rate (`float`, *optional*, defaults to 0.0):
        The dropout probability for dropout.

Example:
```python
>>> from transformers import VanModel, VanConfig

>>> # Initializing a VAN van-base style configuration
>>> configuration = VanConfig()
>>> # Initializing a model from the van-base style configuration
>>> model = VanModel(configuration)
>>> # Accessing the model configuration
>>> configuration = model.config
```


class SKLearnClassifier(ClassifierMixin, SKLBase):

    def _process_target(self, y, reset=False):
        """Classifiers do OHE."""
        target_type = type_of_target(y, raise_unknown=True)
        if target_type not in ['binary', 'multiclass']:
            raise ValueError(f'Only binary and multiclass target types are supported. Target type: {target_type}')
        if reset:
            self._target_encoder = sklearn.pipeline.make_pipeline(TargetReshaper(), sklearn.preprocessing.OneHotEncoder(sparse_output=False)).fit(y)
            self.classes_ = np.unique(y)
            if len(self.classes_) == 1:
                raise ValueError("Classifier can't train when only one class is present.")
        return self._target_encoder.transform(y)

    def _more_tags(self):
        return {'poor_score': True}

    def __sklearn_tags__(self):
        tags = super().__sklearn_tags__()
        tags.classifier_tags.poor_score = True
        return tags

scikit-learn compatible classifier wrapper for Keras models.

Note that there are sources of randomness in model initialization and
training. Refer to [Reproducibility in Keras Models](
https://keras.io/examples/keras_recipes/reproducibility_recipes/) on how to
control randomness.

Args:
    model: `Model`.
        An instance of `Model`, or a callable returning such an object.
        Note that if input is a `Model`, it will be cloned using
        `keras.models.clone_model` before being fitted, unless
        `warm_start=True`.
        The `Model` instance needs to be passed as already compiled.
        If callable, it must accept at least `X` and `y` as keyword
        arguments. Other arguments must be accepted if passed as
        `model_kwargs` by the user.
    warm_start: bool, defaults to `False`.
        Whether to reuse the model weights from the previous fit. If `True`,
        the given model won't be cloned and the weights from the previous
        fit will be reused.
    model_kwargs: dict, defaults to `None`.
        Keyword arguments passed to `model`, if `model` is callable.
    fit_kwargs: dict, defaults to `None`.
        Keyword arguments passed to `model.fit`. These can also be passed
        directly to the `fit` method of the scikit-learn wrapper. The
        values passed directly to the `fit` method take precedence over
        these.

Attributes:
    model_ : `Model`
        The fitted model.
    history_ : dict
        The history of the fit, returned by `model.fit`.
    classes_ : array-like, shape=(n_classes,)
        The classes labels.

Example:
Here we use a function which creates a basic MLP model dynamically
choosing the input and output shapes. We will use this to create our
scikit-learn model.

``` python
from keras.src.layers import Dense, Input, Model

def dynamic_model(X, y, loss, layers=[10]):
    # Creates a basic MLP model dynamically choosing the input and
    # output shapes.
    n_features_in = X.shape[1]
    inp = Input(shape=(n_features_in,))

    hidden = inp
    for layer_size in layers:
        hidden = Dense(layer_size, activation="relu")(hidden)

    n_outputs = y.shape[1] if len(y.shape) > 1 else 1
    out = [Dense(n_outputs, activation="softmax")(hidden)]
    model = Model(inp, out)
    model.compile(loss=loss, optimizer="rmsprop")

    return model
```

You can then use this function to create a scikit-learn compatible model
and fit it on some data.

``` python
from sklearn.datasets import make_classification
from keras.wrappers import SKLearnClassifier

X, y = make_classification(n_samples=1000, n_features=10, n_classes=3)
est = SKLearnClassifier(
    model=dynamic_model,
    model_kwargs={
        "loss": "categorical_crossentropy",
        "layers": [20, 20, 20],
    },
)

est.fit(X, y, epochs=5)
```


class TrueNegatives(_ConfusionMatrixConditionCount):

    def __init__(self, thresholds=None, name=None, dtype=None):
        super(TrueNegatives, self).__init__(confusion_matrix_cond=metrics_utils.ConfusionMatrix.TRUE_NEGATIVES, thresholds=thresholds, name=name, dtype=dtype)

Calculates the number of true negatives.

If `sample_weight` is given, calculates the sum of the weights of
true negatives. This metric creates one local variable, `accumulator`
that is used to keep track of the number of true negatives.

If `sample_weight` is `None`, weights default to 1.
Use `sample_weight` of 0 to mask values.

Args:
  thresholds: (Optional) Defaults to 0.5. A float value or a python
    list/tuple of float threshold values in [0, 1]. A threshold is compared
    with prediction values to determine the truth value of predictions
    (i.e., above the threshold is `true`, below is `false`). One metric
    value is generated for each threshold value.
  name: (Optional) string name of the metric instance.
  dtype: (Optional) data type of the metric result.

Standalone usage:

>>> m = tf.keras.metrics.TrueNegatives()
>>> m.update_state([0, 1, 0, 0], [1, 1, 0, 0])
>>> m.result().numpy()
2.0

>>> m.reset_state()
>>> m.update_state([0, 1, 0, 0], [1, 1, 0, 0], sample_weight=[0, 0, 1, 0])
>>> m.result().numpy()
1.0

Usage with `compile()` API:

```python
model.compile(optimizer='sgd',
              loss='mse',
              metrics=[tf.keras.metrics.TrueNegatives()])
```


def enqueue_tpu_embedding_sparse_tensor_batch(sample_indices, embedding_indices, aggregation_weights, table_ids, device_ordinal, max_sequence_lengths=None, num_features=None, combiners=None, mode_override=None, name=None):
    if mode_override is None:
        mode_override = 'unspecified'
    return gen_tpu_ops.enqueue_tpu_embedding_sparse_tensor_batch(sample_indices=sample_indices, embedding_indices=embedding_indices, aggregation_weights=aggregation_weights, table_ids=table_ids, device_ordinal=device_ordinal, max_sequence_lengths=max_sequence_lengths, combiners=combiners, mode_override=mode_override, num_features=num_features, name=name)

A placeholder op for enqueueing embedding IDs to the TPU.

Args:
  sample_indices: A list of rank 2 Tensors specifying the training example to
    which the corresponding embedding_indices and aggregation_weights values
    belong. It corresponds to sp_ids.indices in embedding_lookup_sparse(). If
    the size of its first dimension is 0, we assume each embedding_indices
    belongs to a different sample. Both int32 and int64 are allowed and will
    be converted to int32 internally.
  embedding_indices: A list of rank 1 Tensors, indices into the embedding
    tables. It corresponds to sp_ids.values in embedding_lookup_sparse(). Both
    int32 and int64 are allowed and will be converted to int32 internally.
  aggregation_weights: A list of rank 1 Tensors containing per training
    example aggregation weights. It corresponds to sp_weights.values in
    embedding_lookup_sparse(). If the size of its first dimension is 0, we
    assume all weights are 1. Both float32 and float64 are allowed and will be
    converted to float32 internally.
  table_ids: A list of integers specifying the identifier of the embedding
    table (offset of TableDescriptor in the TPUEmbeddingConfiguration) to
    lookup the corresponding input. The ith input is looked up using
    table_ids[i]. The size of the table_ids list must be equal to that of
    sample_indices, embedding_indices and aggregation_weights.
  device_ordinal: The TPU device to use. Should be >= 0 and less than the
    number of TPU cores in the task on which the node is placed.
  max_sequence_lengths: A list of integers, the size of which is equal to
    sample_indices. If equal to 0, the corresponding feature is considered to
    be a non-sequence feature, If greater than 0, the corresponding feature is
    a sequence feature with the given maximal length. If None, then we assume
    a list of all zeroes.
  num_features: A list of integers, the size of which is equal to
    sample_indices. If non-empty, entries in this list must be at least 1. For
    each batch element, we will take num_features rows of the input tensor for
    embedding lookup. E.g., when sample_indices is empty, the embedding
    indices must be of shape (batch_size*num_features).
  combiners: A list of string scalars, one for each embedding table that
    specify how to normalize the embedding activations after weighted
    summation. Supported combiners are 'mean', 'sum', or 'sqrtn'. It is
    invalid to have the sum of the weights be 0 for 'mean' or the sum of the
    squared weights be 0 for 'sqrtn'. If combiners isn't passed, the default
    is to use 'sum' for all tables (optional).
  mode_override: A string input that overrides the mode specified in the
    TPUEmbeddingConfiguration. Supported values are {'unspecified',
    'inference', 'train', 'backward_pass_only'}. When set to 'unspecified',
    the mode set in TPUEmbeddingConfiguration is used, otherwise mode_override
    is used (optional).
  name: A name for the operation (optional).

Returns:
  An EnqueueTPUEmbeddingSparseTensorBatch operation.


def get_text_features(self, input_ids: Optional[torch.LongTensor]=None, text_encoder_inputs_embeds: Optional[torch.FloatTensor]=None, attention_mask: Optional[torch.LongTensor]=None) -> torch.FloatTensor:
    outputs = self.text_encoder_model(input_ids=input_ids, inputs_embeds=text_encoder_inputs_embeds, attention_mask=attention_mask)
    return outputs[0]

This method can be used to extract text_embeds from a text. The text embeddings obtained by applying the
projection layer to the pooled output of the CLVP text encoder model.

Args:
    input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
        Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you
        provide it.

        [What are input IDs?](../glossary#input-ids)
    text_encoder_inputs_embeds (`torch.FloatTensor`, *optional*):
        inputs_embeds for the text encoder model passed in place of `input_ids`.
    attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
        Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:

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

        [What are attention masks?](../glossary#attention-mask)

Returns:
    `torch.FloatTensor` of shape `(batch_size, output_dim)`:
        The text embeddings obtained by applying the projection layer to the pooled output of the CLVP Text
        Model.

Examples:

```python
>>> from transformers import ClvpProcessor, ClvpModelForConditionalGeneration

>>> # Define the Text
>>> text = "This is an example text."

>>> # Define processor and model
>>> processor = ClvpProcessor.from_pretrained("susnato/clvp_dev")
>>> model = ClvpModelForConditionalGeneration.from_pretrained("susnato/clvp_dev")

>>> # Generate processor output and text embeds
>>> processor_output = processor(text=text, return_tensors="pt")
>>> text_embeds = model.get_text_features(input_ids=processor_output["input_ids"])
```


def get_resize_output_image_size(image, resolution_max_side: int, input_data_format: Optional[Union[str, ChannelDimension]]=None) -> Tuple[int, int]:
    height, width = get_image_size(image, channel_dim=input_data_format)
    height, width = _resize_output_size_rescale_to_max_len(height, width, max_len=resolution_max_side)
    height, width = _resize_output_size_scale_below_upper_bound(height, width, max_len=MAX_IMAGE_SIZE)
    return (height, width)

Get the output size of the image after resizing given a dictionary specifying the max and min sizes.
Args:
    image (`np.ndarray`):
        Image to resize.
    resolution_max_side (`int`):
        The longest edge of the image will be resized to this value. The shortest edge will be resized to keep the
        input aspect ratio.
    input_data_format (`ChannelDimension` or `str`):
        The channel dimension format of the input image.
Returns:
    The output size of the image after resizing.


def _try_put(self, item):
    try:
        self._event_queue.put(item)
    except QueueClosedError:
        self._internal_close()
        if self._worker.failure_exc_info:
            _, exception, _ = self._worker.failure_exc_info
            raise exception from None

Attempts to enqueue an item to the event queue.

If the queue is closed, this will close the EventFileWriter and reraise the
exception that caused the queue closure, if one exists.

Args:
  item: the item to enqueue


def update_state(self, y_true, y_pred, sample_weight=None):
    y_true = ops.convert_to_tensor(y_true, dtype=self.dtype)
    y_pred = ops.convert_to_tensor(y_pred, dtype='float32')
    y_pred = ops.cast(y_pred >= self.threshold, self.dtype)
    return super().update_state(y_true, y_pred, sample_weight)

Accumulates the confusion matrix statistics.

Before the confusion matrix is updated, the predicted values are
thresholded to be:
    0 for values that are smaller than the `threshold`
    1 for values that are larger or equal to the `threshold`

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

Returns:
    Update op.


def _create_or_get_tensor_values_cache(self, cache_name, graph, shape=None, dtype=dtypes.float32):
    if graph is None:
        raise ValueError('Invalid graph.')
    graph_cache_var = self._cache_variable_for_graph(graph)
    if cache_name not in graph_cache_var:
        if shape is None:
            raise ValueError('shape must be provided at cache creation.')
        if dtype.is_integer:
            init_val = int(_COMPACT_TRACE_ENTRY_INIT_VALUE)
        else:
            init_val = _COMPACT_TRACE_ENTRY_INIT_VALUE
        with graph.as_default() as g, g.name_scope(None):
            graph_cache_var[cache_name] = variable_scope.get_variable(_TT_SNAPSHOT + '_' + self._escape_namescopes(cache_name), shape=shape, dtype=dtype, initializer=init_ops.constant_initializer(init_val), trainable=False, use_resource=True, collections=[_TENSOR_TRACER_STORAGE, ops.GraphKeys.LOCAL_VARIABLES])
    return graph_cache_var[cache_name]

Creates a variable as the cache to store intermediate tensor values.

Args:
  cache_name: Name to be given to the cache (an instance of tf.variable).
  graph: Tensorflow graph.
  shape: A list of dimensions.
  dtype: Data type of created cache.
Returns:
  A ref to newly created or existing cache with the given dimensions.
Raises:
  ValueError:
    (1) If graph is None, or
    (2) shape is None when a new cache needs to be created.


def __init__(self, base: ModelHandler[ExampleT, PredictionT, ModelT]):
    self._base = base
    self._env_vars = getattr(base, '_env_vars', {})

A ModelHandler that skips batching in RunInference.

Args:
  base: An implementation of the underlying model handler.


def get_min_max_value(statistics: calib_stats_pb2.CalibrationStatistics, calib_opts: stablehlo_quant_config_pb2.CalibrationOptions) -> tuple[float, float]:
    calib_method = calib_opts.calibration_method
    if calib_method not in _REGISTRY:
        raise ValueError(f'Unsupported calibration method: {calib_method}')
    calibration_algorithm = _REGISTRY[calib_method](statistics, calib_opts)
    return calibration_algorithm.get_min_max_value()

Calculates min and max from statistics using calibration options.

Args:
  statistics: Collected calibration statistics.
  calib_opts: Calibration options used for calculating min and max.

Returns:
  (min_value, max_value): Min and max calculated using calib_opts.

Raises:
  ValueError: Unsupported calibration method is given.


def _from_tensor_shape(cls, shape: Any, num_row_partitions: int, dtype: dtypes.DType) -> 'DynamicRaggedShape.Spec':
    if dtype != dtypes.int32 and dtype != dtypes.int64:
        raise ValueError('dtype must be tf.int32 or tf.int64')
    shape = tensor_shape.as_shape(shape)
    if shape.rank is None:
        row_partitions = [RowPartitionSpec(dtype=dtype) for _ in range(num_row_partitions)]
        return DynamicRaggedShape.Spec(row_partitions=row_partitions, static_inner_shape=tensor_shape.TensorShape(None), dtype=dtype)
    if shape.rank <= 1:
        if num_row_partitions:
            raise ValueError('num_row_partitions should be zero ' + 'if shape is a scalar or vector.')
        return DynamicRaggedShape.Spec(row_partitions=[], static_inner_shape=shape, dtype=dtype)
    if shape.rank <= num_row_partitions:
        raise ValueError('num_row_partitions must be less than rank')
    num_elements_so_far = tensor_shape.dimension_value(shape[0])
    rp_specs = []
    for i in range(num_row_partitions):
        current_dim = tensor_shape.dimension_value(shape[i + 1])
        if current_dim is None or num_elements_so_far is None:
            nvals = None
        else:
            nvals = num_elements_so_far * current_dim
        rp_specs.append(RowPartitionSpec(nrows=num_elements_so_far, nvals=nvals, uniform_row_length=current_dim, dtype=dtype))
        num_elements_so_far = nvals
    static_inner_shape = tensor_shape.TensorShape([num_elements_so_far]) + shape[num_row_partitions + 1:]
    return DynamicRaggedShape.Spec(row_partitions=rp_specs, static_inner_shape=static_inner_shape, dtype=dtype)

Creates a `DynamicRaggedShape.Spec` corresponding to a `tf.TensorShape`.

It is assumed that this is a `tf.TensorShape` coming from a
`tf.TensorSpec`, not from `RaggedTensor.shape`.

In addition to the shape, we need to know the number of row partitions,
and the dtype used in the shape (tf.int32 or tf.int64).

Within the dimensions that are partitioned, all dimensions are assumed
to be uniform.

Args:
  shape: a TensorShape.
  num_row_partitions: the ragged rank of the RaggedShape.
  dtype: the dtype of the shape (not the tensor); tf.int64 or tf.int32.

Returns:
  a DynamicRaggedShape.Spec representing a TensorShape.


def container(self, container_name) -> Iterator[str]:
    original_container = self._container
    self._container = container_name
    try:
        yield self._container
    finally:
        self._container = original_container

Returns a context manager that specifies the resource container to use.

Stateful operations, such as variables and queues, can maintain their
states on devices so that they can be shared by multiple processes.
A resource container is a string name under which these stateful
operations are tracked. These resources can be released or cleared
with `tf.Session.reset()`.

For example:

```python
with g.container('experiment0'):
  # All stateful Operations constructed in this context will be placed
  # in resource container "experiment0".
  v1 = tf.Variable([1.0])
  v2 = tf.Variable([2.0])
  with g.container("experiment1"):
    # All stateful Operations constructed in this context will be
    # placed in resource container "experiment1".
    v3 = tf.Variable([3.0])
    q1 = tf.queue.FIFOQueue(10, tf.float32)
  # All stateful Operations constructed in this context will be
  # be created in the "experiment0".
  v4 = tf.Variable([4.0])
  q1 = tf.queue.FIFOQueue(20, tf.float32)
  with g.container(""):
    # All stateful Operations constructed in this context will be
    # be placed in the default resource container.
    v5 = tf.Variable([5.0])
    q3 = tf.queue.FIFOQueue(30, tf.float32)

# Resets container "experiment0", after which the state of v1, v2, v4, q1
# will become undefined (such as uninitialized).
tf.Session.reset(target, ["experiment0"])
```

Args:
  container_name: container name string.

Returns:
  A context manager for defining resource containers for stateful ops,
    yields the container name.


def crop_image_to_patches(self, images: np.ndarray, min_patches: int, max_patches: int, use_thumbnail: bool=True, patch_size: Optional[Union[Tuple, int, dict]]=None, data_format: ChannelDimension=None):
    if data_format is None:
        data_format = infer_channel_dimension_format(images)
    images = to_channel_dimension_format(images, ChannelDimension.FIRST, data_format)
    patch_size_height, patch_size_width = (patch_size['height'], patch_size['width'])
    original_height, original_width = images.shape[-2:]
    num_columns, num_rows = get_optimal_tiled_canvas((original_height, original_width), (patch_size_height, patch_size_width), min_patches, max_patches)
    target_width = patch_size_width * num_columns
    target_height = patch_size_height * num_rows
    num_blocks = num_columns * num_rows
    resized_image = self.resize(images, {'height': target_height, 'width': target_width}, data_format=ChannelDimension.FIRST, input_data_format=ChannelDimension.FIRST)
    processed_images = []
    for i in range(num_blocks):
        column = i % num_columns
        row = i // num_columns
        box = (column * patch_size_width, row * patch_size_height, (column + 1) * patch_size_width, (row + 1) * patch_size_height)
        patch_image = resized_image[..., box[1]:box[3], box[0]:box[2]]
        patch_image = to_channel_dimension_format(patch_image, data_format, ChannelDimension.FIRST)
        processed_images.append(patch_image)
    if use_thumbnail and len(processed_images) != 1:
        thumbnail_img = self.resize(images, patch_size, data_format=data_format, input_data_format=ChannelDimension.FIRST)
        processed_images.append(thumbnail_img)
    return processed_images

Crop the image to patches and return a list of cropped images.
The number of patches and their grid arrangement are determined by the original image size,
the target patch size and the minimum and maximum number of patches.
The aspect ratio of the patches grid is chosen to be the closest to the original image aspect ratio.

Args:
    images (`np.ndarray`):
        The image to be cropped.
    min_patches (`int`):
        The minimum number of patches to be extracted from the image.
    max_patches (`int`):
        The maximum number of patches to be extracted from the image.
    use_thumbnail (`bool`, *optional*, defaults to `True`):
        Whether to add a thumbnail image to the list of cropped patches.
    patch_size (`int`, `Tuple[int, int]`, `dict`, *optional*):
        The size of the output patches.
    data_format (`ChannelDimension`, *optional*):
        The format of the image data. If `None`, the format is inferred from the input image.

Returns:
    List[`PIL.Image.Image`] or List[np.ndarray]: The list of cropped images.


def actual_actual_isda(*, start_date, end_date, schedule_info=None, dtype=None, name=None):
    del schedule_info
    with tf.name_scope(name or 'actual_actual_isda'):
        end_date = dt.convert_to_date_tensor(end_date)
        start_date = dt.convert_to_date_tensor(start_date)
        dtype = dtype or tf.float32
        days_in_leap_years, days_in_nonleap_years = du.days_in_leap_and_nonleap_years_between(start_date, end_date)
        days_in_leap_years = tf.cast(days_in_leap_years, dtype=dtype)
        days_in_nonleap_years = tf.cast(days_in_nonleap_years, dtype=dtype)
        return days_in_leap_years / 366 + days_in_nonleap_years / 365

Computes the year fraction between the specified dates.

Computes the year fraction between the dates by dividing the actual number of
days in a leap year by 366 and the actual number of days in a standard year by
365.

When determining whether a leap day is contained in the date range,
'start_date' is excluded and 'end_date' is included.

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

https://en.wikipedia.org/wiki/Day_count_convention#Actual/Actual_ISDA

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

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


def weight_memory_size(weights):
    unique_weights = {id(w): w for w in weights}.values()
    total_memory_size = 0
    for w in unique_weights:
        total_memory_size += _compute_memory_size(w.shape, w.dtype)
    return total_memory_size / 8

Compute the memory footprint for weights based on their dtypes.

Args:
    weights: An iterable contains the weights to compute weight size.

Returns:
    The total memory size (in Bytes) of the weights.


def forward(self, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.Tensor]=None, head_mask: Optional[torch.Tensor]=None, inputs_embeds: Optional[torch.FloatTensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[Tuple, BaseModelOutput]:
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    if input_ids is not None and inputs_embeds is not None:
        raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time')
    elif input_ids is not None:
        input = input_ids
        input_shape = input.shape
        input_ids = input_ids.view(-1, input_shape[-1])
    elif inputs_embeds is not None:
        input = inputs_embeds[:, :, -1]
    else:
        raise ValueError('You have to specify either input_ids or inputs_embeds')
    if inputs_embeds is None:
        inputs_embeds = self.embed_tokens(input_ids)
    embed_pos = self.embed_positions(input)
    hidden_states = inputs_embeds + embed_pos.to(inputs_embeds.device)
    hidden_states = self.layernorm_embedding(hidden_states)
    hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
    attention_mask = self._update_full_mask(attention_mask, inputs_embeds)
    encoder_states = () if output_hidden_states else None
    all_attentions = () if output_attentions else None
    if head_mask is not None:
        if head_mask.size()[0] != len(self.layers):
            raise ValueError(f'The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}.')
    for idx, encoder_layer in enumerate(self.layers):
        if output_hidden_states:
            encoder_states = encoder_states + (hidden_states,)
        to_drop = False
        if self.training:
            dropout_probability = torch.rand([])
            if dropout_probability < self.layerdrop:
                to_drop = True
        if to_drop:
            layer_outputs = (None, None)
        else:
            if self.gradient_checkpointing and self.training:
                layer_outputs = self._gradient_checkpointing_func(encoder_layer.__call__, hidden_states, attention_mask, head_mask[idx] if head_mask is not None else None, output_attentions)
            else:
                layer_outputs = encoder_layer(hidden_states, attention_mask, layer_head_mask=head_mask[idx] if head_mask is not None else None, output_attentions=output_attentions)
            hidden_states = layer_outputs[0]
        if output_attentions:
            all_attentions = all_attentions + (layer_outputs[1],)
    hidden_states = self.layer_norm(hidden_states)
    if output_hidden_states:
        encoder_states = encoder_states + (hidden_states,)
    if not return_dict:
        return tuple((v for v in [hidden_states, encoder_states, all_attentions] if v is not None))
    return BaseModelOutput(last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions)

Args:
    input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
        Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you
        provide it.

        Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
        [`PreTrainedTokenizer.__call__`] for details.

        [What are input IDs?](../glossary#input-ids)
    attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
        Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:

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

        [What are attention masks?](../glossary#attention-mask)
    head_mask (`torch.Tensor` of shape `(encoder_layers, encoder_attention_heads)`, *optional*):
        Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:

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

    inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
        Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
        This is useful if you want more control over how to convert `input_ids` indices into associated vectors
        than the model's internal embedding lookup matrix.
    output_attentions (`bool`, *optional*):
        Whether or not to return the attentions tensors of all attention layers. See `attentions` under
        returned tensors for more detail.
    output_hidden_states (`bool`, *optional*):
        Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
        for more detail.
    return_dict (`bool`, *optional*):
        Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.


def _num_present(losses, weights, per_batch=False):
    if isinstance(weights, float) and weights != 0.0 or (context.executing_eagerly() and weights._rank() == 0 and (not math_ops.equal(weights, 0.0))):
        return _num_elements(losses)
    with ops.name_scope(None, 'num_present', (losses, weights)) as scope:
        weights = math_ops.cast(weights, dtype=dtypes.float32)
        present = array_ops.where(math_ops.equal(weights, 0.0), array_ops.zeros_like(weights), array_ops.ones_like(weights))
        present = weights_broadcast_ops.broadcast_weights(present, losses)
        if per_batch:
            return math_ops.reduce_sum(present, axis=math_ops.range(1, array_ops.rank(present)), keepdims=True, name=scope)
        return math_ops.reduce_sum(present, name=scope)

Computes the number of elements in the loss function induced by `weights`.

A given weights tensor induces different numbers of usable elements in the
`losses` tensor. The `weights` tensor is broadcast across `losses` for all
possible dimensions. For example, if `losses` is a tensor of dimension
`[4, 5, 6, 3]` and `weights` is a tensor of shape `[4, 5]`, then `weights` is,
in effect, tiled to match the shape of `losses`. Following this effective
tile, the total number of present elements is the number of non-zero weights.

Args:
  losses: `Tensor` of shape `[batch_size, d1, ... dN]`.
  weights: `Tensor` of shape `[]`, `[batch_size]` or
    `[batch_size, d1, ... dK]`, where K < N.
  per_batch: Whether to return the number of elements per batch or as a sum
    total.

Returns:
  The number of present (non-zero) elements in the losses tensor. If
    `per_batch` is `True`, the value is returned as a tensor of size
    `[batch_size]`. Otherwise, a single scalar tensor is returned.


def _to_tf_type(dtype):
    return dtypes.as_dtype(dtype)

Converts a native python or numpy type to TF DType.

Args:
  dtype: Could be a python type, a numpy type or a TF DType.

Returns:
  A tensorflow `DType`.


def _prepare_4d_causal_attention_mask_with_cache_position(attention_mask: torch.Tensor, sequence_length: int, target_length: int, dtype: torch.dtype, cache_position: torch.Tensor, batch_size: int, **kwargs):
    if attention_mask is not None and attention_mask.dim() == 4:
        causal_mask = attention_mask
    else:
        min_dtype = torch.finfo(dtype).min
        causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device)
        if sequence_length != 1:
            causal_mask = torch.triu(causal_mask, diagonal=1)
        causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1)
        causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
        if attention_mask is not None:
            causal_mask = causal_mask.clone()
            mask_length = attention_mask.shape[-1]
            padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(causal_mask.device)
            padding_mask = padding_mask == 0
            causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(padding_mask, min_dtype)
    return causal_mask

Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
`(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.

Args:
    attention_mask (`torch.Tensor`):
        A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape
        `(batch_size, 1, query_length, key_value_length)`.
    sequence_length (`int`):
        The sequence length being processed.
    target_length (`int`):
        The target length: when generating with static cache, the mask should be as long as the static cache,
        to account for the 0 padding, the part of the cache that is not filled yet.
    dtype (`torch.dtype`):
        The dtype to use for the 4D attention mask.
    cache_position (`torch.Tensor`):
        Indices depicting the position of the input sequence tokens in the sequence.
    batch_size (`torch.Tensor`):
        Batch size.


def inverse_stft(stfts, frame_length, frame_step, fft_length=None, window_fn=window_ops.hann_window, name=None):
    with ops.name_scope(name, 'inverse_stft', [stfts]):
        stfts = ops.convert_to_tensor(stfts, name='stfts')
        stfts.shape.with_rank_at_least(2)
        frame_length = ops.convert_to_tensor(frame_length, name='frame_length')
        frame_length.shape.assert_has_rank(0)
        frame_step = ops.convert_to_tensor(frame_step, name='frame_step')
        frame_step.shape.assert_has_rank(0)
        if fft_length is None:
            fft_length = _enclosing_power_of_two(frame_length)
        else:
            fft_length = ops.convert_to_tensor(fft_length, name='fft_length')
            fft_length.shape.assert_has_rank(0)
        real_frames = fft_ops.irfft(stfts, [fft_length])
        frame_length_static = tensor_util.constant_value(frame_length)
        if frame_length_static is None or real_frames.shape.ndims is None or real_frames.shape.as_list()[-1] is None:
            real_frames = real_frames[..., :frame_length]
            real_frames_rank = array_ops.rank(real_frames)
            real_frames_shape = array_ops.shape(real_frames)
            paddings = array_ops.concat([array_ops.zeros([real_frames_rank - 1, 2], dtype=frame_length.dtype), [[0, math_ops.maximum(0, frame_length - real_frames_shape[-1])]]], 0)
            real_frames = array_ops.pad(real_frames, paddings)
        elif real_frames.shape.as_list()[-1] > frame_length_static:
            real_frames = real_frames[..., :frame_length_static]
        elif real_frames.shape.as_list()[-1] < frame_length_static:
            pad_amount = frame_length_static - real_frames.shape.as_list()[-1]
            real_frames = array_ops.pad(real_frames, [[0, 0]] * (real_frames.shape.ndims - 1) + [[0, pad_amount]])
        if frame_length_static is not None and real_frames.shape.ndims is not None:
            real_frames.set_shape([None] * (real_frames.shape.ndims - 1) + [frame_length_static])
        if window_fn is not None:
            window = window_fn(frame_length, dtype=stfts.dtype.real_dtype)
            real_frames *= window
        return reconstruction_ops.overlap_and_add(real_frames, frame_step)

Computes the inverse [Short-time Fourier Transform][stft] of `stfts`.

To reconstruct an original waveform, a complementary window function should
be used with `inverse_stft`. Such a window function can be constructed with
`tf.signal.inverse_stft_window_fn`.
Example:

```python
frame_length = 400
frame_step = 160
waveform = tf.random.normal(dtype=tf.float32, shape=[1000])
stft = tf.signal.stft(waveform, frame_length, frame_step)
inverse_stft = tf.signal.inverse_stft(
    stft, frame_length, frame_step,
    window_fn=tf.signal.inverse_stft_window_fn(frame_step))
```

If a custom `window_fn` is used with `tf.signal.stft`, it must be passed to
`tf.signal.inverse_stft_window_fn`:

```python
frame_length = 400
frame_step = 160
window_fn = tf.signal.hamming_window
waveform = tf.random.normal(dtype=tf.float32, shape=[1000])
stft = tf.signal.stft(
    waveform, frame_length, frame_step, window_fn=window_fn)
inverse_stft = tf.signal.inverse_stft(
    stft, frame_length, frame_step,
    window_fn=tf.signal.inverse_stft_window_fn(
       frame_step, forward_window_fn=window_fn))
```

Implemented with TPU/GPU-compatible ops and supports gradients.

Args:
  stfts: A `complex64`/`complex128` `[..., frames, fft_unique_bins]`
    `Tensor` of STFT bins representing a batch of `fft_length`-point STFTs
    where `fft_unique_bins` is `fft_length // 2 + 1`
  frame_length: An integer scalar `Tensor`. The window length in samples.
  frame_step: An integer scalar `Tensor`. The number of samples to step.
  fft_length: An integer scalar `Tensor`. The size of the FFT that produced
    `stfts`. If not provided, uses the smallest power of 2 enclosing
    `frame_length`.
  window_fn: A callable that takes a window length and a `dtype` keyword
    argument and returns a `[window_length]` `Tensor` of samples in the
    provided datatype. If set to `None`, no windowing is used.
  name: An optional name for the operation.

Returns:
  A `[..., samples]` `Tensor` of `float32`/`float64` signals representing
  the inverse STFT for each input STFT in `stfts`.

Raises:
  ValueError: If `stfts` is not at least rank 2, `frame_length` is not scalar,
    `frame_step` is not scalar, or `fft_length` is not scalar.

[stft]: https://en.wikipedia.org/wiki/Short-time_Fourier_transform


def min_max_variable_partitioner(max_partitions=1, axis=0, min_slice_size=256 << 10, bytes_per_string_element=16):

    def _partitioner(shape, dtype):
        """Partitioner that partitions list for a variable of given shape and type.

    Ex: Consider partitioning a variable of type float32 with
      shape=[1024, 1024].
      If `max_partitions` >= 16, this function would return
        [(1024 * 1024 * 4) / (256 * 1024), 1] = [16, 1].
      If `max_partitions` < 16, this function would return
        [`max_partitions`, 1].

    Args:
      shape: Shape of the variable.
      dtype: Type of the variable.

    Returns:
      List of partitions for each axis (currently only one axis can be
      partitioned).

    Raises:
      ValueError: If axis to partition along does not exist for the variable.
    """
        if axis >= len(shape):
            raise ValueError(f'Cannot partition variable along axis {axis} when shape is only {shape}')
        dtype = dtypes.as_dtype(dtype)
        if dtype.base_dtype == dtypes.string:
            bytes_per_element = bytes_per_string_element
        else:
            bytes_per_element = dtype.size
        total_size_bytes = shape.num_elements() * bytes_per_element
        partitions = total_size_bytes / min_slice_size
        partitions_list = [1] * len(shape)
        partitions_list[axis] = max(1, min(shape.dims[axis].value, max_partitions, int(math.ceil(partitions))))
        return partitions_list
    return _partitioner

Partitioner to allocate minimum size per slice.

Returns a partitioner that partitions the variable of given shape and dtype
such that each partition has a minimum of `min_slice_size` slice of the
variable. The maximum number of such partitions (upper bound) is given by
`max_partitions`.

Args:
  max_partitions: Upper bound on the number of partitions. Defaults to 1.
  axis: Axis along which to partition the variable. Defaults to 0.
  min_slice_size: Minimum size of the variable slice per partition. Defaults
    to 256K.
  bytes_per_string_element: If the `Variable` is of type string, this provides
    an estimate of how large each scalar in the `Variable` is.

Returns:
  A partition function usable as the `partitioner` argument to
  `variable_scope` and `get_variable`.


def sigmoid_cross_entropy_loss(inputs: torch.Tensor, labels: torch.Tensor, num_masks: int) -> torch.Tensor:
    criterion = nn.BCEWithLogitsLoss(reduction='none')
    cross_entropy_loss = criterion(inputs, labels)
    loss = cross_entropy_loss.mean(1).sum() / num_masks
    return loss

Args:
    inputs (`torch.Tensor`):
        A float tensor of arbitrary shape.
    labels (`torch.Tensor`):
        A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs
        (0 for the negative class and 1 for the positive class).

Returns:
    loss (`torch.Tensor`): The computed loss.


def _create_grad_indexed_slices_init(grad_output_slices, forward_input):
    assert isinstance(grad_output_slices, indexed_slices.IndexedSlices)
    assert isinstance(forward_input, tensor.Tensor)
    values_out = grad_output_slices.values
    indices_out = grad_output_slices.indices
    if values_out.shape.is_fully_defined():
        values_shape = tensor_shape.TensorShape([0] + values_out.shape.as_list()[1:])
        values = array_ops.zeros(values_shape, dtype=values_out.dtype, name='values_init')
    else:
        if forward_input.dtype == dtypes.resource:
            forward_shape = gen_resource_variable_ops.variable_shape(forward_input)
        else:
            forward_shape = array_ops.shape(forward_input)
        values_shape = array_ops.concat([[0], forward_shape[1:]], 0)
        values = array_ops.zeros(values_shape, dtype=values_out.dtype, name='values_init')
    indices = constant_op.constant([], indices_out.dtype, name='indices_init')
    if forward_input.dtype == dtypes.resource:
        shape = gen_resource_variable_ops.variable_shape(forward_input, name='shape_init')
    else:
        shape = array_ops.shape(forward_input, name='shape_init')
    return indexed_slices.IndexedSlices(values=values, indices=indices, dense_shape=shape)

Creates an IndexedSlices to pass as input to the while grad function.

Args:
  grad_output_slices: IndexedSlices. The corresponding while grad function
    output.
  forward_input: Tensor. The corresponding input to the forward while op.

Returns:
  Zeros IndexedSlices, created in current Graph.


def commandline_parser(parser=None, arguments=None):
    if parser is None:
        parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter, description=textwrap.dedent('        Command line to execute all tasks in a recipe once. ( Common Entry Point )\n\n        This script dispatches all the tasks in a JSON recipe to handlers in sequence.\n        For each task, it calls a subprocess to execute the JSON instructions, waits\n        for the process to complete and dispatches the next task, until all tasks are\n        complete or a critical failure ( exception ) is raised.\n\n        If an exception is raised in any task, all following tasks are not executed by design.\n\n        Example: python run.py [path to recipe file]\n        Caution: This script does NOT check if the last job finished, potentially causing overruns.\n        Notes:\n          - To avoid running the entire script when debugging a single task, the command line\n            can easily replace "all" with the name of any "task" in the json.  For example\n            python tool/recipe.py scripts/say_hello.json\n\n          - Can be easily replaced with the following to run only the "hello" task:\n            python task/hello/run.py scripts/say_hello.json\n\n          - Or specified further to run only the second hello task:\n            python task/hello/run.py scripts/say_hello.json -i 2\n\n    '))
    if arguments is None:
        parser.add_argument('json', help='Path to recipe json file to load.')
    elif '-j' in arguments:
        parser.add_argument('--json', '-j', help='Path to recipe json file to load.')
    if arguments is None or '-p' in arguments:
        parser.add_argument('--project', '-p', help='Cloud ID of Google Cloud Project.', default=None)
    if arguments is None or '-k' in arguments:
        parser.add_argument('--key', '-k', help='API Key of Google Cloud Project.', default=None)
    if arguments is None or '-u' in arguments:
        parser.add_argument('--user', '-u', help='Path to USER credentials json file.', default=None)
    if arguments is None or '-s' in arguments:
        parser.add_argument('--service', '-s', help='Path to SERVICE credentials json file.', default=None)
    if arguments is None or '-c' in arguments:
        parser.add_argument('--client', '-c', help='Path to CLIENT credentials json file.', default=None)
    if arguments is None or '-t' in arguments:
        parser.add_argument('--task', '-t', help='Task number of the task to run starting at 1.', default=None, type=int)
    if arguments is None or '-v' in arguments:
        parser.add_argument('--verbose', '-v', help='Print all the steps as they happen.', action='store_true')
    if arguments is None or '-f' in arguments:
        parser.add_argument('--force', '-force', help='Not used but included for compatiblity with another script.', action='store_true')
    if arguments is None or '-tp' in arguments:
        parser.add_argument('--trace_print', '-tp', help='Execution trace written to stdout.', action='store_true')
    if arguments is None or '-tf' in arguments:
        parser.add_argument('--trace_file', '-tf', help='Execution trace written to file.', action='store_true')
    if arguments is None or '-ni' in arguments:
        parser.add_argument('--no_input', '-ni', help='Raise exception if fields requiring input are in recipe.', action='store_true')
    return parser

Used in StarThinker scripts as entry point for command line calls.

Defines standard parameters used by almost every entry point.

Usage example:

```
import argparse
from starthinker.util.configuration import commandline_parser

if __name__ == "__main__":

  # custom parameters
  parser = argparse.ArgumentParser()
  parser.add_argument('custom', help='custom parameter to be added.')

  # initialize project
  commandline_parser(parser=parser, ['-c', '-u'])

  # access arguments
  print(args.client)
```

Args:
  * parser: (ArgumentParser) optional custom argument parser
  * arguments: (String) optional list of parameters to use when invoking, all set if None

Returns:
  ArgumentParser - parser with added parameters


def create_column_token_type_ids_from_sequences(self, query_ids: List[int], table_values: List[TableValue]) -> List[int]:
    table_column_ids = list(zip(*table_values))[1] if table_values else []
    return [0] * (1 + len(query_ids) + 1) + list(table_column_ids)

Creates the column token type IDs according to the query token IDs and a list of table values.

Args:
    query_ids (`List[int]`): list of token IDs corresponding to the ID.
    table_values (`List[TableValue]`): lift of table values, which are named tuples containing the
        token value, the column ID and the row ID of said token.

Returns:
    `List[int]`: List of ints containing the column token type IDs values.


def forward(self, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.FloatTensor]=None, token_type_ids: Optional[torch.LongTensor]=None, pixel_values: Optional[torch.FloatTensor]=None, pixel_mask: Optional[torch.LongTensor]=None, head_mask: Optional[torch.FloatTensor]=None, inputs_embeds: Optional[torch.FloatTensor]=None, image_embeds: Optional[torch.FloatTensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=True, return_dict: Optional[bool]=None, return_loss: Optional[bool]=None) -> Union[BridgeTowerContrastiveOutput, Tuple[torch.FloatTensor]]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    outputs = self.bridgetower(input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, pixel_values=pixel_values, pixel_mask=pixel_mask, head_mask=head_mask, inputs_embeds=inputs_embeds, image_embeds=image_embeds, output_attentions=output_attentions, output_hidden_states=True, return_dict=return_dict)
    pooler_output = outputs.pooler_output if return_dict else outputs[2]
    hidden_states_txt, hidden_states_img, hidden_states_cross_modal = outputs.hidden_states if return_dict else outputs[3]
    text_embeds = hidden_states_txt[-1]
    image_embeds = hidden_states_img[-1]
    image_embeds_with_ln = self.bridgetower.vision_model.visual.forward_post(image_embeds)
    image_token_type_embeddings = self.bridgetower.token_type_embeddings(torch.full((1,), 1, dtype=torch.long, device=self.bridgetower.token_type_embeddings.weight.device)).expand_as(image_embeds_with_ln)
    image_embeds = self.bridgetower.cross_modal_image_transform(image_embeds_with_ln) + image_token_type_embeddings
    text_embeds = nn.functional.normalize(self.itc_text_head(text_embeds[:, 0, :]), dim=-1, p=2)
    image_embeds = nn.functional.normalize(self.itc_image_head(image_embeds[:, 0, :]), dim=-1, p=2).to(device=text_embeds.device)
    cross_embeds = nn.functional.normalize(self.itc_cross_modal_head(pooler_output), dim=-1, p=2).to(device=text_embeds.device)
    logits = torch.stack([text_embeds, image_embeds, cross_embeds], dim=-2)
    logit_scale = self.logit_scale.exp().to(device=text_embeds.device)
    logits_text_to_image = torch.matmul(text_embeds, image_embeds.t()) * logit_scale
    logits_text_to_cross = torch.matmul(text_embeds, cross_embeds.t()) * logit_scale
    logits_image_to_cross = torch.matmul(image_embeds, cross_embeds.t()) * logit_scale
    itc_loss = None
    if return_loss:
        labels = torch.arange(len(logits), device=logits.device)
        text_to_image_loss = nn.functional.cross_entropy(logits_text_to_image, labels)
        text_to_cross_loss = nn.functional.cross_entropy(logits_text_to_cross, labels)
        image_to_cross_loss = nn.functional.cross_entropy(logits_image_to_cross, labels)
        itc_loss = (text_to_image_loss + text_to_cross_loss + image_to_cross_loss) / 3.0
    if not return_dict:
        output = (logits, text_embeds, image_embeds, cross_embeds) + outputs[3:]
        return (itc_loss,) + output if itc_loss is not None else output
    return BridgeTowerContrastiveOutput(loss=itc_loss, logits=logits, text_embeds=text_embeds, image_embeds=image_embeds, cross_embeds=cross_embeds, hidden_states=outputs.hidden_states, attentions=outputs.attentions)

image_embeds (`torch.FloatTensor` of shape `(batch_size, num_patches, hidden_size)`, *optional*):
    Optionally, instead of passing `pixel_values`, you can choose to directly pass an embedded representation.
    This is useful if you want more control over how to convert `pixel_values` into patch embeddings.
return_loss (`bool`, *optional*):
    Whether or not to return the contrastive loss.

Examples:

```python
>>> from transformers import BridgeTowerProcessor, BridgeTowerForContrastiveLearning
>>> import requests
>>> from PIL import Image
>>> import torch

>>> image_urls = [
...     "https://farm4.staticflickr.com/3395/3428278415_81c3e27f15_z.jpg",
...     "http://images.cocodataset.org/val2017/000000039769.jpg",
... ]
>>> texts = ["two dogs in a car", "two cats sleeping on a couch"]
>>> images = [Image.open(requests.get(url, stream=True).raw) for url in image_urls]

>>> processor = BridgeTowerProcessor.from_pretrained("BridgeTower/bridgetower-large-itm-mlm-itc")
>>> model = BridgeTowerForContrastiveLearning.from_pretrained("BridgeTower/bridgetower-large-itm-mlm-itc")

>>> inputs = processor(images, texts, padding=True, return_tensors="pt")
>>> loss = model(**inputs, return_loss=True).loss

>>> inputs = processor(images, texts[::-1], padding=True, return_tensors="pt")
>>> loss_swapped = model(**inputs, return_loss=True).loss

>>> print("Loss", round(loss.item(), 4))
Loss 0.0019

>>> print("Loss with swapped images", round(loss_swapped.item(), 4))
Loss with swapped images 2.126
```


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)`


class ProphetNetDecoderModelOutput(ModelOutput):
    last_hidden_state: torch.FloatTensor
    last_hidden_state_ngram: Optional[torch.FloatTensor] = None
    past_key_values: Optional[Tuple[torch.FloatTensor]] = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    hidden_states_ngram: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None
    ngram_attentions: Optional[Tuple[torch.FloatTensor]] = None
    cross_attentions: Optional[Tuple[torch.FloatTensor]] = None

Base class for model's outputs that may also contain a past key/values (to speed up sequential decoding).

Args:
    last_hidden_state (`torch.FloatTensor` of shape `(batch_size, decoder_sequence_length, hidden_size)`):
        Sequence of main stream hidden-states at the output of the last layer of the decoder of the model.

        If `past_key_values` is used only the last hidden-state of the sequences of shape `(batch_size, 1,
        hidden_size)` is output.
    last_hidden_state_ngram (`torch.FloatTensor` of shape `(batch_size, ngram * decoder_sequence_length, config.vocab_size)`):
        Sequence of predict stream hidden-states at the output of the last layer of the decoder of the model.
    past_key_values (`List[torch.FloatTensor]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
        List of `torch.FloatTensor` of length `config.n_layers`, with each tensor of shape `(2, batch_size,
        num_attn_heads, decoder_sequence_length, embed_size_per_head)`).

        Contains pre-computed hidden-states (key and values in the attention blocks) of the decoder that can be
        used (see `past_key_values` input) to speed up sequential decoding.
    hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
        shape `(batch_size, decoder_sequence_length, hidden_size)`.

        Hidden-states of main stream of the decoder at the output of each layer plus the initial embedding outputs.
    ngram_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
        shape `(batch_size, ngram * decoder_sequence_length, hidden_size)`.

        Hidden-states of the predict stream of the decoder at the output of each layer plus the initial embedding
        outputs.
    attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
        Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_attn_heads,
        decoder_sequence_length, decoder_sequence_length)`.

        Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the
        self-attention heads.
    ngram_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
        Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_attn_heads,
        decoder_sequence_length, decoder_sequence_length)`.

        Attentions weights of the predict stream of the decoder, after the attention softmax, used to compute the
        weighted average in the
    cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
        Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_attn_heads,
        encoder_sequence_length, decoder_sequence_length)`.

        Attentions weights of the cross-attention layer of the decoder, after the attention softmax, used to
        compute the weighted average in the


def _preprocess_mask_arguments(config: PretrainedConfig, input_embeds: torch.Tensor, attention_mask: Optional[Union[torch.Tensor, BlockMask]], cache_position: torch.Tensor, past_key_values: Optional[Cache], layer_idx: Optional[int]) -> tuple[bool, Optional[Union[torch.Tensor, BlockMask]], int, int]:
    if isinstance(attention_mask, (torch.Tensor, BlockMask)) and len(attention_mask.shape) == 4:
        return (True, attention_mask, None, None)
    if config._attn_implementation not in ALL_MASK_ATTENTION_FUNCTIONS._global_mapping:
        return (True, None, None, None)
    if attention_mask is not None and attention_mask.ndim == 2:
        attention_mask = attention_mask.to(device=cache_position.device, dtype=torch.bool)
    if past_key_values is not None:
        kv_length, kv_offset = past_key_values.get_mask_sizes(cache_position, layer_idx)
    else:
        kv_length, kv_offset = (input_embeds.shape[1], 0)
    return (False, attention_mask, kv_length, kv_offset)

Perform some common pre-processing of the mask arguments we get from the modeling code. Mostly determine the
key-value length and offsets, and if we should early exit or not.

Args:
    config (`PretrainedConfig`):
        The model config.
    input_embeds (`torch.Tensor`):
        The input embeddings of shape (batch_size, query_length, hidden_dim). This is used only to infer the
        batch size, query length and dtype.
    attention_mask (`torch.Tensor`, optional):
        The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length).
        It can also be an already prepared 4D mask, in which case it is returned as-is.
    cache_position (`torch.Tensor`):
        A tensor of shape (query_length,) indicating the current indices of the input sequence elements.
    past_key_values (`Cache`, optional):
        The past key values, if we use a cache.
    layer_idx (`int`, optional):
        If `past_key_values` is not None, this is the layer index of the cache from which to get the key-value
        length and offset. Indeed, for hybrid caches, different layers may return different lengths.

Returns:
    early_exit (`bool`):
        Whether we should early exit mask creation, and return the mask as-is.
    attention_mask (`torch.Tensor` or `BlockMask` or `None`):
        The attention mask to either return immediately, or to use in downstream mask creation.
    kv_length (`int`):
        The size that the key and value states will have during the attention computation.
    kv_offset (`int`):
        An offset to indicate at which first position the key and values states will refer to.