def _benchmarkRunOpPrebuilt(self, name, target, iters):
    times = []
    with ops.Graph().as_default():
        v = variables.Variable(random_ops.random_normal([]))
        with session.Session(target) as sess:
            sess.run(v.initializer)
            runner = sess.make_callable(v.op)
            runner()
            for _ in range(iters):
                start_time = time.time()
                runner()
                end_time = time.time()
                times.append(end_time - start_time)
    print('%s %f' % (name, np.median(times)))
    self.report_benchmark(iters=1, wall_time=np.median(times), name=name)

Runs a microbenchmark to measure the cost of running an op.

    Reports the median cost of running a trivial (Variable) op.

Args:
        name: A human-readable name for logging the output.
        target: The session target to use for the benchmark.
        iters: The number of iterations to perform.


def __init__(self, weights='int8', activations=None, modules_to_not_convert: Optional[List]=None, **kwargs):
    self.quant_method = QuantizationMethod.QUANTO
    self.weights = weights
    self.activations = activations
    self.modules_to_not_convert = modules_to_not_convert
    self.post_init()

This is a wrapper class about all possible attributes and features that you can play with a model that has been
    loaded using `quanto`.

Args:
        weights (`str`, *optional*, defaults to `"int8"`):
            The target dtype for the weights after quantization. Supported values are ("float8","int8","int4","int2")
        activations (`str`, *optional*):
            The target dtype for the activations after quantization. Supported values are (None,"int8","float8")
        modules_to_not_convert (`list`, *optional*, default to `None`):
            The list of modules to not quantize, useful for quantizing models that explicitly require to have
            some modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers).


def make_raw_assign_fn(raw_assign_fn, use_handle=True):

    def assign_fn(var, value, use_locking=False, name=None, read_value=True):
        del use_locking
        handle = var.handle if use_handle else var
        with _maybe_enter_graph(handle), _maybe_on_device(var):
            op = raw_assign_fn(handle, ops.convert_to_tensor(value, dtype=var.dtype), name=name)
            with ops.control_dependencies([op]):
                if read_value:
                    return var._read_variable_op() if use_handle else var.read_value()
                else:
                    return op
    return assign_fn

Wrap `raw_assign_fn` with the proper graph context and device scope.

Args:
        raw_assign_fn: the function to be wrapped.
        use_handle: if True, the `raw_assign_fn` will be applied to the handle of a
        variable; otherwise it will be applied to the variable itself.

Returns:
        The wrapped function.


def forward(self, past_values: torch.Tensor, target_values: Optional[torch.Tensor]=None, past_observed_mask: Optional[bool]=None, output_hidden_states: Optional[bool]=None, output_attentions: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[tuple, PatchTSTForClassificationOutput]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    model_output = self.model(past_values=past_values, past_observed_mask=past_observed_mask, output_hidden_states=output_hidden_states, output_attentions=output_attentions, return_dict=True)
    y_hat = self.head(model_output.last_hidden_state)
    loss_val = None
    if target_values is not None:
        loss = nn.CrossEntropyLoss()
        loss_val = loss(y_hat, target_values)
    if not return_dict:
        outputs = (y_hat,) + model_output[1:-3]
        outputs = (loss_val,) + outputs if loss_val is not None else outputs
        return outputs
    return PatchTSTForClassificationOutput(loss=loss_val, prediction_logits=y_hat, hidden_states=model_output.hidden_states, attentions=model_output.attentions)

past_values (`torch.Tensor` of shape `(bs, sequence_length, num_input_channels)`, *required*):
    Input sequence to the model
    target_values (`torch.Tensor`, *optional*):
    Labels associates with the `past_values`
    past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*):
    Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected
    in `[0, 1]`:

    - 1 for values that are **observed**,
    - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros).

Example:
    ```python
    >>> from transformers import PatchTSTConfig, PatchTSTForClassification

    >>> # classification task with two input channel2 and 3 classes
    >>> config = PatchTSTConfig(
    ...     num_input_channels=2,
    ...     num_targets=3,
    ...     context_length=512,
    ...     patch_length=12,
    ...     stride=12,
    ...     use_cls_token=True,
    ... )
    >>> model = PatchTSTForClassification(config=config)

    >>> # during inference, one only provides past values
    >>> past_values = torch.randn(20, 512, 2)
    >>> outputs = model(past_values=past_values)
    >>> labels = outputs.prediction_logits
    ```


def with_transform(self, transform: MLTransformProvider):
    self._validate_transform(transform)
    self.transforms.append(transform)
    return self

Add a transform to the MLTransform pipeline.

Args:
        transform: A BaseOperation instance.

Returns:
        A MLTransform instance.


def from_np_datetimes(np_datetimes):
    ordinals = tf.constant(np_datetimes, dtype=tf.int32) + _ORDINAL_OF_1_1_1970
    return from_ordinals(ordinals, validate=False)

Creates DateTensor from a Numpy array of dtype datetime64.

Args:
        np_datetimes: Numpy array of dtype datetime64.

Returns:
        DateTensor object.

    #### Example

    ```python
    import datetime
    import numpy as np

    date_tensor_np = np.array(
        [[datetime.date(2019, 3, 25), datetime.date(2020, 6, 2)],
        [datetime.date(2020, 9, 15), datetime.date(2020, 12, 27)]],
        dtype=np.datetime64)

    date_tensor = tff.datetime.dates_from_np_datetimes(date_tensor_np)
    ```


def get_cross_attention_token_mask(input_ids: List[int], image_token_id: int) -> List[List[int]]:
    image_token_locations = [i for i, token in enumerate(input_ids) if token == image_token_id]
    if len(image_token_locations) == 0:
        return []
    if len(image_token_locations) == 1:
        return [[image_token_locations[0], -1]]
    vision_masks = [[loc1, loc2] for loc1, loc2 in zip(image_token_locations[:-1], image_token_locations[1:])]
    vision_masks.append([image_token_locations[-1], len(input_ids)])
    last_mask_end = vision_masks[-1][1]
    for vision_mask in vision_masks[::-1]:
        if vision_mask[0] == vision_mask[1] - 1:
            vision_mask[1] = last_mask_end
        last_mask_end = vision_mask[1]
    return vision_masks

Generate a cross-attention token mask for image tokens in the input sequence.

    This function identifies the positions of image tokens in the input sequence and creates
    a mask that defines which subsequent tokens each image token should attend to.

Args:
        input_ids (List[int]): A list of token ids representing the input sequence.
        image_token_id (int): The id of the token used to represent images in the sequence.

Returns:
        List[List[int]]: A list of [start, end] pairs, where each pair represents the range
        of tokens an image token should attend to.

Note:
        - If no image tokens are present, an empty list is returned.
        - For a single image token, it attends to all subsequent tokens until the end of the sequence.
        - For multiple image tokens, each attends to tokens up to the next image token or the end of the sequence.
        - Consecutive image tokens are treated as a group and attend to all subsequent tokens together.


def update(self, key_states: torch.Tensor, value_states: torch.Tensor, layer_idx: int, cache_kwargs: Optional[Dict[str, Any]]=None) -> Tuple[torch.Tensor, torch.Tensor]:
    key_states = key_states.to(self.key_cache[layer_idx].dtype)
    value_states = value_states.to(self.value_cache[layer_idx].dtype)
    if layer_idx == 0:
        self._seen_tokens += key_states.shape[-2]
        k_out = self.key_cache[0]
        v_out = self.value_cache[0]
    else:
        if self._prefetch_stream is not None:
            torch.cuda.default_stream(self.device).wait_stream(self._prefetch_stream)
        k_out = self._device_key_cache[layer_idx & 1]
        v_out = self._device_value_cache[layer_idx & 1]
    self._prefetch_layer(layer_idx + 1)
    cache_position = cache_kwargs.get('cache_position') if cache_kwargs is not None else None
    if cache_position is None:
        k_out.copy_(key_states)
        v_out.copy_(value_states)
        if layer_idx == 0:
            self.key_cache[layer_idx].copy_(key_states.to(self.offload_device))
            self.value_cache[layer_idx].copy_(value_states.to(self.offload_device))
    else:
        try:
            k_out.index_copy_(2, cache_position, key_states)
            v_out.index_copy_(2, cache_position, value_states)
        except NotImplementedError:
            k_out[:, :, cache_position] = key_states
            v_out[:, :, cache_position] = value_states
        if layer_idx != 0:
            cache_position = cache_position.to(self.offload_device)
            key_states = key_states.to(self.offload_device)
            value_states = value_states.to(self.offload_device)
            try:
                self.key_cache[layer_idx].index_copy_(2, cache_position, key_states)
                self.value_cache[layer_idx].index_copy_(2, cache_position, value_states)
            except NotImplementedError:
                self.key_cache[layer_idx][:, :, cache_position] = key_states
                self.value_cache[layer_idx][:, :, cache_position] = value_states
    return (k_out, v_out)

Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`.
    It is VERY important to index using a tensor, otherwise you introduce a copy to the device.

    Parameters:
        key_states (`torch.Tensor`):
            The new key states to cache.
        value_states (`torch.Tensor`):
            The new value states to cache.
        layer_idx (`int`):
            The index of the layer to cache the states for.
        cache_kwargs (`Dict[str, Any]`, *optional*):
            Additional arguments for the cache subclass. The `OffloadedStaticCache` needs the
            `cache_position` input to know how where to write in the cache.

Returns:
        A tuple containing the updated key and value states.


def to_tensors(self, value: Any) -> List[core.Tensor]:
    del value
    return []

Breaks down a value of this type into Tensors.

    For a TraceType instance, the number of tensors generated for corresponding
    value should be constant.

Args:
        value: A value belonging to this TraceType

Returns:
        List of Tensors.


def _create_valueset_codes_table_if_not_exists(self) -> bigquery.table.Table:
    schema = [bigquery.SchemaField('valueseturi', 'STRING', mode='REQUIRED'), bigquery.SchemaField('valuesetversion', 'STRING', mode='NULLABLE'), bigquery.SchemaField('system', 'STRING', mode='REQUIRED'), bigquery.SchemaField('code', 'STRING', mode='REQUIRED')]
    table = bigquery.Table(self._value_set_codes_table, schema=schema)
    table.clustering_fields = ['valueseturi', 'code']
    return self._client.create_table(table, exists_ok=True)

Creates a table for storing value set code mappings.

    Creates a table named after the `value_set_codes_table` provided at class
    initialization as described by
    https://github.com/FHIR/sql-on-fhir/blob/master/sql-on-fhir.md#valueset-support

    If the table already exists, no action is taken.

Returns:
        An bigquery.Table object representing the created table.


def _filter_ds(dataset, acceptance_dist_ds, initial_dist_ds, class_func, seed, name=None) -> DatasetV2:

    def maybe_warn_on_large_rejection(accept_dist, initial_dist):
        proportion_rejected = math_ops.reduce_sum((1 - accept_dist) * initial_dist)
        return cond.cond(math_ops.less(proportion_rejected, 0.5), lambda: accept_dist, lambda: logging_ops.Print(accept_dist, [proportion_rejected, initial_dist, accept_dist], message='Proportion of examples rejected by sampler is high: ', summarize=100, first_n=10))
    acceptance_dist_ds = DatasetV2.zip((acceptance_dist_ds, initial_dist_ds), name=name).map(maybe_warn_on_large_rejection, name=name)

    def _gather_and_copy(acceptance_prob, data):
        if isinstance(data, tuple):
            class_val = class_func(*data)
        else:
            class_val = class_func(data)
        return (class_val, array_ops.gather(acceptance_prob, class_val), data)
    current_probabilities_and_class_and_data_ds = DatasetV2.zip((acceptance_dist_ds, dataset), name=name).map(_gather_and_copy, name=name)

    def _reject(unused_class_val, p, unused_data):
        return random_ops.random_uniform([], seed=seed, dtype=p.dtype) < p
    filtered_ds = current_probabilities_and_class_and_data_ds.filter(_reject, name=name)
    return filtered_ds.map(lambda class_value, _, data: (class_value, data), name=name)

Filters a dataset based on per-class acceptance probabilities.

Args:
        dataset: The dataset to be filtered.
        acceptance_dist_ds: A dataset of acceptance probabilities.
        initial_dist_ds: A dataset of the initial probability distribution, given or
        estimated.
        class_func: A function mapping an element of the input dataset to a scalar
        `tf.int32` tensor. Values should be in `[0, num_classes)`.
        seed: (Optional.) Python integer seed for the resampler.
        name: (Optional.) A name for the tf.data operation.

Returns:
        A dataset of (class value, data) after filtering.


def testGradient(self, params, indices, expected_out, out_grad, expected_grad, params_ragged_rank=None):
    if context.executing_eagerly():
        return
    params = ragged_factory_ops.constant(params, dtype=dtypes.float32, ragged_rank=params_ragged_rank)
    indices = constant_op.constant(indices, dtype=dtypes.int32)
    out_ragged_rank = params.ragged_rank + indices.shape.ndims - 1
    out_grad = ragged_factory_ops.constant(out_grad, dtype=dtypes.float32, ragged_rank=out_ragged_rank)
    expected_out = ragged_factory_ops.constant(expected_out, dtype=dtypes.float32, ragged_rank=out_ragged_rank)
    expected_grad = ragged_factory_ops.constant(expected_grad, dtype=dtypes.float32, ragged_rank=params.ragged_rank)
    out = ragged_gather_ops.gather(params, indices)
    self.assertAllClose(out, expected_out)
    grads = gradients_impl.gradients(out.flat_values, params.nested_row_splits + (params.flat_values, indices), out_grad.flat_values)
    param_nested_splits_grads = grads[:-2]
    params_flat_values_grad = grads[-2]
    indices_grad = grads[-1]
    self.assertEqual(indices_grad, None)
    for splits_grad in param_nested_splits_grads:
        self.assertEqual(splits_grad, None)
    self.assertIsInstance(params_flat_values_grad, indexed_slices.IndexedSlices)
    params_flat_values_grad = ops.convert_to_tensor(params_flat_values_grad)
    params_grad = params.with_flat_values(params_flat_values_grad)
    self.assertAllClose(params_grad, expected_grad, atol=2e-06, rtol=2e-06)

Tests that ragged_gather generates the right gradient.

Args:
        params: The `params` that should be passed to `gather`.
        indices: The `indices` that should be passed to `gather`.
        expected_out: The expected value of `gather(params, indices)`.
        `expected_out.shape = indices.shape + params.shape[1:]`.
        out_grad: The value that should be fed in as the gradient for `out`
        when testing the gradient of `ragged_gather`.  Must have the same
        shape as `expected_out`.
        expected_grad: The expected gradient for that should be returned for
        `params`.  Must have hte same shape as `params`.
        params_ragged_rank: The ragged_rank of `params`.


def compute_token_logits(sequence_output, temperature, output_weights, output_bias):
    logits = (torch.einsum('bsj,j->bs', sequence_output, output_weights) + output_bias) / temperature
    return logits

Computes logits per token

Args:
        sequence_output (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
            Also known as last_hidden_state. Sequence of hidden-states at the output of the last layer of the model.
        temperature (`float`):
            Temperature for the Bernoulli distribution.
        output_weights (`torch.FloatTensor` of shape `(hidden_size,)`):
            Weights of the linear layer for cell selection.
        output_bias (`torch.FloatTensor` of shape `()`):
            Bias of the linear layer for cell selection

Returns:
        logits (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): Logits per token.


def pad_to_square(self, images: 'torch.Tensor', background_color: Union[int, Tuple[int, int, int]]=0) -> 'torch.Tensor':
    height, width = get_image_size(images, ChannelDimension.FIRST)
    if height == width:
        return images
    num_channels = images.shape[1] if len(images.shape) == 4 else images.shape[0]
    if isinstance(background_color, int):
        background_color = [background_color] + [0] * (num_channels - 1)
    elif len(background_color) != num_channels:
        raise ValueError(f'background_color must have no more than {num_channels} elements to match the number of channels')
    max_dim = max(height, width)
    paste_x_left = (max_dim - width) // 2
    paste_y_left = (max_dim - height) // 2
    paste_x_right = max_dim - width - paste_x_left
    paste_y_right = max_dim - height - paste_y_left
    padded_images = F.pad(images, padding=[paste_x_left, paste_y_left, paste_x_right, paste_y_right], fill=background_color)
    return padded_images

Pads an image to a square based on the longest edge.

Args:
        images (`np.ndarray`):
            The images to pad.
        background_color (`int` or `Tuple[int, int, int]`, *optional*, defaults to 0):
            The color to use for the padding. Can be an integer for single channel or a
            tuple of integers representing for multi-channel images. If passed as integer
            in mutli-channel mode, it will default to `0` in subsequent channels.

Returns:
        `torch.Tensor`: The padded images.


def __init__(self, default: typing.Optional[bool]=MISSING_VALUE, is_noneable: bool=False, frozen: bool=False):
    super().__init__(bool, default, is_noneable=is_noneable, frozen=frozen)

Constructor.

Args:
        default: Default value for the value spec.
        is_noneable: If True, None is acceptable.
        frozen: If True, values other than the default value is not accceptable.


def _pad(self, encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding], max_length: Optional[int]=None, padding_strategy: PaddingStrategy=PaddingStrategy.DO_NOT_PAD, pad_to_multiple_of: Optional[int]=None, padding_side: Optional[str]=None, return_attention_mask: Optional[bool]=None) -> dict:
    if return_attention_mask is None:
        return_attention_mask = 'attention_mask' in self.model_input_names
    if padding_strategy == PaddingStrategy.LONGEST:
        max_length = len(encoded_inputs['input_ids'])
    if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
        max_length = (max_length // pad_to_multiple_of + 1) * pad_to_multiple_of
    needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len(encoded_inputs['input_ids']) != max_length
    if return_attention_mask and 'attention_mask' not in encoded_inputs:
        encoded_inputs['attention_mask'] = [1] * len(encoded_inputs['input_ids'])
    if needs_to_be_padded:
        difference = max_length - len(encoded_inputs['input_ids'])
        padding_side = padding_side if padding_side is not None else self.padding_side
        if padding_side == 'right':
            if return_attention_mask:
                encoded_inputs['attention_mask'] = encoded_inputs['attention_mask'] + [0] * difference
            if 'token_type_ids' in encoded_inputs:
                encoded_inputs['token_type_ids'] = encoded_inputs['token_type_ids'] + [[self.pad_token_type_id] * 7] * difference
            if 'labels' in encoded_inputs:
                encoded_inputs['labels'] = encoded_inputs['labels'] + [0] * difference
            if 'numeric_values' in encoded_inputs:
                encoded_inputs['numeric_values'] = encoded_inputs['numeric_values'] + [float('nan')] * difference
            if 'numeric_values_scale' in encoded_inputs:
                encoded_inputs['numeric_values_scale'] = encoded_inputs['numeric_values_scale'] + [1.0] * difference
            if 'special_tokens_mask' in encoded_inputs:
                encoded_inputs['special_tokens_mask'] = encoded_inputs['special_tokens_mask'] + [1] * difference
            encoded_inputs['input_ids'] = encoded_inputs['input_ids'] + [self.pad_token_id] * difference
        elif padding_side == 'left':
            if return_attention_mask:
                encoded_inputs['attention_mask'] = [0] * difference + encoded_inputs['attention_mask']
            if 'token_type_ids' in encoded_inputs:
                encoded_inputs['token_type_ids'] = [[self.pad_token_type_id] * 7] * difference + encoded_inputs['token_type_ids']
            if 'labels' in encoded_inputs:
                encoded_inputs['labels'] = [0] * difference + encoded_inputs['labels']
            if 'numeric_values' in encoded_inputs:
                encoded_inputs['numeric_values'] = [float('nan')] * difference + encoded_inputs['numeric_values']
            if 'numeric_values_scale' in encoded_inputs:
                encoded_inputs['numeric_values_scale'] = [1.0] * difference + encoded_inputs['numeric_values_scale']
            if 'special_tokens_mask' in encoded_inputs:
                encoded_inputs['special_tokens_mask'] = [1] * difference + encoded_inputs['special_tokens_mask']
            encoded_inputs['input_ids'] = [self.pad_token_id] * difference + encoded_inputs['input_ids']
        else:
            raise ValueError('Invalid padding strategy:' + str(padding_side))
    return encoded_inputs

Pad encoded inputs (on left/right and up to predefined length or max length in the batch)

Args:
        encoded_inputs:
            Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`).
        max_length: maximum length of the returned list and optionally padding length (see below).
            Will truncate by taking into account the special tokens.
        padding_strategy: PaddingStrategy to use for padding.

            - PaddingStrategy.LONGEST Pad to the longest sequence in the batch
            - PaddingStrategy.MAX_LENGTH: Pad to the max length (default)
            - PaddingStrategy.DO_NOT_PAD: Do not pad
            The tokenizer padding sides are defined in self.padding_side:

                - 'left': pads on the left of the sequences
                - 'right': pads on the right of the sequences
        pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value.
            This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability
            `>= 7.5` (Volta).
        padding_side:
            The side on which the model should have padding applied. Should be selected between ['right', 'left'].
            Default value is picked from the class attribute of the same name.
        return_attention_mask:
            (optional) Set to False to avoid returning attention mask (default: set to model specifics)


def _lookup_dependency(self, name, cached_dependencies=None):
    if cached_dependencies:
        return cached_dependencies.get(name)
    return self._self_unconditional_dependency_names.get(name)

Look up a dependency by name.

    May be overridden to include conditional dependencies.

Args:
        name: The local name of the dependency.
        cached_dependencies: Optional dict containing all computed dependencies
        returned by `self._trackable_children()`.

Returns:
        A `Trackable` object, or `None` if no dependency by this name was
        found.


def __init__(self, **kwargs):
    super(_Merge, self).__init__(**kwargs)
    self.supports_masking = True

Initializes a Merge layer.

Args:
        **kwargs: standard layer keyword arguments.


def __init__(self, runtime_list, metric_id):
    value = self._prepare_runtime_metrics(runtime_list)
    submit_timestamp = time.time()
    label = runtime_list[0].key.metric.namespace + '_' + RUNTIME_METRIC
    super().__init__(submit_timestamp, metric_id, value, None, label)

The Distribution Metric in ready-to-publish format.

Args:
        runtime_list: list of distributions metrics from MetricResult
        with runtime name
        metric_id(uuid): unique id to identify test run


def __setattr__(self, name: str, value: Any) -> None:
    if name.startswith('_'):
        super().__setattr__(name, value)
    else:
        self[name] = value

Set attribute of this Dict.

    NOTE(daiyip): When setting attributes, public attributes (not started with
    '_') are set as dict fields, while private attributes (started with '_') are
    set on the object instance.

Args:
        name: Name of attribute.
        value: Value of attribute.


def pad(x, pad_width, mode='constant', constant_values=None):
    return Pad(pad_width, mode=mode)(x, constant_values=constant_values)

Pad a tensor.

Args:
        x: Tensor to pad.
        pad_width: Number of values padded to the edges of each axis.
            `((before_1, after_1), ...(before_N, after_N))` unique pad
            widths for each axis.
            `((before, after),)` yields same before and after pad for
            each axis.
            `(pad,)` or `int` is a shortcut for `before = after = pad`
            width for all axes.
        mode: One of `"constant"`, `"edge"`, `"linear_ramp"`,
            `"maximum"`, `"mean"`, `"median"`, `"minimum"`,
            `"reflect"`, `"symmetric"`, `"wrap"`, `"empty"`,
            `"circular"`. Defaults to `"constant"`.
        constant_values: value to pad with if `mode == "constant"`.
            Defaults to `0`. A `ValueError` is raised if not None and
            `mode != "constant"`.

Note:
        Torch backend only supports modes `"constant"`, `"reflect"`,
        `"symmetric"` and `"circular"`.
        Only Torch backend supports `"circular"` mode.

Note:
        Tensorflow backend only supports modes `"constant"`, `"reflect"`
        and `"symmetric"`.

Returns:
        Padded tensor.


def sparse_categorical_crossentropy(y_true, y_pred, from_logits=False, axis=-1):
    y_pred = tensor_conversion.convert_to_tensor_v2_with_dispatch(y_pred)
    y_true = math_ops.cast(y_true, y_pred.dtype)
    return backend.sparse_categorical_crossentropy(y_true, y_pred, from_logits=from_logits, axis=axis)

Computes the sparse categorical crossentropy loss.

    Standalone usage:

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

Args:
        y_true: Ground truth values.
        y_pred: The predicted values.
        from_logits: Whether `y_pred` is expected to be a logits tensor. By default,
        we assume that `y_pred` encodes a probability distribution.
        axis: Defaults to -1. The dimension along which the entropy is
        computed.

Returns:
        Sparse categorical crossentropy loss value.


def _get_dependent_variables(input_ops, output_ops):
    output_ops = nest.map_structure(gen_array_ops.identity, output_ops)
    inbetween_ops = op_selector.get_backward_walk_ops(seed_ops=output_ops, stop_at_ts=input_ops, inclusive=False, only_differentiable=True)
    var_ops = (op for op in inbetween_ops if op.type in VAR_OP_TYPES)
    var_names = (op.name for op in var_ops)
    tf_vars = (get_variable_by_name(var_name) for var_name in var_names)
    tf_vars = [v for v in tf_vars if v is not None]
    return tf_vars

Finds variables involved in the subgraph between input_ops and output_ops.

Args:
        input_ops: Flattened list of input ops
        output_ops: Flattened list of output ops

Returns:
        A list of variables


def _default_layout(self, layout: layout_lib.Layout):
    previous_default = None
    previous_graph_size = None
    graph = None
    self._register_mesh(layout.mesh)
    try:
        previous_default = self._current_output_layout
        self._current_output_layout = layout.to_string().encode('utf-8')
        _pywrap_dtensor_device.ExperimentalSetDefaultLayout(self._device_info, self._current_output_layout)
        if context.executing_eagerly():
            with ops.device(self.name):
                yield
        else:
            graph = ops.get_default_graph()
            previous_graph_size = len(graph.get_operations())
            yield
    finally:
        if graph is not None:
            for operation in graph.get_operations()[previous_graph_size:]:
                operation._set_attr('_layout', attr_value_pb2.AttrValue(list=attr_value_pb2.AttrValue.ListValue(s=[self._current_output_layout])))
                operation._set_attr('_mesh', attr_value_pb2.AttrValue(s=layout.mesh.to_string().encode('utf-8')))
        self._current_output_layout = previous_default
        if self._current_output_layout is None:
            _pywrap_dtensor_device.ExperimentalClearDefaultLayout(self._device_info)
        else:
            _pywrap_dtensor_device.ExperimentalSetDefaultLayout(self._device_info, self._current_output_layout.decode('utf-8'))

Sets a default output layout for all ops in the scope.

    Note: This is an internal helper method, which is not user facing api.

    Useful for requesting a specific layout for ops which would have no inferred
    layout, e.g. tf.zeros.

    Caveats:

    - Currently only affects the first output of an op. For Op with multiple
        outputs, this does not support yet.

    - All Ops in the scope will be attached with the same layout. This might not
        be valid as the rank is different. The current suggestion is: Try to wrap
        the raw op wheneven possible.

Args:
        layout: A Layout for the outputs of all operations in this scope.

Yields:
        Nothing.


def create_feed_dict_from_input_data(input_data: RepresentativeSample, signature_def: meta_graph_pb2.SignatureDef) -> Mapping[str, np.ndarray]:
    feed_dict = {}
    for input_key, input_value in input_data.items():
        input_tensor_name = signature_def.inputs[input_key].name
        value = input_value
        if isinstance(input_value, core.Tensor):
            value = input_value.eval()
        feed_dict[input_tensor_name] = value
    return feed_dict

Constructs a feed_dict from input data.

    Note: This function should only be used in graph mode.

    This is a helper function that converts an 'input key -> input value' mapping
    to a feed dict. A feed dict is an 'input tensor name -> input value' mapping
    and can be directly passed to the `feed_dict` argument of `sess.run()`.

Args:
        input_data: Input key -> input value mapping. The input keys should match
        the input keys of `signature_def`.
        signature_def: A SignatureDef representing the function that `input_data` is
        an input to.

Returns:
        Feed dict, which is intended to be used as input for `sess.run`. It is
        essentially a mapping: input tensor name -> input value. Note that the input
        value in the feed dict is not a `Tensor`.


def extract_model_metrics(model):
    if getattr(model, '_compile_metrics', None):
        return {m.name: m for m in model._compile_metric_functions}
    return None

Convert metrics from a Keras model `compile` API to dictionary.

    This is used for converting Keras models to SavedModels.

Args:
        model: A `tf.keras.Model` object.

Returns:
        Dictionary mapping metric names to metric instances. May return `None` if
        the model does not contain any metrics.


def _convert_bytes_to_cc_source(data, array_name, max_line_width=80, include_guard=None, include_path=None, use_tensorflow_license=False):
    starting_pad = '   '
    array_lines = []
    array_line = starting_pad
    for value in bytearray(data):
        if len(array_line) + 4 > max_line_width:
            array_lines.append(array_line + '\n')
            array_line = starting_pad
        array_line += ' 0x%02x,' % value
    if len(array_line) > len(starting_pad):
        array_lines.append(array_line + '\n')
    array_values = ''.join(array_lines)
    if include_guard is None:
        include_guard = 'TENSORFLOW_LITE_UTIL_' + array_name.upper() + '_DATA_H_'
    if include_path is not None:
        include_line = '#include "{include_path}"\n'.format(include_path=include_path)
    else:
        include_line = ''
    if use_tensorflow_license:
        license_text = '\n/* Copyright {year} The TensorFlow Authors. All Rights Reserved.\n\nLicensed under the Apache License, Version 2.0 (the "License");\nyou may not use this file except in compliance with the License.\nYou may obtain a copy of the License at\n\n    http://www.apache.org/licenses/LICENSE-2.0\n\nUnless required by applicable law or agreed to in writing, software\ndistributed under the License is distributed on an "AS IS" BASIS,\nWITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\nSee the License for the specific language governing permissions and\nlimitations under the License.\n==============================================================================*/\n'.format(year=datetime.date.today().year)
    else:
        license_text = ''
    source_template = '{license_text}\n// This is a binary file that has been converted into a C++ data array using the\n// //tensorflow/lite/experimental/acceleration/compatibility/convert_binary_to_cc_source.py\n// script. This form is useful for compiling into a binary to simplify\n// deployment on mobile devices\n\n{include_line}\n// We need to keep the data array aligned on some architectures.\n#ifdef __has_attribute\n#define HAVE_ATTRIBUTE(x) __has_attribute(x)\n#else\n#define HAVE_ATTRIBUTE(x) 0\n#endif\n#if HAVE_ATTRIBUTE(aligned) || (defined(__GNUC__) && !defined(__clang__))\n#define DATA_ALIGN_ATTRIBUTE __attribute__((aligned(16)))\n#else\n#define DATA_ALIGN_ATTRIBUTE\n#endif\n\nextern const unsigned char {array_name}[] DATA_ALIGN_ATTRIBUTE = {{\n{array_values}}};\nextern const int {array_name}_len = {array_length};\n'
    source_text = source_template.format(array_name=array_name, array_length=len(data), array_values=array_values, license_text=license_text, include_line=include_line)
    header_template = '\n{license_text}\n\n// This is a binary file that has been converted into a C++ data array using the\n// //tensorflow/lite/experimental/acceleration/compatibility/convert_binary_to_cc_source.py\n// script. This form is useful for compiling into a binary to simplify\n// deployment on mobile devices\n\n#ifndef {include_guard}\n#define {include_guard}\n\nextern const unsigned char {array_name}[];\nextern const int {array_name}_len;\n\n#endif  // {include_guard}\n'
    header_text = header_template.format(array_name=array_name, include_guard=include_guard, license_text=license_text)
    return (source_text, header_text)

Returns strings representing a C++ constant array containing `data`.

Args:
        data: Byte array that will be converted into a C++ constant.
        array_name: String to use as the variable name for the constant array.
        max_line_width: The longest line length, for formatting purposes.
        include_guard: Name to use for the include guard macro definition.
        include_path: Optional path to include in the source file.
        use_tensorflow_license: Whether to include the standard TensorFlow Apache2
        license in the generated files.

Returns:
        Text that can be compiled as a C++ source file to link in the data as a
        literal array of values.
        Text that can be used as a C++ header file to reference the literal array.


def recipe_video(config, auth_read, sheet, tab, project, dataset, table):
    sheets(config, {'__comment__': 'Copy the tamplate sheet to the users sheet.  If it already exists, nothing happens.', 'auth': auth_read, 'template': {'sheet': 'https://docs.google.com/spreadsheets/d/1BXRHWz-1P3gNS92WZy-3sPZslU8aalXa8heOgygWEFs/edit#gid=0', 'tab': 'Video'}, 'sheet': sheet, 'tab': tab})
    video(config, {'__comment__': 'Read video effects and values from sheet and/or bigquery.', 'auth': auth_read, 'sheets': {'sheet': sheet, 'tab': tab}, 'bigquery': {'project': project, 'dataset': dataset, 'table': table}})

Add images, text, and audio to videos.

Args:
        auth_read (authentication) - Credentials used for reading data.
        sheet (string) - Name or URL of sheet.
        tab (string) - Name of sheet tab.
        project (string) - Google Cloud Project Identifier.
        dataset (string) - Name of dataset.
        table (string) - Name of table.


def __init__(self, nbits: int=4, group_size: int=64, view_as_float: bool=False, axis: Optional[int]=None, dynamic_config: Optional[dict]=None, skip_modules: List[str]=['lm_head'], **kwargs):
    if is_hqq_available():
        from hqq.core.quantize import BaseQuantizeConfig as HQQBaseQuantizeConfig
    else:
        raise ImportError('A valid HQQ version (>=0.2.1) is not available. Please follow the instructions to install it: `https://github.com/mobiusml/hqq/`.')
    for deprecated_key in ['quant_zero', 'quant_scale', 'offload_meta']:
        if deprecated_key in kwargs:
            logger.info(deprecated_key + ' is deprecated. This parameter will be ignored in quantization settings.')
    if axis is None:
        axis = 1
        logger.info('Setting axis=1 as faster backends such as TorchAO or BitBlas are only compatible with it.')
    if axis not in [0, 1]:
        raise ValueError('Invalid axis value. Only 0 and 1 are allowed.')
    if dynamic_config is not None:
        self.quant_config = {}
        for key in dynamic_config:
            self.quant_config[key] = HQQBaseQuantizeConfig(**dynamic_config[key])
    else:
        self.quant_config = HQQBaseQuantizeConfig(**{'nbits': nbits, 'group_size': group_size, 'view_as_float': view_as_float, 'axis': axis})
    self.quant_method = QuantizationMethod.HQQ
    self.skip_modules = skip_modules
    self.post_init()

This is wrapper around hqq's BaseQuantizeConfig.

Args:
        nbits (`int`, *optional*, defaults to 4):
            Number of bits. Supported values are (8, 4, 3, 2, 1).
        group_size (`int`, *optional*, defaults to 64):
            Group-size value. Supported values are any value that is divisible by weight.shape[axis]).
        view_as_float (`bool`, *optional*, defaults to `False`):
            View the quantized weight as float (used in distributed training) if set to `True`.
        axis (`Optional[int]`, *optional*):
            Axis along which grouping is performed. Supported values are 0 or 1.
        dynamic_config (dict, *optional*):
            Parameters for dynamic configuration. The key is the name tag of the layer and the value is a quantization config.
            If set, each layer specified by its id will use its dedicated quantization configuration.
        skip_modules (`List[str]`, *optional*, defaults to `['lm_head']`):
            List of `nn.Linear` layers to skip.
        kwargs (`Dict[str, Any]`, *optional*):
            Additional parameters from which to initialize the configuration object.


def conv3d(self, input_tensor: core.Tensor) -> Mapping[str, core.Tensor]:
    out = nn_ops.conv3d(input_tensor, self.filters, strides=[1, 1, 2, 1, 1], dilations=[1, 1, 1, 1, 1], padding=padding, data_format='NDHWC')
    if has_bias:
        out = nn_ops.bias_add(out, self.bias)
    if activation_fn is not None:
        out = activation_fn(out)
    return {'output': out}

Performs a 3D convolution operation.

Args:
        input_tensor: Input tensor to perform convolution on.

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


def _SwitchRefOrTensor(data, pred, name='Switch'):
    data = ops.convert_to_tensor_or_composite(data, name='data')
    with ops.colocate_with(data, ignore_existing=True):
        if isinstance(data, tensor_lib.Tensor):
            if data.dtype._is_ref_dtype:
                return ref_switch(data, pred, name=name)
        return switch(data, pred, name=name)

Forwards `data` to an output determined by `pred`.

    If `pred` is false, the `data` input is forwarded to the first output.
    Otherwise, the data goes to the second output.

    This op handles `Tensor`s and `IndexedSlices`.

Args:
        data: The tensor to be forwarded to the appropriate output.
        pred: A scalar that specifies which output port will receive data.
        name: A name for this operation (optional).

Returns:
        `(output_false, output_true)`: If `pred` is true, data will be forwarded to
        `output_true`, otherwise it goes to `output_false`.

Raises:
        TypeError: if data is not a Tensor or IndexedSlices


def load_wav_file(filename):
    with tf.compat.v1.Session(graph=tf.Graph()) as sess:
        wav_filename_placeholder = tf.compat.v1.placeholder(tf.string, [])
        wav_loader = io_ops.read_file(wav_filename_placeholder)
        wav_decoder = tf.audio.decode_wav(wav_loader, desired_channels=1)
        return sess.run(wav_decoder, feed_dict={wav_filename_placeholder: filename}).audio.flatten()

Loads an audio file and returns a float PCM-encoded array of samples.

Args:
        filename: Path to the .wav file to load.

Returns:
        Numpy array holding the sample data as floats between -1.0 and 1.0.


def run_query(query: str) -> None:
    try:
        result = parse_query(query)
    except Exception as e:
        result = f'ERROR: {type(e).__name__}: {e.__str__()}.'
        return result
    return filter_records(convert_to_dataframe(result), query)

Run a query and display the result.

Args:
        query (str): The query to be executed.


def min(x, axis=None, keepdims=False, initial=None):
    if any_symbolic_tensors((x,)):
        return Min(axis=axis, keepdims=keepdims, initial=initial).symbolic_call(x)
    return backend.numpy.min(x, axis=axis, keepdims=keepdims, initial=initial)

Return the minimum of a tensor or minimum along an axis.

Args:
        x: Input tensor.
        axis: Axis or axes along which to operate. By default, flattened input
            is used.
        keepdims: If this is set to `True`, the axes which are reduced are left
            in the result as dimensions with size one. Defaults to `False`.
        initial: The maximum value of an output element. Defaults to `None`.

Returns:
        Minimum of `x`.


def __init__(self, config: GraniteMoeConfig):
    super(GraniteMoeMoE, self).__init__()
    self.input_size = config.hidden_size
    self.hidden_size = config.intermediate_size
    self.activation = ACT2FN[config.hidden_act]
    self.input_linear = GraniteMoeParallelExperts(config.num_local_experts, self.input_size, self.hidden_size * 2)
    self.output_linear = GraniteMoeParallelExperts(config.num_local_experts, self.hidden_size, self.input_size)
    self.router = GraniteMoeTopKGating(input_size=self.input_size, num_experts=config.num_local_experts, top_k=config.num_experts_per_tok)

A Sparsely gated mixture of experts layer with 1-layer Feed-Forward networks as experts.

Args:
        config:
            Configuration object with model hyperparameters.


def _input_array(num_dims):
    formatter = '{:0%db}' % num_dims
    strings = [formatter.format(i) for i in range(2 ** num_dims)]
    return np.array(strings, dtype='S%d' % num_dims).reshape([2] * num_dims)

Creates an ndarray where each element is the binary of its linear index.

Args:
        num_dims: The number of dimensions to create.

Returns:
        An ndarray of shape [2] * num_dims.


def _SanitizedMRO(obj):
    return_list = []
    for cls in tf_inspect.getmro(obj):
        if cls.__name__ == '_NewClass':
            continue
        str_repr = _NormalizeType(str(cls))
        return_list.append(str_repr)
        if 'tensorflow' not in str_repr and 'keras' not in str_repr:
            break
        if 'StubOutForTesting' in str_repr:
            break
    return return_list

Get a list of superclasses with minimal amount of non-TF classes.

    Based on many parameters like python version, OS, protobuf implementation
    or changes in google core libraries the list of superclasses of a class
    can change. We only return the first non-TF class to be robust to non API
    affecting changes. The Method Resolution Order returned by `tf_inspect.getmro`
    is still maintained in the return value.

Args:
        obj: A python routine for us the create the sanitized arspec of.

Returns:
        list of strings, string representation of the class names.


def _test_tensorflow_vs_numpy(self, x_np):
    y_np = self._total_variation_np(x_np)
    self._test(x_np, y_np)

Test the TensorFlow implementation against a numpy implementation.

Args:
        x_np: Numpy array with 3 or 4 dimensions.


def Colorize(self, string, color, justify=None):
    if justify:
        string = justify(string)
    if self._csi and color in self._ANSI_COLOR:
        return '{csi}{color_code}{string}{csi}{reset_code}'.format(csi=self._csi, color_code=self._ANSI_COLOR[color], reset_code=self._ANSI_COLOR_RESET, string=string)
    return string

Generates a colorized string, optionally justified.

Args:
        string: The string to write.
        color: The color name -- must be in _ANSI_COLOR.
        justify: The justification function, no justification if None. For
        example, justify=lambda s: s.center(10)

Returns:
        str, The colorized string that can be printed to the console.


def create_database_view(self, view: views.View, view_name: str) -> None:
    dataset = f'{self._view_dataset.project}.{self._view_dataset.dataset_id}'
    view_sql = f'CREATE OR REPLACE VIEW `{dataset}.{view_name}` AS\n{self.to_sql(view)}'
    self._client.query(view_sql).result()

Creates a BigQuery view with the given name in the runner's view_dataset.

Args:
        view: the FHIR view that creates
        view_name: the view name passed to the CREATE OR REPLACE VIEW statement.

Raises:
        google.cloud.exceptions.GoogleAPICallError if the job failed.


def sum(x, axis=None, keepdims=False):
    if any_symbolic_tensors((x,)):
        return Sum(axis=axis, keepdims=keepdims).symbolic_call(x)
    return backend.numpy.sum(x, axis=axis, keepdims=keepdims)

Sum of a tensor over the given axes.

Args:
        x: Input tensor.
        axis: Axis or axes along which the sum is computed. The default is to
            compute the sum of the flattened tensor.
        keepdims: If this is set to `True`, the axes which are reduced are left
            in the result as dimensions with size one.

Returns:
        Output tensor containing the sum.


def run_inference(self, batch: Sequence[numpy.ndarray], model: Union[xgboost.Booster, xgboost.XGBModel], inference_args: Optional[dict[str, Any]]=None) -> Iterable[PredictionResult]:
    return self._inference_fn(batch, model, inference_args)

Runs inferences on a batch of 2d numpy arrays.

Args:
        batch: A sequence of examples as 2d numpy arrays. Each
        row in an array is a single example. The dimensions
        must match the dimensions of the data used to train
        the model.
        model: XGBoost booster or XBGModel (sklearn interface). Must
        implement predict(X). Where the parameter X is a 2d numpy array.
        inference_args: Any additional arguments for an inference.

Returns:
        An Iterable of type PredictionResult.


def __init__(self, feature_extractor, tokenizer):
    super().__init__(feature_extractor, tokenizer)
    self.current_processor = self.feature_extractor
    self._in_target_context_manager = False

Constructs a Speech2Text processor which wraps a Speech2Text feature extractor and a Speech2Text tokenizer into a
    single processor.

    [`Speech2TextProcessor`] offers all the functionalities of [`Speech2TextFeatureExtractor`] and
    [`Speech2TextTokenizer`]. See the [`~Speech2TextProcessor.__call__`] and [`~Speech2TextProcessor.decode`] for more
    information.

Args:
        feature_extractor (`Speech2TextFeatureExtractor`):
            An instance of [`Speech2TextFeatureExtractor`]. The feature extractor is a required input.
        tokenizer (`Speech2TextTokenizer`):
            An instance of [`Speech2TextTokenizer`]. The tokenizer is a required input.


def forward(self, input_ids: Optional[torch.LongTensor]=None, pixel_values: Optional[torch.FloatTensor]=None, attention_mask: Optional[torch.Tensor]=None, position_ids: Optional[torch.LongTensor]=None, token_type_ids: Optional[torch.Tensor]=None, return_loss: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, interpolate_pos_encoding: bool=False, return_dict: Optional[bool]=None) -> Union[Tuple, AltCLIPOutput]:
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    text_outputs = self.text_model(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    vision_outputs = self.vision_model(pixel_values=pixel_values, output_attentions=output_attentions, output_hidden_states=output_hidden_states, interpolate_pos_encoding=interpolate_pos_encoding, return_dict=return_dict)
    image_embeds = vision_outputs[1]
    image_embeds = self.visual_projection(image_embeds)
    text_embeds = text_outputs[1]
    text_embeds = self.text_projection(text_embeds)
    image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True)
    text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True)
    logit_scale = self.logit_scale.exp()
    logits_per_text = torch.matmul(text_embeds, image_embeds.t()) * logit_scale
    logits_per_image = logits_per_text.T
    loss = None
    if return_loss:
        loss = clip_loss(logits_per_text)
    if not return_dict:
        output = (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)
        return (loss,) + output if loss is not None else output
    return AltCLIPOutput(loss=loss, logits_per_image=logits_per_image, logits_per_text=logits_per_text, text_embeds=text_embeds, image_embeds=image_embeds, text_model_output=text_outputs, vision_model_output=vision_outputs)

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

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

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


def get_size_with_aspect_ratio(image_size: Tuple[int, int], size: int, max_size: Optional[int]=None, mod_size: int=16) -> Tuple[int, int]:
    height, width = image_size
    raw_size = None
    if max_size is not None:
        min_original_size = float(min((height, width)))
        max_original_size = float(max((height, width)))
        if max_original_size / min_original_size * size > max_size:
            raw_size = max_size * min_original_size / max_original_size
            size = int(round(raw_size))
    if width < height:
        ow = size
        if max_size is not None and raw_size is not None:
            oh = int(raw_size * height / width)
        else:
            oh = int(size * height / width)
    elif height <= width and height == size or (width <= height and width == size):
        oh, ow = (height, width)
    else:
        oh = size
        if max_size is not None and raw_size is not None:
            ow = int(raw_size * width / height)
        else:
            ow = int(size * width / height)
    if mod_size is not None:
        ow_mod = np.mod(ow, mod_size)
        oh_mod = np.mod(oh, mod_size)
        ow = ow - ow_mod
        oh = oh - oh_mod
    return (oh, ow)

Computes the output image size given the input image size and the desired output size with multiple of divisible_size.

Args:
        image_size (`Tuple[int, int]`):
            The input image size.
        size (`int`):
            The desired output size.
        max_size (`int`, *optional*):
            The maximum allowed output size.
        mod_size (`int`, *optional*):
            The size to make multiple of mod_size.


def _call_unittest_assertion(assertion_method, *args, msg=None, extras=None, **kwargs):
    my_msg = None
    try:
        assertion_method(*args, **kwargs)
    except AssertionError as e:
        my_msg = str(e)
        if msg:
            my_msg = f'{my_msg} {msg}'
    if my_msg is not None:
        raise signals.TestFailure(my_msg, extras=extras)

Wrapper for converting a unittest assertion into a Mobly one.

Args:
        assertion_method: unittest.TestCase assertion method to call.
        *args: Positional arguments for the assertion call.
        msg: A string that adds additional info about the failure.
        extras: An optional field for extra information to be included in
        test result.
        **kwargs: Keyword arguments for the assertion call.


def get_temp_dir():
    return _googletest.GetTempDir()

Returns a temporary directory for use during tests.

    There is no need to delete the directory after the test.

    @compatibility(TF2)
    This function is removed in TF2. Please use `TestCase.get_temp_dir` instead
    in a test case.
    Outside of a unit test, obtain a temporary directory through Python's
    `tempfile` module.
    @end_compatibility

Returns:
        The temporary directory.


def __init__(self, missing_modules: Collection[str]=()):
    if os.getenv('TYPESHED_HOME'):
        self._store = ExternalTypeshedFs(missing_file=self.MISSING_FILE)
    else:
        self._store = InternalTypeshedFs(missing_file=self.MISSING_FILE)
    self._missing = self._load_missing().union(missing_modules)
    self._stdlib_versions = self._load_stdlib_versions()
    self._third_party_packages = self._load_third_party_packages()

Initializer.

Args:
        missing_modules: A collection of modules in the format
        'stdlib/module_name', which will be combined with the contents of
        MISSING_FILE to form a set of missing modules for which pytype will not
        report errors.


def initialize_multi_client_cluster(job_name: str, dtensor_jobs: List[str], client_id: int, collective_leader: str, port: Optional[int]=None, gpu_use_nccl_communication: bool=False, enable_coordination_service: bool=True):
    assert context.executing_eagerly()
    if not collective_leader.startswith('/job:'):
        collective_leader = '/job:' + collective_leader
    context.context().configure_collective_ops(use_nccl_communication=gpu_use_nccl_communication, collective_leader=collective_leader)
    if enable_coordination_service:
        context.context().configure_coordination_service(service_type='standalone', service_leader=collective_leader)
    config_proto = context.get_config()
    cluster_def = cluster_pb2.ClusterDef()
    cluster_def.job.add(name=job_name, tasks=dict(enumerate(dtensor_jobs)))
    server_def = tensorflow_server_pb2.ServerDef(cluster=cluster_def, default_session_config=config_proto, job_name=job_name, task_index=client_id, protocol=remote_utils.get_default_communication_protocol(), port=port)
    server_def.default_session_config.rpc_options.num_channels_per_target = 4
    server_def.default_session_config.experimental.recv_buf_max_chunk = -1
    logging.info('Enabling collectives with server_def: %s', server_def)
    context.context().enable_collective_ops(server_def)
    context.ensure_initialized()

Initialize GRPC servers and collectives for multi-client DTensor setup.

    This function can be used to initialize a multi-client cluster and enable
    collective ops. GRPC servers are necessary in the multi-client mode, even
    when the number of clientis is 1.

    NOTE: this function must be called in an eager context.

Args:
        job_name: The job name used by all clients in the DTensor cluster.
        dtensor_jobs: A list of the DTensor client jobs participating in the
        cluster. Must be strings of the form "hostname:port".
        client_id: The ID of the DTensor client this function is being called in.
        collective_leader: The job/task that will be used to run collectives.
        port: The port this client's GRPC server will run on. If omitted, use the
        port from dtensor_jobs for this client.
        gpu_use_nccl_communication: if True, configure TensorFlow to use NCCL by
        default.
        enable_coordination_service: If true, enable distributed coordination
        service to make sure that workers know the devices on each other, a
        prerequisite for data transfer through cross-worker rendezvous.

Raises:
        RuntimeError: If running inside a tf.function.


def __init__(self, mapping: PythonDict[Hashable, trace.TraceType], placeholder_type: Optional[Type[Any]]=None):
    self.mapping = mapping
    self._placeholder_type = placeholder_type

Represents a dictionary of TraceType objects.

Attributes:
        mapping: A mapping from keys to corresponding TraceTypes of the dict values.


def save_subset_weights_to_hdf5_group(f, weights):
    weight_values = [backend.convert_to_numpy(w) for w in weights]
    weight_names = [str(w.path).encode('utf8') for w in weights]
    save_attributes_to_hdf5_group(f, 'weight_names', weight_names)
    for name, val in zip(weight_names, weight_values):
        param_dset = f.create_dataset(name, val.shape, dtype=val.dtype)
        if not val.shape:
            param_dset[()] = val
        else:
            param_dset[:] = val

Save top-level weights of a model to a HDF5 group.

Args:
        f: HDF5 group.
        weights: List of weight variables.


def _join_test_cases_n_ary(self, test_cases: Iterator[tuple[Iterator[str], ...]], delimiters: ListOrTuple[str]) -> tuple[Iterator[str], ...]:
    output_stream_tuple, delimiter_stream_tuple = zip(*((io.StringIO(), io.StringIO(delimiter)) for delimiter in delimiters))
    for delimiter_stream in delimiter_stream_tuple:
        delimiter_stream.seek(0, os.SEEK_END)
    for test_case_stream_tuple in test_cases:
        for test_case_stream, delimiter_stream, output_stream in zip(test_case_stream_tuple, delimiter_stream_tuple, output_stream_tuple):
            output_stream.writelines(delimiter_stream)
            output_stream.writelines(test_case_stream)
            delimiter_stream.seek(0)
    for output_stream in output_stream_tuple:
        output_stream.seek(0)
    return output_stream_tuple

Joins each output[i] from all `test_cases`, delimited by `delimiters[i]`.

    Given an iterator to tuples of iterators, this function concatenates along
    the first "dimension" and returns a tuple of concatenated iterators. The
    concatenated iterators comprising `output_stream_tuple[i]` are delimited by
    `delimiters[i]`.

    For example, given the following input:
        `test_cases = [("a", "b"), ("c", "d")]`
        `delimiters = ("-", "_")`
    This function will return a tuple of two iterators:
        `[("a-", "b-"), ("c_", "d_")]`

Args:
        test_cases: An iterator over the test cases to join. Each test case is a
        tuple of N iterators yielding lines of text.
        delimiters: A list or tuple of N delimiter strings, one for each member of
        the returned tuple `output_stream_tuple`.

Returns:
        A tuple of N output-stream iterators. For each tuple-element index `i`,
        `output_stream_tuple[i]` is the concatenation of `test_cases[:][i]`,
        delimited by `delimiters[i]`.


def take(self, count, name=None) -> 'DatasetV2':
    from tensorflow.python.data.ops import take_op
    return take_op._take(self, count, name=name)

Creates a `Dataset` with at most `count` elements from this dataset.

    >>> dataset = tf.data.Dataset.range(10)
    >>> dataset = dataset.take(3)
    >>> [a.item() for a in dataset.as_numpy_iterator()]
    [0, 1, 2]

Args:
        count: A `tf.int64` scalar `tf.Tensor`, representing the number of
        elements of this dataset that should be taken to form the new dataset.
        If `count` is -1, or if `count` is greater than the size of this
        dataset, the new dataset will contain all elements of this dataset.
        name: (Optional.) A name for the tf.data operation.

Returns:
        A new `Dataset` with the transformation applied as described above.


def invoke_process_element(self, sdf_invoker, output_processor, element, restriction, watermark_estimator_state, *args, **kwargs):
    assert isinstance(sdf_invoker, DoFnInvoker)

    class CheckpointState(object):

        def __init__(self):
            self.checkpointed = None
            self.residual_restriction = None
    checkpoint_state = CheckpointState()

    def initiate_checkpoint():
        with self._checkpoint_lock:
            if checkpoint_state.checkpointed:
                return
            checkpoint_state.checkpointed = object()
        split = sdf_invoker.try_split(0)
        if split:
            _, checkpoint_state.residual_restriction = split
        else:
            checkpoint_state.checkpointed = None
    output_processor.reset()
    Timer(self._max_duration, initiate_checkpoint).start()
    sdf_invoker.invoke_process(element, additional_args=args, restriction=restriction, watermark_estimator_state=watermark_estimator_state)
    assert output_processor.output_iter is not None
    output_count = 0
    process_continuation = None
    for output in output_processor.output_iter:
        assert not process_continuation
        if isinstance(output, ProcessContinuation):
            initiate_checkpoint()
            process_continuation = output
            continue
        yield output
        output_count += 1
        if self._max_num_outputs and output_count >= self._max_num_outputs:
            initiate_checkpoint()
    result = SDFProcessElementInvoker.Result(residual_restriction=checkpoint_state.residual_restriction) if checkpoint_state.residual_restriction else SDFProcessElementInvoker.Result()
    yield result

Invokes `process()` method of a Splittable `DoFn` for a given element.

Args:
        sdf_invoker: a `DoFnInvoker` for the Splittable `DoFn`.
        element: the element to process

Returns:
        a `SDFProcessElementInvoker.Result` object.


def _call_wrapped_cell(self, inputs, state, cell_call_fn, **kwargs):
    raise NotImplementedError

Calls the wrapped cell and performs the wrapping logic.

    This method is called from the wrapper's `call` or `__call__` methods.

Args:
        inputs: A tensor with wrapped cell's input.
        state: A tensor or tuple of tensors with wrapped cell's state.
        cell_call_fn: Wrapped cell's method to use for step computation (cell's
        `__call__` or 'call' method).
        **kwargs: Additional arguments.

Returns:
        A pair containing:
        - Output: A tensor with cell's output.
        - New state: A tensor or tuple of tensors with new wrapped cell's state.


def _generate_schedule(self, cpn_frequency, roll_convention):
    if self._first_coupon_date is None and self._penultimate_coupon_date is None:
        cpn_dates = dates.PeriodicSchedule(start_date=self._start_date, end_date=self._end_date, tenor=cpn_frequency, roll_convention=roll_convention).dates()
        is_regular_cpn = tf.constant(True, dtype=bool, shape=cpn_dates[:, :-1].shape)
    elif self._first_coupon_date is not None:
        cpn_dates = dates.PeriodicSchedule(start_date=self._first_coupon_date, end_date=self._end_date, tenor=cpn_frequency, roll_convention=roll_convention).dates()
        cpn_dates = dates.DateTensor.concat([self._start_date.expand_dims(-1), cpn_dates], axis=1)
        is_irregular_cpn = tf.constant(False, dtype=bool, shape=self._start_date.shape)
        is_regular_cpn = tf.concat([tf.expand_dims(is_irregular_cpn, axis=-1), tf.constant(True, dtype=bool, shape=cpn_dates[:, :-2].shape)], axis=1)
    else:
        cpn_dates = dates.PeriodicSchedule(start_date=self._start_date, end_date=self._penultimate_coupon_date, backward=True, tenor=cpn_frequency, roll_convention=roll_convention).dates()
        cpn_dates = dates.DateTensor.concat([cpn_dates, self._end_date.expand_dims(-1)], axis=1)
        is_irregular_cpn = tf.constant(False, dtype=bool, shape=self._end_date.shape)
        is_regular_cpn = tf.concat([tf.constant(True, dtype=bool, shape=cpn_dates[:, :-2].shape), tf.expand_dims(is_irregular_cpn, axis=-1)], axis=1)
    return (cpn_dates, is_regular_cpn)

Method to generate coupon dates.

Args:
        cpn_frequency: A `PeriodTensor` specifying the frequency of coupon
        payments.
        roll_convention: Scalar of type `BusinessDayConvention` specifying how
        dates are rolled if they fall on holidays.

Returns:
        A tuple containing the generated date schedule and a boolean `Tensor`
        of the same shape as the schedule specifying whether the coupons are
        regular coupons.


def breakpoints(self):
    return self._breakpoints

Get a set of the currently-activated breakpoints.

Returns:
        A `set` of 3-tuples: (node_name, output_slot, debug_op), e.g.,
        {("MatMul", 0, "DebugIdentity")}.


def _merge_with(self, other: 'DynamicRaggedShape.Spec') -> 'DynamicRaggedShape.Spec':
    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_rp = [a._merge_with(b) for a, b in zip(a._row_partitions, b._row_partitions)]
    new_static_inner_shape = a._static_inner_shape.merge_with(b._static_inner_shape)
    dtype = b.dtype if a.dtype == dtypes.int32 else dtypes.int64
    return DynamicRaggedShape.Spec(new_rp, new_static_inner_shape, dtype=dtype)

Merges all information between two specs.

    Specs are expected to represent the same information modulo
    num_row_partitons.

    If the specs are of different ranks, then fail.

Args:
        other: another Spec of the same rank.

Returns:
        a Spec with the union of information.


def debug_string(self, with_typing: bool=False, indent: int=0) -> str:
    operand_name = f'{self} '
    operand_prints = ''.join(('\n' + op.debug_string(with_typing, indent + 1) for op in self.operands))
    type_print = f' type={self.return_type}' if with_typing else ''
    return f'{'| ' * indent}+ {operand_name}<{self.__class__.__name__}{type_print}> ({operand_prints})'

Builds a string describing the expression tree starting from this node.

Args:
        with_typing: If true, includes the type each node evaluates to.
        indent: The initial number of spaces to use as indentation for the debug
        string.

Returns:
        A string which recursively describes this node and its operands.


def FromTimestampToHttp(self, ts):
    ts = time.gmtime(ts)
    return time.strftime('%a, %d %b %Y %H:%M:%S GMT', ts)

Converts internal nss_cache timestamp to HTTP timestamp.

Args:
        ts: number of seconds since epoch

Returns:
        HTTP format timestamp string


def _get_trainable_state(self):
    layers = self._flatten_layers(include_self=False, recursive=False)
    trainable_state = {self: self.trainable}
    for l in layers:
        trainable_state.update(l._get_trainable_state())
    return trainable_state

Get the `trainable` state of each sublayer.

Returns:
        A dict mapping all sublayers to their `trainable` value.


def _unique_parameters(self) -> 'list[cfg.Variable]':
    return []

Get unique parameter subtypes as variables.

    This will retrieve 'children' of this value that contribute to the
    type of it. So it will retrieve type parameters, but not attributes. To
    keep the number of possible combinations reasonable, when we encounter
    multiple instances of the same type, we include only one.

Returns:
        A list of variables.


def call(self, hidden_states: tf.Tensor, attention_mask: tf.Tensor | None=None, encoder_hidden_states: tf.Tensor | None=None, encoder_attention_mask: tf.Tensor | None=None, layer_head_mask: tf.Tensor | None=None, cross_attn_layer_head_mask: tf.Tensor | None=None, past_key_value: Tuple[tf.Tensor] | None=None, training: Optional[bool]=False) -> Tuple[tf.Tensor, tf.Tensor, Tuple[Tuple[tf.Tensor]]]:
    residual = hidden_states
    hidden_states = self.self_attn_layer_norm(hidden_states)
    self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
    hidden_states, self_attn_weights, present_key_value = self.self_attn(hidden_states=hidden_states, past_key_value=self_attn_past_key_value, attention_mask=attention_mask, layer_head_mask=layer_head_mask)
    hidden_states = self.dropout(hidden_states, training=training)
    hidden_states = residual + hidden_states
    cross_attn_present_key_value = None
    cross_attn_weights = None
    if encoder_hidden_states is not None:
        residual = hidden_states
        hidden_states = self.encoder_attn_layer_norm(hidden_states)
        cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
        hidden_states, cross_attn_weights, cross_attn_present_key_value = self.encoder_attn(hidden_states=hidden_states, key_value_states=encoder_hidden_states, attention_mask=encoder_attention_mask, layer_head_mask=cross_attn_layer_head_mask, past_key_value=cross_attn_past_key_value)
        hidden_states = self.dropout(hidden_states, training=training)
        hidden_states = residual + hidden_states
        present_key_value = present_key_value + cross_attn_present_key_value
    residual = hidden_states
    hidden_states = self.final_layer_norm(hidden_states)
    hidden_states = self.activation_fn(self.fc1(hidden_states))
    hidden_states = self.activation_dropout(hidden_states, training=training)
    hidden_states = self.fc2(hidden_states)
    hidden_states = self.dropout(hidden_states, training=training)
    hidden_states = residual + hidden_states
    return (hidden_states, self_attn_weights, cross_attn_weights, present_key_value)

Args:
    hidden_states (`tf.Tensor`): input to the layer of shape *(batch, seq_len, embed_dim)*
    attention_mask (`tf.Tensor`): attention mask of size
        *(batch, 1, tgt_len, src_len)* where padding elements are indicated by very large negative values.
    encoder_hidden_states (`tf.Tensor`):
        cross attention input to the layer of shape *(batch, seq_len, embed_dim)*
    encoder_attention_mask (`tf.Tensor`): encoder attention mask of size
        *(batch, 1, tgt_len, src_len)* where padding elements are indicated by very large negative values.
    layer_head_mask (`tf.Tensor`): mask for attention heads in a given layer of size
        *(decoder_attention_heads,)*
    cross_attn_layer_head_mask (`tf.Tensor`): mask for heads of the cross-attention module.
        *(decoder_attention_heads,)*
    past_key_value (`Tuple(tf.Tensor)`): cached past key and value projection states


def __init__(self, name='mean', dtype=None):
    super(Mean, self).__init__(reduction=metrics_utils.Reduction.WEIGHTED_MEAN, name=name, dtype=dtype)

Computes the (weighted) mean of the given values.

    For example, if values is [1, 3, 5, 7] then the mean is 4.
    If the weights were specified as [1, 1, 0, 0] then the mean would be 2.

    This metric creates two variables, `total` and `count` that are used to
    compute the average of `values`. This average is ultimately returned as `mean`
    which is an idempotent operation that simply divides `total` by `count`.

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

Args:
        name: (Optional) string name of the metric instance.
        dtype: (Optional) data type of the metric result.

    Standalone usage:

    >>> m = tf.keras.metrics.Mean()
    >>> m.update_state([1, 3, 5, 7])
    >>> m.result().numpy()
    4.0
    >>> m.reset_state()
    >>> m.update_state([1, 3, 5, 7], sample_weight=[1, 1, 0, 0])
    >>> m.result().numpy()
    2.0

    Usage with `compile()` API:

    ```python
    model.add_metric(tf.keras.metrics.Mean(name='mean_1')(outputs))
    model.compile(optimizer='sgd', loss='mse')
    ```


def __init__(self, quantile_tracker: Optional[QuantileTracker]=None):
    self._quantile_tracker = quantile_tracker or BufferedSlidingQuantileTracker(DEFAULT_WINDOW_SIZE, 0.5)
    assert self._quantile_tracker._q == 0.5, 'quantile_tracker must be initialized with q = 0.5'

Tracks the median of a stream of values using a quantile tracker.

    This wrapper class encapsulates a `QuantileTracker` configured specifically
    for the 0.5 quantile (median).

Args:
        quantile_tracker: An optional `QuantileTracker` instance. If not provided,
        a `BufferedSlidingQuantileTracker` with a default window size 1000 and
        q=0.5 is created.

Raises:
        AssertionError: If the provided quantile_tracker is not initialized with
        q=0.5.


def _check_trt_version_compatibility():
    if not _pywrap_py_utils.is_tensorrt_enabled():
        logging.error('Tensorflow needs to be built with TensorRT support enabled to allow TF-TRT to operate.')
        raise RuntimeError('Tensorflow has not been built with TensorRT support.')
    if platform.system() == 'Windows':
        logging.warn('Windows support is provided experimentally. No guarantee is made regarding functionality or engineering support. Use at your own risk.')
    linked_version = _pywrap_py_utils.get_linked_tensorrt_version()
    loaded_version = _pywrap_py_utils.get_loaded_tensorrt_version()
    logging.info('Linked TensorRT version: %s', str(linked_version))
    logging.info('Loaded TensorRT version: %s', str(loaded_version))

    def raise_trt_version_deprecated(version_type, trt_version):
        assert version_type in ['linked', 'loaded'], "Incorrect value received for version_type: %s. Accepted: ['linked', 'loaded']" % version_type
        logging.error('The {version_type} version of TensorRT: `{trt_version}` has now been removed. Please upgrade to TensorRT 7 or more recent.'.format(version_type=version_type, trt_version=trt_utils.version_tuple_to_string(trt_version)))
        raise RuntimeError('Incompatible %s TensorRT versions' % version_type)
    if not trt_utils.is_linked_tensorrt_version_greater_equal(7, 0, 0):
        raise_trt_version_deprecated('linked', linked_version)
    if not trt_utils.is_loaded_tensorrt_version_greater_equal(7, 0, 0):
        raise_trt_version_deprecated('loaded', loaded_version)
    if loaded_version[0] != linked_version[0] or not trt_utils.is_loaded_tensorrt_version_greater_equal(*linked_version):
        logging.error('Loaded TensorRT %s but linked TensorFlow against TensorRT %s. A few requirements must be met:\n\t-It is required to use the same major version of TensorRT during compilation and runtime.\n\t-TensorRT does not support forward compatibility. The loaded version has to be equal or more recent than the linked version.', trt_utils.version_tuple_to_string(loaded_version), trt_utils.version_tuple_to_string(linked_version))
        raise RuntimeError('Incompatible TensorRT major version')
    elif loaded_version != linked_version:
        logging.info('Loaded TensorRT %s and linked TensorFlow against TensorRT %s. This is supported because TensorRT minor/patch upgrades are backward compatible.', trt_utils.version_tuple_to_string(loaded_version), trt_utils.version_tuple_to_string(linked_version))

Check compatibility of TensorRT version.

Raises:
        RuntimeError: if the TensorRT library version is incompatible.


def default_session(session):
    return _default_session_stack.get_controller(session)

Python "with" handler for defining a default session.

    This function provides a means of registering a session for handling
    Tensor.eval() and Operation.run() calls. It is primarily intended for use
    by session.Session, but can be used with any object that implements
    the Session.run() interface.

    Use with the "with" keyword to specify that Tensor.eval() and Operation.run()
    invocations within the scope of a block should be executed by a particular
    session.

    The default session applies to the current thread only, so it is always
    possible to inspect the call stack and determine the scope of a default
    session. If you create a new thread, and wish to use the default session
    in that thread, you must explicitly add a "with ops.default_session(sess):"
    block in that thread's function.

Example:
        The following code examples are equivalent:

        # 1. Using the Session object directly:
        sess = ...
        c = tf.constant(5.0)
        sess.run(c)

        # 2. Using default_session():
        sess = ...
        with ops.default_session(sess):
        c = tf.constant(5.0)
        result = c.eval()

        # 3. Overriding default_session():
        sess = ...
        with ops.default_session(sess):
        c = tf.constant(5.0)
        with ops.default_session(...):
        c.eval(session=sess)

Args:
        session: The session to be installed as the default session.

Returns:
        A context manager for the default session.


def repeat(x, repeats, axis=None):
    if any_symbolic_tensors((x,)):
        return Repeat(repeats, axis=axis).symbolic_call(x)
    return backend.numpy.repeat(x, repeats, axis=axis)

Repeat each element of a tensor after themselves.

Args:
        x: Input tensor.
        repeats: The number of repetitions for each element.
        axis: The axis along which to repeat values. By default, use
            the flattened input array, and return a flat output array.

Returns:
        Output tensor.


def _get_truncated_table_rows(self, query_tokens: List[str], tokenized_table: TokenizedTable, num_rows: int, num_columns: int, max_length: int, truncation_strategy: Union[str, TapasTruncationStrategy]) -> Tuple[int, int]:
    if not isinstance(truncation_strategy, TapasTruncationStrategy):
        truncation_strategy = TapasTruncationStrategy(truncation_strategy)
    if max_length is None:
        max_length = self.model_max_length
    if truncation_strategy == TapasTruncationStrategy.DROP_ROWS_TO_FIT:
        while True:
            num_tokens = self._get_max_num_tokens(query_tokens, tokenized_table, num_rows=num_rows, num_columns=num_columns, max_length=max_length)
            if num_tokens is not None:
                break
            num_rows -= 1
            if num_rows < 1:
                break
    elif truncation_strategy != TapasTruncationStrategy.DO_NOT_TRUNCATE:
        raise ValueError(f'Unknown truncation strategy {truncation_strategy}.')
    return (num_rows, num_tokens or 1)

Truncates a sequence pair in-place following the strategy.

Args:
        query_tokens (`List[str]`):
            List of strings corresponding to the tokenized query.
        tokenized_table (`TokenizedTable`):
            Tokenized table
        num_rows (`int`):
            Total number of table rows
        num_columns (`int`):
            Total number of table columns
        max_length (`int`):
            Total maximum length.
        truncation_strategy (`str` or [`TapasTruncationStrategy]`):
            Truncation strategy to use. Seeing as this method should only be called when truncating, the only
            available strategy is the `"drop_rows_to_fit"` strategy.

Returns:
        `Tuple(int, int)`: tuple containing the number of rows after truncation, and the number of tokens available
        for each table element.


def squared_hinge(y_true, y_pred):
    y_pred = ops.convert_to_tensor(y_pred)
    y_true = ops.cast(y_true, y_pred.dtype)
    y_true = convert_binary_labels_to_hinge(y_true)
    return ops.mean(ops.square(ops.maximum(1.0 - y_true * y_pred, 0.0)), axis=-1)

Computes the squared hinge loss between `y_true` & `y_pred`.

    Formula:

    ```python
    loss = mean(square(maximum(1 - y_true * y_pred, 0)), axis=-1)
    ```

Args:
        y_true: The ground truth values. `y_true` values are expected to be -1
            or 1. If binary (0 or 1) labels are provided we will convert them
            to -1 or 1 with shape = `[batch_size, d0, .. dN]`.
        y_pred: The predicted values with shape = `[batch_size, d0, .. dN]`.

Returns:
        Squared hinge loss values with shape = `[batch_size, d0, .. dN-1]`.

Example:
    >>> y_true = np.random.choice([-1, 1], size=(2, 3))
    >>> y_pred = np.random.random(size=(2, 3))
    >>> loss = keras.losses.squared_hinge(y_true, y_pred)


def learn_one(self, x: beam.Row) -> None:
    if len(x.__dict__) != 1:
        raise ValueError('RobustZScore.learn_one expected univariate input, but got %s', str(x))
    v = next(iter(x))
    self._mad_tracker.push(v)

Updates the `MadTracker` with a new data point.

Args:
        x: A `beam.Row` containing a single numerical value.


def _get_default_initializer(self, name, shape=None, dtype=dtypes.float32):
    del shape
    if dtype.is_floating:
        initializer = init_ops.glorot_uniform_initializer()
        initializing_from_value = False
    elif dtype.is_integer or dtype.is_unsigned or dtype.is_bool or (dtype == dtypes.string):
        initializer = init_ops.zeros_initializer()
        initializing_from_value = False
    else:
        raise ValueError('An initializer for variable %s of %s is required' % (name, dtype.base_dtype))
    return (initializer, initializing_from_value)

Provide a default initializer and a corresponding value.

Args:
        name: see get_variable.
        shape: see get_variable.
        dtype: see get_variable.

Returns:
        initializer and initializing_from_value. See get_variable above.

Raises:
        ValueError: When giving unsupported dtype.


def test_fail(self, e=None):
    self._test_end(TestResultEnums.TEST_RESULT_FAIL, e)

To mark the test as failed in this record.

    Only test_fail does instance check because we want 'assert xxx' to also
    fail the test same way assert_true does.

Args:
        e: An exception object. It can be an instance of AssertionError or
        mobly.base_test.TestFailure.


def set_save_handler(save_handler: Optional[Callable[..., Any]]) -> Optional[Callable[..., Any]]:
    if save_handler and (not callable(save_handler)):
        raise ValueError('`save_handler` must be callable.')
    global _SAVE_HANDLER
    old_handler = _SAVE_HANDLER
    _SAVE_HANDLER = save_handler
    return old_handler

Sets global save handler.

Args:
        save_handler: A callable object that takes at least one argument as value to
        save. `symbolic.save` method will pass through all arguments to this
        handler and return its return value.

Returns:
        Previous global save handler.


def call(self, pixel_values: tf.Tensor, labels: tf.Tensor | None=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[Tuple, TFSemanticSegmenterOutput]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    if labels is not None and (not self.config.num_labels > 1):
        raise ValueError('The number of labels should be greater than one')
    outputs = self.segformer(pixel_values, output_attentions=output_attentions, output_hidden_states=True, return_dict=return_dict)
    encoder_hidden_states = outputs.hidden_states if return_dict else outputs[1]
    logits = self.decode_head(encoder_hidden_states)
    loss = None
    if labels is not None:
        loss = self.hf_compute_loss(logits=logits, labels=labels)
    logits = tf.transpose(logits, perm=[0, 3, 1, 2])
    if not return_dict:
        if output_hidden_states:
            output = (logits,) + outputs[1:]
        else:
            output = (logits,) + outputs[2:]
        return (loss,) + output if loss is not None else output
    return TFSemanticSegmenterOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states if output_hidden_states else None, attentions=outputs.attentions)

labels (`tf.Tensor` of shape `(batch_size, height, width)`, *optional*):
    Ground truth semantic segmentation maps for computing the loss. Indices should be in `[0, ...,
    config.num_labels - 1]`. If `config.num_labels > 1`, a (per-pixel) classification loss is computed
    (Cross-Entropy).

Returns:

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

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

    >>> image_processor = AutoImageProcessor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
    >>> model = TFSegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")

    >>> inputs = image_processor(images=image, return_tensors="tf")
    >>> outputs = model(**inputs, training=False)
    >>> # logits are of shape (batch_size, num_labels, height/4, width/4)
    >>> logits = outputs.logits
    >>> list(logits.shape)
    [1, 150, 128, 128]
    ```


def __neg__(self: EventSetOrNode) -> EventSetOrNode:
    from temporian.core.operators.scalar import multiply_scalar
    return multiply_scalar(input=self, value=-1)

Negates an [`EventSet`][temporian.EventSet] element-wise.

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

        ```

Returns:
        Negated EventSet.


def create_token_type_ids_from_sequences(self, token_ids_0: List[int], token_ids_1: Optional[List[int]]=None) -> List[int]:
    sep = [self.sep_token_id]
    cls = [self.cls_token_id]
    if token_ids_1 is None:
        return len(cls + token_ids_0 + sep) * [0]
    return len(cls + token_ids_0 + sep + sep + token_ids_1 + sep) * [0]

Create a mask from the two sequences passed to be used in a sequence-pair classification task. Longformer does not
    make use of token type ids, therefore a list of zeros is returned.

Args:
        token_ids_0 (`List[int]`):
            List of IDs.
        token_ids_1 (`List[int]`, *optional*):
            Optional second list of IDs for sequence pairs.

Returns:
        `List[int]`: List of zeros.


def forward(self, hidden_states: torch.Tensor):
    if hidden_states.size(-1) != self.dim_norm:
        raise AssertionError('hidden_states.size(-1) != self.dim_norm')
    old_dtype = hidden_states.dtype
    variance = hidden_states.to(torch.float32).pow(2).mean(dim=-1, keepdim=True)
    hidden_states = (hidden_states * torch.rsqrt(variance + self.eps)).to(old_dtype) * self.weight
    return hidden_states

Args:
    hidden_states (`torch.Tensor` of shape `(batch, seq_len, dim_in)`)


def _gelu_new(x):
    x = tf.convert_to_tensor(x)
    pi = tf.cast(math.pi, x.dtype)
    coeff = tf.cast(0.044715, x.dtype)
    cdf = 0.5 * (1.0 + tf.tanh(tf.sqrt(2.0 / pi) * (x + coeff * tf.pow(x, 3))))
    return x * cdf

Gaussian Error Linear Unit. This is a smoother version of the GELU. Original paper: https://huggingface.co/papers/1606.0841

Args:
        x: float Tensor to perform activation

Returns:
        `x` with the GELU activation applied.


def forward(self, masks_queries_logits: torch.Tensor, class_queries_logits: torch.Tensor, mask_labels: torch.Tensor, class_labels: torch.Tensor) -> List[Tuple[Tensor]]:
    indices: List[Tuple[np.array]] = []
    batch_size = masks_queries_logits.shape[0]
    for i in range(batch_size):
        pred_probs = class_queries_logits[i].softmax(-1)
        pred_mask = masks_queries_logits[i]
        cost_class = -pred_probs[:, class_labels[i]]
        target_mask = mask_labels[i].to(pred_mask)
        target_mask = target_mask[:, None]
        pred_mask = pred_mask[:, None]
        point_coordinates = torch.rand(1, self.num_points, 2, device=pred_mask.device)
        target_coordinates = point_coordinates.repeat(target_mask.shape[0], 1, 1)
        target_mask = sample_point(target_mask, target_coordinates, align_corners=False).squeeze(1)
        pred_coordinates = point_coordinates.repeat(pred_mask.shape[0], 1, 1)
        pred_mask = sample_point(pred_mask, pred_coordinates, align_corners=False).squeeze(1)
        cost_mask = pair_wise_sigmoid_cross_entropy_loss(pred_mask, target_mask)
        cost_dice = pair_wise_dice_loss(pred_mask, target_mask)
        cost_matrix = self.cost_mask * cost_mask + self.cost_class * cost_class + self.cost_dice * cost_dice
        cost_matrix = torch.minimum(cost_matrix, torch.tensor(10000000000.0))
        cost_matrix = torch.maximum(cost_matrix, torch.tensor(-10000000000.0))
        cost_matrix = torch.nan_to_num(cost_matrix, 0)
        assigned_indices: Tuple[np.array] = linear_sum_assignment(cost_matrix.cpu())
        indices.append(assigned_indices)
    matched_indices = [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]
    return matched_indices

Params:
    masks_queries_logits (`torch.Tensor`):
        A tensor of dim `batch_size, num_queries, num_labels` with the classification logits.
    class_queries_logits (`torch.Tensor`):
        A tensor of dim `batch_size, num_queries, height, width` with the predicted masks.
    class_labels (`torch.Tensor`):
        A tensor of dim `num_target_boxes` (where num_target_boxes is the number of ground-truth objects in the
        target) containing the class labels.
    mask_labels (`torch.Tensor`):
        A tensor of dim `num_target_boxes, height, width` containing the target masks.

Returns:
    matched_indices (`List[Tuple[Tensor]]`): A list of size batch_size, containing tuples of (index_i, index_j)
    where:
        - index_i is the indices of the selected predictions (in order)
        - index_j is the indices of the corresponding selected labels (in order)
    For each batch element, it holds:
        len(index_i) = len(index_j) = min(num_queries, num_target_boxes).


def __init__(self, write_config: Dict[str, Any], *, schema_config: Optional[SchemaConfig]=None):
    if 'table' not in write_config:
        raise ValueError("write_config must be provided with 'table' specified")
    self.write_config = write_config
    self.schema_config = schema_config

Configuration for writing vectors to BigQuery using managed transforms.

    Supports both default schema (id, embedding, content, metadata columns) and
    custom schemas through SchemaConfig.

    Example with default schema:
        >>> config = BigQueryVectorWriterConfig(
        ...     write_config={'table': 'project.dataset.embeddings'})

    Example with custom schema:
        >>> schema_config = SchemaConfig(
        ...   schema={
        ...     'fields': [
        ...       {'name': 'id', 'type': 'STRING'},
        ...       {'name': 'embedding', 'type': 'FLOAT64', 'mode': 'REPEATED'},
        ...       {'name': 'source_url', 'type': 'STRING'}
        ...     ]
        ...   },
        ...   chunk_to_dict_fn=lambda chunk: {
        ...       'id': chunk.id,
        ...       'embedding': chunk.embedding.dense_embedding,
        ...       'source_url': chunk.metadata.get('url')
        ...   }
        ... )
        >>> config = BigQueryVectorWriterConfig(
        ...   write_config={'table': 'project.dataset.embeddings'},
        ...   schema_config=schema_config
        ... )

Args:
        write_config: BigQuery write configuration dict. Must include 'table'.
            Other options like create_disposition, write_disposition can be
            specified.
        schema_config: Optional configuration for custom schema and row
            conversion.
            If not provided, uses default schema with id, embedding, content and
            metadata columns.

Raises:
        ValueError: If write_config doesn't include table specification.


def make_pool_tests(pool_op_in, allow_fully_quantize=False):
    pool_op = pool_op_in

    def f(options, expected_tf_failures=0):
        """Actual function that generates examples.

    Args:
      options: An Options instance.
      expected_tf_failures: number of expected tensorflow failures.
    """
        test_parameters = [{'ksize': [[2, 1, 1, 2], [1, 1, 1, 1], [1, 1, 2, 1], [1, 10, 11, 1]], 'strides': [[2, 1, 1, 2], [1, 1, 1, 1], [1, 1, 2, 1], [1, 10, 11, 1]], 'input_shape': [[], [1, 1, 1, 1], [1, 15, 14, 1], [3, 15, 14, 3]], 'padding': ['SAME', 'VALID'], 'data_format': ['NHWC'], 'fully_quantize': [False], 'quant_16x8': [False]}, {'ksize': [[2, 1, 1, 2], [1, 1, 1, 1], [1, 1, 2, 1], [1, 10, 11, 1]], 'strides': [[2, 1, 1, 2], [1, 1, 1, 1], [1, 1, 2, 1], [1, 10, 11, 1]], 'input_shape': [[], [1, 1, 1, 1], [1, 15, 14, 1], [3, 15, 14, 3]], 'padding': ['SAME', 'VALID'], 'data_format': ['NHWC'], 'fully_quantize': [True], 'quant_16x8': [False]}, {'ksize': [[1, 1, 1, 1]], 'strides': [[1, 1, 1, 1]], 'input_shape': [[1, 1, 1, 1]], 'padding': ['SAME', 'VALID'], 'data_format': ['NHWC'], 'fully_quantize': [True], 'quant_16x8': [True]}]
        if not allow_fully_quantize:
            test_parameters = [test_parameter for test_parameter in test_parameters if True not in test_parameter['fully_quantize']]

        def build_graph(parameters):
            input_tensor = tf.compat.v1.placeholder(dtype=tf.float32, name='input', shape=parameters['input_shape'])
            out = pool_op(input_tensor, ksize=parameters['ksize'], strides=parameters['strides'], data_format=parameters['data_format'], padding=parameters['padding'])
            return ([input_tensor], [out])

        def build_inputs(parameters, sess, inputs, outputs):
            if allow_fully_quantize:
                input_values = create_tensor_data(tf.float32, parameters['input_shape'], min_value=-1, max_value=1)
            else:
                input_values = create_tensor_data(tf.float32, parameters['input_shape'])
            return ([input_values], sess.run(outputs, feed_dict=dict(zip(inputs, [input_values]))))
        make_zip_of_tests(options, test_parameters, build_graph, build_inputs, expected_tf_failures=expected_tf_failures)
    return f

Make a set of tests to do average pooling.

Args:
        pool_op_in: TensorFlow pooling operation to test  i.e. `tf.nn.avg_pool2d`.
        allow_fully_quantize: bool, whether fully_quantize is allowed.

Returns:
        A function representing the true generator (after curried pool_op_in).


def _GetMaps(name, commands, default_options):
    global_options = copy.copy(default_options)
    options_map = collections.defaultdict(lambda: copy.copy(default_options))
    subcommands_map = collections.defaultdict(set)
    for command in commands:
        if len(command) == 1:
            if _IsOption(command[0]):
                global_options.add(command[0])
            else:
                subcommands_map[name].add(command[0])
        elif command:
            subcommand = command[-2]
            arg = _FormatForCommand(command[-1])
            if _IsOption(arg):
                args_map = options_map
            else:
                args_map = subcommands_map
            args_map[subcommand].add(arg)
            args_map[subcommand.replace('_', '-')].add(arg)
    return (global_options, options_map, subcommands_map)

Returns sets of subcommands and options for each command.

Args:
        name: The first token in the commands, also the name of the command.
        commands: A list of all possible commands that tab completion can complete
        to. Each command is a list or tuple of the string tokens that make up
        that command.
        default_options: A dict of options that can be used with any command. Use
        this if there are flags that can always be appended to a command.

Returns:
        global_options: A set of all options of the first token of the command.
        subcommands_map: A dict storing set of subcommands for each
        command/subcommand.
        options_map: A dict storing set of options for each subcommand.


def convert_variable_to_constant(self, incoming_edge, tensor_data):
    index = incoming_edge.destination.index
    for edge in self.outgoing_edges:
        if edge.source.index == index:
            edge.destination.convertible.convert_variable_to_constant(edge, tensor_data)
    function = self.converted_self().function
    function.signature.input_arg[index].type = tensor_data.dtype
    if '_input_shapes' in function.attr:
        function.attr['_input_shapes'].list.shape[index].unknown_rank = True
        del function.attr['_input_shapes'].list.shape[index].dim[:]
    arg_attrs = function.arg_attr[index].attr
    if '_output_shapes' in arg_attrs:
        arg_attrs['_output_shapes'].list.shape[0].unknown_rank = True
        del arg_attrs['_output_shapes'].list.shape[0].dim[:]

Converts one function argument into a constant.

Args:
        incoming_edge: The edge into the argument to be converted.
        tensor_data: The constant value.