def byte_swap_tflite_model_obj(model, from_endiness, to_endiness):
    if model is None:
        return
    buffer_swapped = []
    types_of_16_bits = [schema_fb.TensorType.FLOAT16, schema_fb.TensorType.INT16, schema_fb.TensorType.UINT16]
    types_of_32_bits = [schema_fb.TensorType.FLOAT32, schema_fb.TensorType.INT32, schema_fb.TensorType.COMPLEX64, schema_fb.TensorType.UINT32]
    types_of_64_bits = [schema_fb.TensorType.INT64, schema_fb.TensorType.FLOAT64, schema_fb.TensorType.COMPLEX128, schema_fb.TensorType.UINT64]
    for subgraph in model.subgraphs:
        for tensor in subgraph.tensors:
            if tensor.buffer > 0 and tensor.buffer < len(model.buffers) and (tensor.buffer not in buffer_swapped) and (model.buffers[tensor.buffer].data is not None):
                if tensor.type == schema_fb.TensorType.STRING:
                    byte_swap_string_content(model.buffers[tensor.buffer], from_endiness, to_endiness)
                elif tensor.type in types_of_16_bits:
                    byte_swap_buffer_content(model.buffers[tensor.buffer], 2, from_endiness, to_endiness)
                elif tensor.type in types_of_32_bits:
                    byte_swap_buffer_content(model.buffers[tensor.buffer], 4, from_endiness, to_endiness)
                elif tensor.type in types_of_64_bits:
                    byte_swap_buffer_content(model.buffers[tensor.buffer], 8, from_endiness, to_endiness)
                else:
                    continue
                buffer_swapped.append(tensor.buffer)

Byte swaps the buffers field in a TFLite model.

Args:
        model: TFLite model object of from_endiness format.
        from_endiness: The original endianness format of the buffers in model.
        to_endiness: The destined endianness format of the buffers in model.


def broadcast(tensor):
    _check_device(tensor)
    with ops.device(tensor.device):
        return gen_nccl_ops.nccl_broadcast(input=tensor, shape=tensor.shape)

Returns a tensor that can be efficiently transferred to other devices.

Args:
        tensor: The tensor to send; must be assigned to a GPU device.

Returns:
        A tensor with the value of `src_tensor`, which can be used as input to
        ops on other GPU devices.


def update_state(self, y_true, y_pred, sample_weight=None):
    deps = []
    if not self._built:
        self._build(tensor_shape.TensorShape(y_pred.shape))
    if self.multi_label or self.label_weights is not None:
        shapes = [(y_true, ('N', 'L'))]
        if self.multi_label:
            shapes.extend([(self.true_positives, ('T', 'L')), (self.true_negatives, ('T', 'L')), (self.false_positives, ('T', 'L')), (self.false_negatives, ('T', 'L'))])
        if self.label_weights is not None:
            shapes.append((self.label_weights, ('L',)))
        deps = [check_ops.assert_shapes(shapes, message='Number of labels is not consistent.')]
    label_weights = None if self.multi_label else self.label_weights
    if self._from_logits:
        y_pred = activations.sigmoid(y_pred)
    with ops.control_dependencies(deps):
        return metrics_utils.update_confusion_matrix_variables({metrics_utils.ConfusionMatrix.TRUE_POSITIVES: self.true_positives, metrics_utils.ConfusionMatrix.TRUE_NEGATIVES: self.true_negatives, metrics_utils.ConfusionMatrix.FALSE_POSITIVES: self.false_positives, metrics_utils.ConfusionMatrix.FALSE_NEGATIVES: self.false_negatives}, y_true, y_pred, self._thresholds, thresholds_distributed_evenly=self._thresholds_distributed_evenly, sample_weight=sample_weight, multi_label=self.multi_label, label_weights=label_weights)

Accumulates confusion matrix statistics.

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

Returns:
        Update op.


def _get_base_converter_args(self):
    args = {'input_format': constants.TENSORFLOW_GRAPHDEF, 'allow_custom_ops': self.allow_custom_ops, 'debug_info': self._debug_info, 'target_ops': self.target_spec.supported_ops, 'select_user_tf_ops': self.target_spec.experimental_select_user_tf_ops, 'supported_backends': self.target_spec.experimental_supported_backends, 'unfold_batchmatmul': self.unfold_batchmatmul, 'legalize_custom_tensor_list_ops': self.legalize_custom_tensor_list_ops, 'lower_tensor_list_ops': self._experimental_lower_tensor_list_ops, 'unfold_large_splat_constant': self._experimental_unfold_large_splat_constant, 'default_to_single_batch_in_tensor_list_ops': self._experimental_default_to_single_batch_in_tensor_list_ops, 'tf_quantization_mode': self._experimental_tf_quantization_mode, 'experimental_enable_resource_variables': self.experimental_enable_resource_variables, 'enable_dynamic_update_slice': self._experimental_enable_dynamic_update_slice, 'preserve_assert_op': self._experimental_preserve_assert_op, 'guarantee_all_funcs_one_use': self._experimental_guarantee_all_funcs_one_use, 'allow_all_select_tf_ops': self._experimental_allow_all_select_tf_ops, 'disable_fuse_mul_and_fc': self._experimental_disable_fuse_mul_and_fc, 'quantization_options': self._experimental_quantization_options, 'ir_dump_dir': self.ir_dump_dir, 'ir_dump_pass_regex': self.ir_dump_pass_regex, 'ir_dump_func_regex': self.ir_dump_func_regex, 'enable_timing': self.enable_timing, 'print_ir_before': self.print_ir_before, 'print_ir_after': self.print_ir_after, 'print_ir_module_scope': self.print_ir_module_scope, 'elide_elementsattrs_if_larger': self.elide_elementsattrs_if_larger, 'use_buffer_offset': self._experimental_use_buffer_offset, 'reduce_type_precision': self._experimental_reduce_type_precision, 'use_stablehlo_quantizer': self.experimental_use_stablehlo_quantizer, 'stablehlo_quantizer_config': self.experimental_stablehlo_quantizer_config, 'qdq_conversion_mode': self._experimental_qdq_conversion_mode, 'strict_qdq_mode': self._experimental_strict_qdq, 'disable_per_channel_quantization_for_dense_layers': self._experimental_disable_per_channel_quantization_for_dense_layers, 'enable_composite_direct_lowering': self._experimental_enable_composite_direct_lowering, 'model_origin_framework': self.model_origin_framework, 'canonicalizing_inf_as_min_max_float': self.canonicalizing_inf_as_min_max_float, 'serialize_debug_metadata': self.serialize_debug_metadata, 'unsafe_fuse_dynamic_shaped_broadcast': self._experimental_unsafe_fuse_dynamic_shaped_broadcast}
    if self.saved_model_dir:
        args.update({'saved_model_dir': self.saved_model_dir, 'saved_model_version': self._saved_model_version, 'saved_model_tags': self._saved_model_tags, 'saved_model_exported_names': self._saved_model_exported_names})
    if self._experimental_quantization_options:
        logging.warning('Configs from custom methods in experimental_quantization_options may not produce a valid tflite model. Note that currently this option only supports StableHLO path. Setting this option in TFLite path will be a no-op.')
    if self.experimental_use_stablehlo_quantizer:
        self._assign_stablehlo_quantization_config_or_populate_default(args)
    elif self.experimental_stablehlo_quantizer_config is not None:
        raise ValueError('QuantizationConfig should be provided only when experimental_use_stablehlo_quantizer is set to true.')
    return args

Returns the base converter args.

Returns:
        {key str: val}


def __init__(self, qas_id, question_text, context_text, answer_text, start_position_character, title, answers=[], is_impossible=False):
    self.qas_id = qas_id
    self.question_text = question_text
    self.context_text = context_text
    self.answer_text = answer_text
    self.title = title
    self.is_impossible = is_impossible
    self.answers = answers
    self.start_position, self.end_position = (0, 0)
    doc_tokens = []
    char_to_word_offset = []
    prev_is_whitespace = True
    for c in self.context_text:
        if _is_whitespace(c):
            prev_is_whitespace = True
        else:
            if prev_is_whitespace:
                doc_tokens.append(c)
            else:
                doc_tokens[-1] += c
            prev_is_whitespace = False
        char_to_word_offset.append(len(doc_tokens) - 1)
    self.doc_tokens = doc_tokens
    self.char_to_word_offset = char_to_word_offset
    if start_position_character is not None and (not is_impossible):
        self.start_position = char_to_word_offset[start_position_character]
        self.end_position = char_to_word_offset[min(start_position_character + len(answer_text) - 1, len(char_to_word_offset) - 1)]

A single training/test example for the Squad dataset, as loaded from disk.

Args:
        qas_id: The example's unique identifier
        question_text: The question string
        context_text: The context string
        answer_text: The answer string
        start_position_character: The character position of the start of the answer
        title: The title of the example
        answers: None by default, this is used during evaluation. Holds answers as well as their start positions.
        is_impossible: False by default, set to True if the example has no possible answer.


def __init__(self, q1_tracker: Optional[QuantileTracker]=None, q3_tracker: Optional[QuantileTracker]=None, **kwargs):
    if 'threshold_criterion' not in kwargs:
        kwargs['threshold_criterion'] = FixedThreshold(1.5)
    super().__init__(**kwargs)
    self._q1_tracker = q1_tracker or BufferedSlidingQuantileTracker(DEFAULT_WINDOW_SIZE, 0.25)
    assert self._q1_tracker._q == 0.25, 'q1_tracker must be initialized with q = 0.25'
    self._q3_tracker = q3_tracker or SecondaryBufferedQuantileTracker(self._q1_tracker, 0.75)
    assert self._q3_tracker._q == 0.75, 'q3_tracker must be initialized with q = 0.75'

Interquartile Range (IQR) anomaly detector.

    This class implements an anomaly detection algorithm based on the
    Interquartile Range (IQR) [#]_ . It calculates the IQR using quantile trackers
    for Q1 (25th percentile) and Q3 (75th percentile) and scores data points based
    on their deviation from these quartiles.

    The score is calculated as follows:

        * If a data point is above Q3, the score is (value - Q3) / IQR.
        * If a data point is below Q1, the score is (Q1 - value) / IQR.
        * If a data point is within the IQR (Q1 <= value <= Q3), the score is 0.
        Initializes the IQR anomaly detector.

Args:
        q1_tracker: Optional QuantileTracker for Q1 (25th percentile). If None, a
        BufferedSlidingQuantileTracker with a default window size is used.
        q3_tracker: Optional QuantileTracker for Q3 (75th percentile). If None, a
        SecondaryBufferedQuantileTracker based on q1_tracker is used.
        threshold_criterion: Optional ThresholdFn to apply on the score. Defaults
        to `FixedThreshold(1.5)` since outliers are commonly defined as data
        points that fall below Q1 - 1.5 IQR or above Q3 + 1.5 IQR.
        **kwargs: Additional keyword arguments.

    .. [#] https://en.wikipedia.org/wiki/Interquartile_range


def dequantize_higgs(model, current_key_name=None):
    with torch.no_grad():
        for name, module in model.named_children():
            if current_key_name is None:
                current_key_name = []
            current_key_name.append(name)
            if isinstance(module, HiggsLinear):
                in_features = module.in_features
                out_features = module.out_features
                model._modules[name] = torch.nn.Linear(in_features, out_features, bias=module.bias is not None, device=module.scales.device, dtype=module.scales.dtype)
                model._modules[name].weight.data = module(torch.eye(in_features, device=module.scales.device, dtype=module.scales.dtype)).T.contiguous()
            if len(list(module.children())) > 0:
                _ = dequantize_higgs(module, current_key_name=current_key_name)
            current_key_name.pop(-1)
        return model

Dequantizes the HiggsLinear layers in the given model by replacing them with standard torch.nn.Linear layers.

Args:
        model (torch.nn.Module): The model containing HiggsLinear layers to be dequantized.
        current_key_name (list, optional): A list to keep track of the current module names during recursion. Defaults to None.

Returns:
        torch.nn.Module: The model with HiggsLinear layers replaced by torch.nn.Linear layers.


def add_extension_to_message(extension: message.Message, msg: message.Message) -> None:
    desc = msg.DESCRIPTOR
    fields_by_url = {get_inlined_extension_url(field): field for field in desc.fields if field.name != 'id'}
    id_field = desc.fields_by_name.get('id')
    if proto_utils.field_is_set(extension, id_field):
        proto_utils.set_value_at_field(msg, id_field, cast(Any, extension).id)
    if proto_utils.field_is_set(extension, 'value'):
        if len(fields_by_url) != 1:
            raise fhir_errors.InvalidFhirError(f'Expected a single field, found {len(fields_by_url)}; {desc.full_name} is an invalid extension type.')
        field = list(fields_by_url.items())[0][1]
        if proto_utils.field_is_repeated(field):
            raise fhir_errors.InvalidFhirError(f'Expected {field.full_name} to be a singular field. {desc.full_name} is an invalid extension type.')
        _add_extension_value_to_message(extension, msg, field)
        return
    child_extensions = proto_utils.get_value_at_field(extension, 'extension')
    for child_extension in child_extensions:
        field = fields_by_url.get(child_extension.url.value)
        if field is None:
            raise ValueError(f'Message of type: {desc.full_name} has no field with name: {child_extension.url.value}.')
        if proto_utils.field_is_set(child_extension, 'value'):
            _add_extension_value_to_message(child_extension, msg, field)
            continue
        if not proto_utils.field_is_repeated(field):
            if proto_utils.field_is_set(msg, field):
                raise ValueError(f'Field: {field.full_name} is already set on message: {desc.full_name}.')
            if proto_utils.field_content_length(child_extension, 'extension') > 1:
                raise ValueError(f'Cardinality mismatch between field: {field.full_name} and extension: {desc.full_name}.')
        child_message = proto_utils.set_in_parent_or_add(msg, field)
        add_extension_to_message(child_extension, child_message)

Recursively parses extension and adds to message.

Args:
        extension: The FHIR extension to serialize and add.
        msg: The message to add the extension onto

Raises:
        InvalidFhirError: In the event that a value is set on the extension, but the
        corresponding message field to copy it to is repeated (extension values are
        singular only).


def early_stop_by_value(step_values: List[Tuple[int, float]], metric: Union[str, Callable[[pg.tuning.Measurement], float]]='reward', maximize: bool=True):
    assert isinstance(step_values, list), step_values
    for v in step_values:
        if not isinstance(v, tuple) or len(v) != 2 or (not isinstance(v[0], int)) or (not isinstance(v[1], numbers.Number)):
            raise ValueError(f'Invalid definition in `step_values`: {v}. Expect a tuple of 2 elements: (step: int, threshold: float).')

    def _cmp(x, y) -> bool:
        return x < y if maximize else x > y

    def _value(m: pg.tuning.Measurement) -> float:
        if isinstance(metric, str):
            return m.reward if metric == 'reward' else m.metrics[metric]
        assert callable(metric), metric
        return metric(m)

    def _make_predicate(threshold: float):

        def _predicate(m: pg.tuning.Measurement, unused_history):
            v = _value(m)
            ret = _cmp(v, threshold)
            return ret
        return _predicate
    return StepWise([(step, _make_predicate(threshold)) for step, threshold in step_values])

Step-wise early stopping policy based on the value of reward/metric.

    Example::

        policy = early_stop_by_value([
        # Stop at step 1 if trial reward is less than 0.2.
        (1, 0.2),

        # Stop at step 2 if trial reward is less than 0.8.
        (2, 0.8),
        ])()

Args:
        step_values: A list of tuple (gating step, value threshold).
        gating step - At which step this rule will be triggered.
        value threshold - A float number indicating the threshold value for
            early stopping.
        metric: Based on which metric the value should be compared against.
        Use str for metric name or a callable object that takes a measurement
        object at a given step as input and returns a float value.
        maximize: If True, reward or metric value below the threshold will be
        stopped, otherwise trials with values above the threshold will be stopped.

Returns:
        A `StepWise` early stopping policy.


def __rmul__(self, other):
    return self * other

Returns the product of `self` and `other`.

Args:
        other: Another Dimension, or a value accepted by `as_dimension`.

Returns:
        A Dimension whose value is the product of `self` and `other`.


def reparameterization_type(self):
    return self._reparameterization_type

Describes how samples from the distribution are reparameterized.

    Currently this is one of the static instances
    `distributions.FULLY_REPARAMETERIZED`
    or `distributions.NOT_REPARAMETERIZED`.

Returns:
        An instance of `ReparameterizationType`.


def CanSplit(self, must_split):
    current = self.next_token
    previous = current.previous_token
    if current.is_pseudo:
        return False
    if not must_split and subtypes.DICTIONARY_KEY_PART in current.subtypes and (subtypes.DICTIONARY_KEY not in current.subtypes) and (not style.Get('ALLOW_MULTILINE_DICTIONARY_KEYS')):
        return False
    if not must_split and subtypes.DICTIONARY_VALUE in current.subtypes and (not style.Get('ALLOW_SPLIT_BEFORE_DICT_VALUE')):
        return False
    if previous and previous.value == '(' and (current.value == ')'):
        token = previous.previous_token
        while token:
            prev = token.previous_token
            if not prev or prev.name not in {'NAME', 'DOT'}:
                break
            token = token.previous_token
        if token and subtypes.DICTIONARY_VALUE in token.subtypes:
            if not style.Get('ALLOW_SPLIT_BEFORE_DICT_VALUE'):
                return False
    if previous and previous.value == '.' and (current.value == '.'):
        return False
    return current.can_break_before

Determine if we can split before the next token.

Args:
        must_split: (bool) A newline was required before this token.

Returns:
        True if the line can be split before the next token.


def to_code(entity, recursive=True, experimental_optional_features=None):
    source = tf_inspect.getsource(to_graph(entity, recursive=recursive, experimental_optional_features=experimental_optional_features))
    return textwrap.dedent(source)

Returns the source code generated by AutoGraph, as a string.

    Example usage:

    >>> def f(x):
    ...   if x < 0:
    ...     x = -x
    ...   return x
    >>> tf.autograph.to_code(f)
    "...def tf__f(x):..."

    Also see: `tf.autograph.to_graph`.

    Note: If a function has been decorated with `tf.function`, pass its
    underlying Python function, rather than the callable that `tf.function
    creates:

    >>> @tf.function
    ... def f(x):
    ...   if x < 0:
    ...     x = -x
    ...   return x
    >>> tf.autograph.to_code(f.python_function)
    "...def tf__f(x):..."

Args:
        entity: Python callable or class to convert.
        recursive: Whether to recursively convert any functions that the converted
        function may call.
        experimental_optional_features: `None`, a tuple of, or a single
        `tf.autograph.experimental.Feature` value.

Returns:
        The converted code as string.


def year_month_day_to_ordinal(year, month, day):
    with tf.compat.v1.name_scope(None, 'ymd2o', [year, month, day]):
        year = tf.convert_to_tensor(year, tf.int32, name='year')
        month = tf.convert_to_tensor(month, tf.int32, name='month')
        day = tf.convert_to_tensor(day, tf.int32, name='day')
        year -= tf.compat.v2.where(month <= 2, 1, 0)
        month += tf.compat.v2.where(month > 2, -3, 9)
        era = year // _YEARS_IN_ERA
        year_of_era = year % _YEARS_IN_ERA
        day_of_year = _days_in_year_before_month(month) + day - 1
        day_of_era = year_of_era * _DAYS_IN_YEAR + year_of_era // 4 - year_of_era // 100 + day_of_year
        return era * _DAYS_IN_ERA + day_of_era + _ORDINAL_OF_1_3_0000

Calculates ordinals Tensor given years, months and dates.

Args:
        year: Tensor of int32 type. Elements should be positive.
        month: Tensor of int32 type of same shape as `year`. Elements should be in
        range `[1, 12]`.
        day: Tensor of int32 type of same shape as `year`. Elements should be in
        range `[1, 31]` and represent valid dates together with corresponding
        elements of `month` and `year` Tensors.

Returns:
        Tensor of int32 type. Each element is number of days since 1 Jan 0001. 1 Jan
        0001 has `ordinal = 1`.


def smart_constant_value(pred):
    if isinstance(pred, tensor.Tensor):
        pred_value = tensor_util.constant_value(pred)
        if pred_value is None:
            pred_value = tensor_util.try_evaluate_constant(pred)
    elif pred in {0, 1}:
        pred_value = bool(pred)
    elif isinstance(pred, bool):
        pred_value = pred
    else:
        raise TypeError(f'Argument `pred` must be a Tensor, or a Python bool, or 1 or 0. Received: pred={pred} of type {type(pred).__name__}')
    return pred_value

Return the bool value for `pred`, or None if `pred` had a dynamic value.

Args:
        pred: A scalar, either a Python bool or tensor.

Returns:
        True or False if `pred` has a constant boolean value, None otherwise.

Raises:
        TypeError: If `pred` is not a Tensor or bool.


def FormatFiles(filenames, lines, style_config=None, no_local_style=False, in_place=False, print_diff=False, parallel=False, quiet=False, verbose=False, print_modified=False):
    changed = False
    if parallel:
        import concurrent.futures
        import multiprocessing
        workers = min(multiprocessing.cpu_count(), len(filenames))
        with concurrent.futures.ProcessPoolExecutor(workers) as executor:
            future_formats = [executor.submit(_FormatFile, filename, lines, style_config, no_local_style, in_place, print_diff, quiet, verbose, print_modified) for filename in filenames]
            for future in concurrent.futures.as_completed(future_formats):
                changed |= future.result()
    else:
        for filename in filenames:
            changed |= _FormatFile(filename, lines, style_config, no_local_style, in_place, print_diff, quiet, verbose, print_modified)
    return changed

Format a list of files.

Args:
        filenames: (list of unicode) A list of files to reformat.
        lines: (list of tuples of integers) A list of tuples of lines, [start, end],
        that we want to format. The lines are 1-based indexed. This argument
        overrides the 'args.lines'. It can be used by third-party code (e.g.,
        IDEs) when reformatting a snippet of code.
        style_config: (string) Style name or file path.
        no_local_style: (string) If style_config is None don't search for
        directory-local style configuration.
        in_place: (bool) Modify the files in place.
        print_diff: (bool) Instead of returning the reformatted source, return a
        diff that turns the formatted source into reformatter source.
        parallel: (bool) True if should format multiple files in parallel.
        quiet: (bool) True if should output nothing.
        verbose: (bool) True if should print out filenames while processing.
        print_modified: (bool) True if should print out filenames of modified files.

Returns:
        True if the source code changed in any of the files being formatted.


def save_randomly_initialized_version(config_name: str, save_dir: str, **config_kwargs):
    cfg = AutoConfig.from_pretrained(config_name, **config_kwargs)
    model = AutoModelForSeq2SeqLM.from_config(cfg)
    model.save_pretrained(save_dir)
    AutoTokenizer.from_pretrained(config_name).save_pretrained(save_dir)
    return model

Save a randomly initialized version of a model using a pretrained config.

Args:
        config_name: which config to use
        save_dir: where to save the resulting model and tokenizer
        config_kwargs: Passed to AutoConfig

    Usage::
        save_randomly_initialized_version("facebook/bart-large-cnn", "distilbart_random_cnn_6_3", encoder_layers=6, decoder_layers=3, num_beams=3)


def get_prompt_embeddings(self, input_points: tf.Tensor | None=None, input_labels: tf.Tensor | None=None, input_boxes: tf.Tensor | None=None, input_masks: tf.Tensor | None=None):
    prompt_output = self.prompt_encoder(input_points=input_points, input_labels=input_labels, input_boxes=input_boxes, input_masks=input_masks)
    return prompt_output

Returns the prompt embeddings by passing the input points, labels, boxes and masks through the prompt encoder.

Args:
        input_points (`tf.Tensor` of shape `(batch_size, point_batch_size, num_points_per_image, 2)`):
            Optional input points for the prompt encoder. The padding of the point is automatically done by the
            processor. `point_batch_size` refers to the number of masks that we want the model to predict per
            point. The model will output `point_batch_size` times 3 masks in total.
        input_labels (`tf.Tensor` of shape `(batch_size, point_batch_size, num_points_per_image)`):
            Optional input labels for the prompt encoder. The padding of the labels is automatically done by the
            processor, or can be fed by the user.
        input_boxes (`tf.Tensor` of shape `(batch_size, num_boxes_per_image, 4)`):
            Optional input boxes for the prompt encoder. The padding of the boxes is automatically done by the
            processor. users can also pass manually the input boxes.
        input_masks (`tf.Tensor` of shape `(batch_size, image_size, image_size)`):
            Optional input masks for the prompt encoder.


def forward(self, input_ids: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None, position_ids: Optional[torch.Tensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[Tuple, ClapTextModelOutput]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    text_outputs = self.text_model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    pooled_output = text_outputs[1] if not return_dict else text_outputs.pooler_output
    text_embeds = self.text_projection(pooled_output)
    if not return_dict:
        outputs = (text_embeds, text_outputs[0]) + text_outputs[2:]
        return tuple((output for output in outputs if output is not None))
    return ClapTextModelOutput(text_embeds=text_embeds, last_hidden_state=text_outputs.last_hidden_state, hidden_states=text_outputs.hidden_states, attentions=text_outputs.attentions)

Examples:

    ```python
    >>> from transformers import AutoTokenizer, ClapTextModelWithProjection

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

    >>> inputs = tokenizer(["a sound of a cat", "a sound of a dog"], padding=True, return_tensors="pt")

    >>> outputs = model(**inputs)
    >>> text_embeds = outputs.text_embeds
    ```


def For(start, limit, delta, inputs, body, name=None, hostmem=None, rewrite_with_while=None):
    if rewrite_with_while:
        return _ForUsingWhile(start, limit, delta, inputs, body, name, hostmem)
    if body.captured_inputs:
        ret = gen_functional_ops._for(start, limit, delta, inputs + body.captured_inputs, _LoopBodyCaptureWrapper(body), name=name)
        ret = ret[:-len(body.captured_inputs)]
    else:
        ret = gen_functional_ops._for(start, limit, delta, inputs, body, name=name)
    if hostmem:
        num_for_params = 3
        input_attr = attr_value_pb2.AttrValue()
        input_attr.list.i.extend([num_for_params + i for i in hostmem])
        ret[0].op._set_attr('_input_hostmem', input_attr)
        output_attr = attr_value_pb2.AttrValue()
        output_attr.list.i.extend(hostmem)
        ret[0].op._set_attr('_output_hostmem', output_attr)
    return ret

out = input; for i in range(start, limit, delta) out = body(i, out).

Args:
        start: A `Tensor` of type `int32`.
        limit: A `Tensor` of type `int32`.
        delta: A `Tensor` of type `int32`.
        inputs: A list of `Tensor` objects. A list of input tensors whose types are
        T.
        body: A function takes a list of tensors and returns another list of
        tensors. Both lists have the same types as (int32, T...).
        name: A name for the operation (optional).
        hostmem: A list of integer. If i is in the list, inputs[i] is a host memory
        tensor. In other words, (i+1)-th argument of the body function is
        expecting a host memory.
        rewrite_with_while: If True, using While op to implement the For.

Returns:
        A list of `Tensor` objects. Has the same type as `input`.
        A list of output tensors whose types are T.


def extract_sequences(x, sequence_length, sequence_stride):
    if any_symbolic_tensors((x,)):
        return ExtractSequences(sequence_length, sequence_stride).symbolic_call(x)
    return backend.math.extract_sequences(x, sequence_length, sequence_stride)

Expands the dimension of last axis into sequences of `sequence_length`.

    Slides a window of size `sequence_length` over the last axis of the input
    with a stride of `sequence_stride`, replacing the last axis with
    `[num_sequences, sequence_length]` sequences.

    If the dimension along the last axis is N, the number of sequences can be
    computed by:

    `num_sequences = 1 + (N - sequence_length) // sequence_stride`

Args:
        x: Input tensor.
        sequence_length: An integer representing the sequences length.
        sequence_stride: An integer representing the sequences hop size.

Returns:
        A tensor of sequences with shape [..., num_sequences, sequence_length].

Example:
    >>> x = keras.ops.convert_to_tensor([1, 2, 3, 4, 5, 6])
    >>> extract_sequences(x, 3, 2)
    array([[1, 2, 3],
        [3, 4, 5]])


def set_callback_parameters(callback_list, model, do_validation=False, batch_size=None, epochs=None, steps_per_epoch=None, samples=None, verbose=1, mode=ModeKeys.TRAIN):
    metric_names = model.metrics_names
    for cbk in callback_list:
        if isinstance(cbk, (BaseLogger, ProgbarLogger)):
            cbk.stateful_metrics = metric_names[1:]
    callback_metrics = []
    if mode != ModeKeys.PREDICT:
        callback_metrics = copy.copy(metric_names)
        if do_validation:
            callback_metrics += ['val_' + n for n in metric_names]
    callback_params = {'batch_size': batch_size, 'epochs': epochs, 'steps': steps_per_epoch, 'samples': samples, 'verbose': verbose, 'do_validation': do_validation, 'metrics': callback_metrics}
    callback_list.set_params(callback_params)

Sets callback parameters.

Args:
        callback_list: CallbackList instance.
        model: Model being trained.
        do_validation: Whether or not validation loop will be run.
        batch_size: Number of samples per batch.
        epochs: Number of epoch to train.
        steps_per_epoch: Number of batches to run per training epoch.
        samples: Number of training samples.
        verbose: int, 0 or 1. Keras logging verbosity to pass to ProgbarLogger.
        mode: String. One of ModeKeys.TRAIN, ModeKeys.TEST, or ModeKeys.PREDICT.
        Which loop mode to configure callbacks for.


def remove_hallucinated_references(self, text: str) -> str:
    lines = text.split('\n')
    if len(lines) == 0:
        return ''
    clean_lines = remove_numbers(lines)
    slices = get_slices(lines, clean_lines)
    to_delete = []
    for slice in slices:
        to_delete.append(remove_slice_from_lines(lines, clean_lines, slice))
    for to_delete in reversed(to_delete):
        text = text.replace(to_delete, '\n\n[MISSING_PAGE_POST]\n\n')
    text = re.sub('## References\\n+\\[MISSING_PAGE_POST(:\\d+)?\\]', '\n\n[MISSING_PAGE_POST\\1]', text)
    return text

Remove hallucinated or missing references from the text.

    This function identifies and removes references that are marked as missing or hallucinated from the input text.

Args:
        text (`str`):
            The input text containing references.

Returns:
        `str`: The text with hallucinated references removed.


def tpu_core_locations_to_ids(self, tpu_core_locations):
    return _pywrap_dtensor_device.TPUCoreLocationsToIDs(context.context()._handle, self._device_info, tpu_core_locations)

Translates TPU core locations to TPU core IDs.

Args:
        tpu_core_locations: A list of TPU core locations. Each one is a list of
        four unsigned integers, [x, y, z, core].

Returns:
        A list of corresponding TPU core IDs.


def add_debug_tensor_watch(run_options, node_name, output_slot=0, debug_ops='DebugIdentity', debug_urls=None, tolerate_debug_op_creation_failures=False, global_step=-1):
    watch_opts = run_options.debug_options.debug_tensor_watch_opts
    run_options.debug_options.global_step = global_step
    watch = watch_opts.add()
    watch.tolerate_debug_op_creation_failures = tolerate_debug_op_creation_failures
    watch.node_name = node_name
    watch.output_slot = output_slot
    if isinstance(debug_ops, str):
        debug_ops = [debug_ops]
    watch.debug_ops.extend(debug_ops)
    if debug_urls:
        if isinstance(debug_urls, str):
            debug_urls = [debug_urls]
        watch.debug_urls.extend(debug_urls)

Add watch on a `Tensor` to `RunOptions`.

    N.B.:
        1. Under certain circumstances, the `Tensor` may not get actually watched
        (e.g., if the node of the `Tensor` is constant-folded during runtime).
        2. For debugging purposes, the `parallel_iteration` attribute of all
        `tf.while_loop`s in the graph are set to 1 to prevent any node from
        being executed multiple times concurrently. This change does not affect
        subsequent non-debugged runs of the same `tf.while_loop`s.

Args:
        run_options: An instance of `config_pb2.RunOptions` to be modified.
        node_name: (`str`) name of the node to watch.
        output_slot: (`int`) output slot index of the tensor from the watched node.
        debug_ops: (`str` or `list` of `str`) name(s) of the debug op(s). Can be a
        `list` of `str` or a single `str`. The latter case is equivalent to a
        `list` of `str` with only one element.
        For debug op types with customizable attributes, each debug op string can
        optionally contain a list of attribute names, in the syntax of:
        debug_op_name(attr_name_1=attr_value_1;attr_name_2=attr_value_2;...)
        debug_urls: (`str` or `list` of `str`) URL(s) to send debug values to,
        e.g., `file:///tmp/tfdbg_dump_1`, `grpc://localhost:12345`.
        tolerate_debug_op_creation_failures: (`bool`) Whether to tolerate debug op
        creation failures by not throwing exceptions.
        global_step: (`int`) Optional global_step count for this debug tensor
        watch.


def scatter_sub(self, sparse_delta, use_locking=False, name=None):
    raise NotImplementedError

Subtracts `tf.IndexedSlices` from this variable.

Args:
        sparse_delta: `tf.IndexedSlices` to be subtracted from this variable.
        use_locking: If `True`, use locking during the operation.
        name: the name of the operation.

Returns:
        The updated variable.

Raises:
        TypeError: if `sparse_delta` is not an `IndexedSlices`.


def _preprocess_token_ids(self, token_ids, skip_special_tokens: bool=False):
    if skip_special_tokens:
        prompt_token_id = self.convert_tokens_to_ids('<|startofprev|>')
        decoder_start_token_id = self.convert_tokens_to_ids('<|startoftranscript|>')
        token_ids = self._strip_prompt(token_ids, prompt_token_id, decoder_start_token_id)
    return token_ids

Pre-process the token ids for decoding by removing the prompt tokens ids and timestamp token ids.

Args:
        token_ids (`Union[int, List[int], np.ndarray, torch.Tensor, tf.Tensor]`):
            List of tokenized input ids. Typically, obtained using the `__call__` method of the tokenizer.
        skip_special_tokens (`bool`, *optional*, defaults to `False`):
            Whether or not to remove special tokens from the token ids. If `True`, the prompt token ids will be
            removed.


def recipe_sheets_copy(config, auth_read, from_sheet, from_tab, to_sheet, to_tab):
    sheets(config, {'auth': auth_read, 'template': {'sheet': from_sheet, 'tab': from_tab}, 'sheet': to_sheet, 'tab': to_tab})

Copy tab from a sheet to a sheet.

Args:
        auth_read (authentication) - Credentials used for reading data.
        from_sheet (string) - NA
        from_tab (string) - NA
        to_sheet (string) - NA
        to_tab (string) - NA


def __init__(self, conf, map_name, automount_mountpoint=None):
    super(Cache, self).__init__()
    self.log = logging.getLogger(__name__)
    self.conf = conf
    self.output_dir = conf.get('dir', '.')
    self.automount_mountpoint = automount_mountpoint
    self.map_name = map_name
    if map_name == config.MAP_PASSWORD:
        self.data = passwd.PasswdMap()
    elif map_name == config.MAP_SSHKEY:
        self.data = sshkey.SshkeyMap()
    elif map_name == config.MAP_GROUP:
        self.data = group.GroupMap()
    elif map_name == config.MAP_SHADOW:
        self.data = shadow.ShadowMap()
    elif map_name == config.MAP_NETGROUP:
        self.data = netgroup.NetgroupMap()
    elif map_name == config.MAP_AUTOMOUNT:
        self.data = automount.AutomountMap()
    else:
        raise error.UnsupportedMap('Cache does not support %s' % map_name)

Initialise the Cache object.

Args:
        conf: A dictionary of key/value pairs
        map_name: A string representation of the map type
        automount_mountpoint: A string containing the automount mountpoint,
        used only by automount maps.

Raises:
        UnsupportedMap: for map types we don't know about


def get(self):
    return self._get_helper(self._sorted_items, self._q)

Returns the current quantile value using the sorted list.

    Calculates the quantile using linear interpolation on the sorted values.

Returns:
        float: The calculated quantile value. Returns NaN if the window is empty.


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

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

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

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


def keyword_args_only(func):
    decorator_utils.validate_callable(func, 'keyword_args_only')

    @functools.wraps(func)
    def new_func(*args, **kwargs):
        """Keyword args only wrapper."""
        if args:
            raise ValueError(f'The function {func.__name__} only accepts keyword arguments. Do not pass positional arguments. Received the following positional arguments: {args}')
        return func(**kwargs)
    return new_func

Decorator for marking specific function accepting keyword args only.

    This decorator raises a `ValueError` if the input `func` is called with any
    non-keyword args. This prevents the caller from providing the arguments in
    wrong order.

Args:
        func: The function or method needed to be decorated.

Returns:
        Decorated function or method.

Raises:
        ValueError: If `func` is not callable.


def get_operation_by_name(self, name) -> 'Operation':
    if not isinstance(name, str):
        raise TypeError('Operation names are strings (or similar), not %s.' % type(name).__name__)
    op = cast(Operation, self.as_graph_element(name, allow_tensor=False, allow_operation=True))
    return op

Returns the `Operation` with the given `name`.

    This method may be called concurrently from multiple threads.

Args:
        name: The name of the `Operation` to return.

Returns:
        The `Operation` with the given `name`.

Raises:
        TypeError: If `name` is not a string.
        KeyError: If `name` does not correspond to an operation in this graph.


def visit_boolean_op(self, boolean_logic: _evaluation.BooleanOperatorNode) -> _sql_data_types.Select:
    lhs_result = self.visit(boolean_logic.left)
    rhs_result = self.visit(boolean_logic.right)
    if lhs_result.sql_data_type != _sql_data_types.Boolean:
        lhs_result = lhs_result.is_not_null()
    if rhs_result.sql_data_type != _sql_data_types.Boolean:
        rhs_result = rhs_result.is_not_null()
    lhs_subquery = lhs_result.as_operand()
    rhs_subquery = rhs_result.as_operand()
    if boolean_logic.op == _ast.BooleanLogic.Op.IMPLIES:
        sql_value = f'NOT {lhs_subquery} OR {rhs_subquery}'
    elif boolean_logic.op == _ast.BooleanLogic.Op.XOR:
        sql_value = f'{lhs_subquery} <> {rhs_subquery}'
    else:
        sql_value = f'{lhs_subquery} {boolean_logic.op.upper()} {rhs_subquery}'
    return _sql_data_types.Select(select_part=_sql_data_types.RawExpression(sql_value, _sql_data_type=_sql_data_types.Boolean, _sql_alias='logic_'), from_part=None, sql_dialect=_sql_data_types.SqlDialect.SPARK)

Translates a FHIRPath Boolean logic operation to Spark SQL.

    Note that evaluation for Boolean logic is only supported for Boolean
    operands of scalar cardinality.

Args:
        boolean_logic: The FHIRPath AST `BooleanLogic` node.

Returns:
        A compiled Spark SQL expression.


def compute_output_signature(self, input_signature):

    def check_type_return_shape(s):
        if not isinstance(s, tensor.TensorSpec):
            raise TypeError('Only TensorSpec signature types are supported, but saw signature entry: {}.'.format(s))
        return s.shape
    input_shape = nest.map_structure(check_type_return_shape, input_signature)
    output_shape = self.compute_output_shape(input_shape)
    dtype = self._compute_dtype
    if dtype is None:
        input_dtypes = [s.dtype for s in nest.flatten(input_signature)]
        dtype = input_dtypes[0]
    return nest.map_structure(lambda s: tensor.TensorSpec(dtype=dtype, shape=s), output_shape)

Compute the output tensor signature of the layer based on the inputs.

    Unlike a TensorShape object, a TensorSpec object contains both shape
    and dtype information for a tensor. This method allows layers to provide
    output dtype information if it is different from the input dtype.
    For any layer that doesn't implement this function,
    the framework will fall back to use `compute_output_shape`, and will
    assume that the output dtype matches the input dtype.

Args:
        input_signature: Single TensorSpec or nested structure of TensorSpec
        objects, describing a candidate input for the layer.

Returns:
        Single TensorSpec or nested structure of TensorSpec objects, describing
        how the layer would transform the provided input.

Raises:
        TypeError: If input_signature contains a non-TensorSpec object.


def matches(self, applied_ptransform):
    raise NotImplementedError

Determines whether the given AppliedPTransform matches.

    Note that the matching will happen *after* Runner API proto translation.
    If matching is done via type checks, to/from_runner_api[_parameter] methods
    must be implemented to preserve the type (and other data) through proto
    serialization.

    Consider URN-based translation instead.

Args:
        applied_ptransform: AppliedPTransform to be matched.

Returns:
        a bool indicating whether the given AppliedPTransform is a match.


def _GetAxisFromLabel(subscripts, label):
    splits = subscripts.split(ellipsis)
    index = splits[0].find(label)
    if index != -1:
        return index
    if len(splits) < 2:
        return None
    index = splits[1].find(label)
    if index != -1:
        return index - len(splits[1])
    return None

Returns the axis (possibly negative) corresponding to a label.

    Returns the axis index of the axis label if it is before an ellipsis (or if
    the ellipsis is not present), and the negative index if it occurs after the
    ellipsis. E.g. index of `b` in `ab...cd`, is `1`, but that of `c` is `-2`.

    For multiple occurrences, returns the leftmost one. If not found, returns
    None.

Args:
        subscripts: A string denoting the einsum subscript (e.g. `ab...cd`)
        label: The single character axis label.


def _list_profile_sort_key(profile_datum, sort_by):
    if sort_by == SORT_OPS_BY_OP_NAME:
        return profile_datum.node_exec_stats.node_name
    elif sort_by == SORT_OPS_BY_OP_TYPE:
        return profile_datum.op_type
    elif sort_by == SORT_OPS_BY_LINE:
        return profile_datum.file_line_func
    elif sort_by == SORT_OPS_BY_OP_TIME:
        return profile_datum.op_time
    elif sort_by == SORT_OPS_BY_EXEC_TIME:
        return profile_datum.node_exec_stats.all_end_rel_micros
    else:
        return profile_datum.node_exec_stats.all_start_micros

Get a profile_datum property to sort by in list_profile command.

Args:
        profile_datum: A `ProfileDatum` object.
        sort_by: (string) indicates a value to sort by.
        Must be one of SORT_BY* constants.

Returns:
        profile_datum property to sort by.


def __init__(self, meshes: List[layout_lib.Mesh], is_async=True, in_flight_nodes_limit=8):
    if any((not isinstance(mesh, layout_lib.Mesh) for mesh in meshes)):
        raise TypeError('Expected a flat list of Mesh objects, got {}'.format(meshes))
    global _next_device_number
    ctx = context.context()
    with _next_device_number_lock:
        self.name = '{}/device:CUSTOM:{}'.format(ctx.host_address_space(), _next_device_number)
        _next_device_number += 1
    device, device_info = _pywrap_dtensor_device.Allocate(self.name, is_async, in_flight_nodes_limit)
    context.register_custom_device(device, self.name, device_info)
    self._device_info = device_info
    self._current_output_layout = None
    self._current_default_mesh = None
    self._meshes = set()
    self._mesh_lock = threading.Lock()
    for mesh in meshes:
        self._register_mesh(mesh)

Create a new DTensorDevice which executes ops on `underlying_device`.

Args:
        meshes: A list of `Mesh` objects indicating groups of devices to execute
        on. These may also be registered lazily.
        is_async: Indicates whether DTensor operations on this client will return
        immediately (with "non-ready" handles) or block until executed. This is
        on by default and is exposed as an option for ease of debugging.
        in_flight_nodes_limit: Indicates the limit of in-flight nodes before
        enqueueing of async operations to DTensorDevice is blocked. This limit
        is per mesh. 0 for no limits from DTensor. Default is 8.


def get_video_features(self, pixel_values: torch.FloatTensor, vision_feature_layer: Optional[Union[int, List[int]]]=None, vision_feature_select_strategy: Optional[str]=None):
    vision_feature_layer = vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer
    vision_feature_select_strategy = vision_feature_select_strategy if vision_feature_select_strategy is not None else self.config.vision_feature_select_strategy
    batch_size, frames, channels, height, width = pixel_values.shape
    pixel_values = pixel_values.reshape(batch_size * frames, channels, height, width)
    video_features = self.vision_tower(pixel_values, output_hidden_states=True)
    if isinstance(vision_feature_layer, int):
        selected_video_features = video_features.hidden_states[vision_feature_layer]
    else:
        hs_pool = [video_features.hidden_states[layer_idx] for layer_idx in vision_feature_layer]
        selected_video_features = torch.cat(hs_pool, dim=-1)
    if vision_feature_select_strategy == 'default':
        selected_video_features = selected_video_features[:, 1:]
    elif vision_feature_select_strategy == 'full':
        selected_video_features = selected_video_features
    video_features = self.vision_resampler(selected_video_features)
    video_features = self.multi_modal_projector(video_features)
    video_features = torch.split(video_features, frames, dim=0)
    return video_features

Obtains video last hidden states from the vision tower and apply multimodal projection.

Args:
        pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_frames, channels, height, width)`)
        The tensors corresponding to the input video.
        vision_feature_layer (`Union[int, List[int]]`, *optiona;*):
            The index of the layer to select the vision feature. If multiple indices are provided,
            the vision feature of the corresponding indices will be concatenated to form the
            vision features.
        vision_feature_select_strategy (`str`, *optional*):
            The feature selection strategy used to select the vision feature from the vision backbone.
            Can be one of `"default"` or `"full"`

Returns:
        video_features (List[`torch.Tensor`]): List of video feature tensor, each contains all the visual feature of all patches
        and are of shape `(num_videos, video_length, embed_dim)`).


def cd(path):
    if not path:
        yield
        return
    curdir = path_utils.getcwd()
    os.chdir(path)
    try:
        yield
    finally:
        os.chdir(curdir)

Context manager. Change the directory, and restore it afterwards.

    Example usage:
        with cd("/path"):
        ...

Args:
        path: The directory to change to. If empty, this function is a no-op.

Yields:
        Executes your code, in a changed directory.


def _MakeParseFn(fn, metadata):
    fn_spec = inspectutils.GetFullArgSpec(fn)
    num_required_args = len(fn_spec.args) - len(fn_spec.defaults)
    required_kwonly = set(fn_spec.kwonlyargs) - set(fn_spec.kwonlydefaults)

    def _ParseFn(args):
        """Parses the list of `args` into (varargs, kwargs), remaining_args."""
        kwargs, remaining_kwargs, remaining_args = _ParseKeywordArgs(args, fn_spec)
        parsed_args, kwargs, remaining_args, capacity = _ParseArgs(fn_spec.args, fn_spec.defaults, num_required_args, kwargs, remaining_args, metadata)
        if fn_spec.varargs or fn_spec.varkw:
            capacity = True
        extra_kw = set(kwargs) - set(fn_spec.kwonlyargs)
        if fn_spec.varkw is None and extra_kw:
            raise FireError('Unexpected kwargs present:', extra_kw)
        missing_kwonly = set(required_kwonly) - set(kwargs)
        if missing_kwonly:
            raise FireError('Missing required flags:', missing_kwonly)
        if fn_spec.varargs is not None:
            varargs, remaining_args = (remaining_args, [])
        else:
            varargs = []
        for index, value in enumerate(varargs):
            varargs[index] = _ParseValue(value, None, None, metadata)
        varargs = parsed_args + varargs
        remaining_args += remaining_kwargs
        consumed_args = args[:len(args) - len(remaining_args)]
        return ((varargs, kwargs), consumed_args, remaining_args, capacity)
    return _ParseFn

Creates a parse function for fn.

Args:
        fn: The function or class to create the parse function for.
        metadata: Additional metadata about the component the parse function is for.

Returns:
        A parse function for fn. The parse function accepts a list of arguments
        and returns (varargs, kwargs), remaining_args. The original function fn
        can then be called with fn(*varargs, **kwargs). The remaining_args are
        the leftover args from the arguments to the parse function.


def _FindNodeWithStandaloneLineParent(node):
    if pytree_utils.NodeName(node.parent) in _STANDALONE_LINE_NODES:
        return node
    else:
        return _FindNodeWithStandaloneLineParent(node.parent)

Find a node whose parent is a 'standalone line' node.

    See the comment above _STANDALONE_LINE_NODES for more details.

Args:
        node: node to start from

Returns:
        Suitable node that's either the node itself or one of its ancestors.


def _encode_reference_type_constraints(self, builder: expressions.Builder, elem: message.Message) -> List[validation_pb2.SqlRequirement]:
    field_name = _last_path_token(builder)
    constraint_key = f'{field_name}-resource-type-exclusivity'
    if constraint_key in self._options.skip_keys:
        return []
    element_definition = cast(Any, elem)
    type_codes = _utils.element_type_codes(element_definition)
    if type_codes != ['Reference']:
        return []
    allowed_reference_types = [target_profile.value for target_profile in element_definition.type[0].target_profile]
    if len(allowed_reference_types) <= 1:
        return []
    num_references_exist: expressions.Builder = _num_fields_exist((builder.getReferenceKey(reference_type) for reference_type in sorted(allowed_reference_types)))
    constraint: expressions.Builder = num_references_exist <= 1
    if _fhir_path_data_types.is_collection(builder.return_type):
        constraint: expressions.Builder = builder.all(constraint)
    constraint_sql = self._encode_fhir_path_builder_constraint(constraint, builder.get_parent_builder())
    if constraint_sql is None:
        return []
    reference_type_path = self._abs_path_invocation(builder)
    column_name = f'{_path_to_sql_column_name(reference_type_path)}_{_key_to_sql_column_name(constraint_key)}'
    parent_path = self._abs_path_invocation(builder.get_parent_builder())
    description = f'Reference type {reference_type_path} links to multiple resources or to resources of a type restricted by the profile.'
    return [validation_pb2.SqlRequirement(column_name=column_name, sql_expression=constraint_sql.sql, severity=validation_pb2.ValidationSeverity.SEVERITY_ERROR, type=validation_pb2.ValidationType.VALIDATION_TYPE_REFERENCE_TYPE, element_path=parent_path, description=description, fhir_path_key=constraint_key, fhir_path_expression=constraint_sql.builder.fhir_path, fields_referenced_by_expression=[field_name])]

Generates constraints for reference types.

    Ensures that a reference type only has a value for one of the resourceId
    columns across each of the possible resources the reference can link.

Args:
        builder: The builder to the reference type for which to encode
        constraints.
        elem: Element definition of the builder.

Returns:
        A constraint enforcing the above requirements for the given reference
        type.


def expand(self, pcoll: beam.PCollection[Chunk]) -> beam.PTransform[Chunk, Any]:
    write_transform = self.database_config.create_write_transform()
    return pcoll | write_transform

Creates and applies the database-specific write transform.

Args:
        pcoll: PCollection of Chunks with embeddings to write to the
            vector database. Each Chunk must have:
            - An embedding
            - An ID
            - Metadata used to filter results as specified by database config

Returns:
        Result of writing to database (implementation specific).


def __init__(self, dtype, ragged_rank, row_splits_dtype=dtypes.int64):
    row_splits_dtype = dtypes.as_dtype(row_splits_dtype)
    self._dtype = dtype
    self._ragged_rank = ragged_rank
    self._row_splits_dtype = row_splits_dtype

Initializes a RaggedTensorType object.

Args:
        dtype: data type of the `RaggedTensor`'s inner values.
        ragged_rank: ragged_rank of the declared `RaggedTensor`.
        row_splits_dtype: data type for the `RaggedTensor`'s row splits.
        One of: `tf.int32` or `tf.int64`.


def emit_pid(self, name: str, pid: int) -> None:
    event = {}
    event['name'] = 'process_name'
    event['ph'] = 'M'
    event['pid'] = pid
    event['args'] = {'name': name}
    self._metadata.append(event)

Adds a process metadata event to the trace.

Args:
        name:  The process name as a string.
        pid:  Identifier of the process as an integer.


def check_missing_deps() -> Iterator[None]:
    try:
        yield
    except ImportError as e:
        reraise_utils.reraise(e, suffix='\nEach etils sub-modules require deps to be installed separately (e.g. `from etils import ecolab` -> `pip install etils[ecolab]`)')

Raise a better error message in case of `ImportError`.

    Usage:

    ```python
    from etils.epy import _internal

    with _internal.check_missing_deps():
        # pylint: disable=g-import-not-at-top
        import xyz
        # pylint: enable=g-import-not-at-top
    ```

Yields:
        None


def dummy_inputs(self) -> Dict[str, tf.Tensor]:
    dummies = {}
    for key, spec in self.input_signature.items():
        dummy_shape = [dim if dim is not None else 2 for dim in spec.shape]
        if spec.shape[0] is None:
            dummy_shape[0] = 1
        dummies[key] = tf.ones(shape=dummy_shape, dtype=spec.dtype)
        if key == 'token_type_ids':
            dummies[key] = tf.zeros_like(dummies[key])
    if self.config.add_cross_attention and 'encoder_hidden_states' in inspect.signature(self.call).parameters:
        if 'encoder_hidden_states' not in dummies:
            if self.main_input_name == 'input_ids':
                dummies['encoder_hidden_states'] = tf.ones(shape=(1, 2, self.config.hidden_size), dtype=tf.float32, name='encoder_hidden_states')
            else:
                raise NotImplementedError("Model has cross-attention but we couldn't infer the shape for the encoder hidden states. Please manually override dummy_inputs!")
    return dummies

Dummy inputs to build the network.

Returns:
        `Dict[str, tf.Tensor]`: The dummy inputs.


def create_accumulator(self, *args, **kwargs):
    raise NotImplementedError(str(self))

Return a fresh, empty accumulator for the combine operation.

Args:
        *args: Additional arguments and side inputs.
        **kwargs: Additional arguments and side inputs.


def GetMap(self, cache_info, data):
    for obj in json.loads(cache_info.read()):
        key = obj.get('Key', '')
        value = obj.get('Value', '')
        if not value or not key:
            continue
        map_entry = self._ReadEntry(key, value)
        if map_entry is None:
            self.log.warning('Could not create entry from line %r in cache, skipping', value)
            continue
        if not data.Add(map_entry):
            self.log.warning('Could not add entry %r read from line %r in cache', map_entry, value)
    return data

Returns a map from a cache.

Args:
        cache_info: file like object containing the cache.
        data: a Map to populate.

Returns:
        A child of Map containing the cache data.


def get_resize_output_image_size(video, resolution_max_side: int) -> tuple[int, int]:
    height, width = video.size()[-2:]
    resolution_max_side = min(MAX_IMAGE_SIZE, resolution_max_side)
    resolution_max_side = max(height, width) if resolution_max_side is None else resolution_max_side
    aspect_ratio = width / height
    if width >= height:
        width = resolution_max_side
        height = int(width / aspect_ratio)
        if height % 2 != 0:
            height += 1
    elif height > width:
        height = resolution_max_side
        width = int(height * aspect_ratio)
        if width % 2 != 0:
            width += 1
    height = max(height, 1)
    width = max(width, 1)
    return (height, width)

Get the output size of the video after resizing given a dictionary specifying the max and min sizes.

Args:
        video (`np.ndarray`):
            Video to resize.
        resolution_max_side (`int`):
            The longest edge of the video will be resized to this value. The shortest edge will be resized to keep the
            input aspect ratio.

Returns:
        The output size of the video after resizing.


def get_temp_dir(self) -> str:
    if not self._tempdir:
        self._tempdir = tempfile.mkdtemp(dir=googletest.GetTempDir())
    return self._tempdir

Returns a unique temporary directory for the test to use.

    If you call this method multiple times during in a test, it will return the
    same folder. However, across different runs the directories will be
    different. This will ensure that across different runs tests will not be
    able to pollute each others environment.
    If you need multiple unique directories within a single test, you should
    use tempfile.mkdtemp as follows:
        tempfile.mkdtemp(dir=self.get_temp_dir()):

Returns:
        string, the path to the unique temporary directory created for this test.


def _merge_precomputed_encodings(self, other, validate=True):
    if self is other or (self._row_splits is other._row_splits and self._row_lengths is other._row_lengths and (self._value_rowids is other._value_rowids) and (self._nrows is other._nrows) and (self._nvals is other._nvals) and (self._uniform_row_length is other._uniform_row_length)):
        return self
    nrows, nrows_validated = _merge_tensors(self._nrows, other._nrows, 'nrows', validate)
    nvals, _ = _merge_tensors(self._nvals, other._nvals, 'nvals', validate)
    uniform_row_length, uniform_row_length_validated = _merge_tensors(self._uniform_row_length, other._uniform_row_length, 'uniform_row_length', validate)
    if uniform_row_length_validated and nrows_validated:
        validate = False
    row_splits, row_splits_validated = _merge_tensors(self._row_splits, other._row_splits, 'row_splits', validate)
    if row_splits_validated:
        validate = False
    row_lengths, row_lengths_validated = _merge_tensors(self._row_lengths, other._row_lengths, 'row_lengths', validate)
    if row_lengths_validated:
        validate = False
    value_rowids, value_rowids_validated = _merge_tensors(self._value_rowids, other._value_rowids, 'value_rowids', validate)
    if value_rowids_validated and nrows_validated:
        validate = False
    if row_splits is self._row_splits and row_lengths is self._row_lengths and (value_rowids is self._value_rowids) and (nrows is self._nrows) and (uniform_row_length is self._uniform_row_length):
        return self
    if row_splits is other._row_splits and row_lengths is other._row_lengths and (value_rowids is other._value_rowids) and (nrows is other._nrows) and (uniform_row_length is other._uniform_row_length):
        return other
    return RowPartition(row_splits=row_splits, row_lengths=row_lengths, value_rowids=value_rowids, nrows=nrows, uniform_row_length=uniform_row_length, nvals=nvals, internal=_row_partition_factory_key)

Returns a RowPartition that merges encodings from `self` and `other`.

    Requires that `self` and `other` describe the same partition.

Args:
        other: A `RowPartition` that encodes the same partition as `self`.
        validate: If true, then add runtime checks to verify that `self` and
        `other` encode the same row partition.

Returns:
        A `RowPartition`.


def report_fetch(config, auth, account, report_id=None, name=None, timeout=60):
    if config.verbose:
        print('DCM REPORT FILE', report_id or name)
    if report_id is None:
        report = report_get(config, auth, account, name=name)
        if report is None:
            raise Exception('Report does not exist:', name)
        else:
            report_id = report['id']
    running = False
    while timeout >= 0:
        for file_json in report_files(config, auth, account, report_id):
            if file_json.get('status') in ('PROCESSING', 'QUEUED', None):
                if config.verbose:
                    print('REPORT PROCESSING WILL WAIT')
                running = True
                if timeout > 0:
                    break
            elif file_json.get('status') == 'REPORT_AVAILABLE':
                if config.verbose:
                    print('REPORT DONE')
                return file_json
        if not running:
            break
        if timeout > 0:
            if config.verbose:
                print('WAITING MINUTES', timeout)
            sleep(60)
        timeout -= 1
    if config.verbose:
        print('NO REPORT FILES')
    return running

Retrieves most recent DCM file JSON by name or ID, if in progress, waits for it to complete.

    Bulletproofing:
    https://developers.google.com/doubleclick-advertisers/v3.2/files/get

    Timeout is in minutes ( retries will happen at 1 minute interval, default
    total time is 60 minutes )

Args:
        * auth: (string) Either user or service.
        * account: (string) [account:advertiser@profile] token.
        * report_id: (int) ID of DCm report to fetch ( either or name ).
        * name: (string) Name of report to fetch ( either or report_id ).
        * timeout: (int) Minutes to wait for in progress report before giving up.

Returns:
        * Report JSON if report exists and is ready.
        * True if report is in progress but not ready.
        * False if report does not exist.


def find_all_build_files(dir: str) -> List[Tuple[str, str]]:
    build_file_dirs = []
    for root, _, files in os.walk(dir):
        for file in files:
            if file in BUILD_FILENAMES:
                root = root.strip('./')
                build_file_dirs.append((root, file))
    return build_file_dirs

List all the BUILD files.

Returns:
        The list of (directory, filename) of all BUILD files.


def __init__(self, min_batch_size: int, max_batch_size: int, *, inference_fn: TensorRTInferenceFn=_default_tensorRT_inference_fn, large_model: bool=False, model_copies: Optional[int]=None, max_batch_duration_secs: Optional[int]=None, **kwargs):
    self.min_batch_size = min_batch_size
    self.max_batch_size = max_batch_size
    self.max_batch_duration_secs = max_batch_duration_secs
    self.inference_fn = inference_fn
    if 'engine_path' in kwargs:
        self.engine_path = kwargs.get('engine_path')
    elif 'onnx_path' in kwargs:
        self.onnx_path = kwargs.get('onnx_path')
    self._env_vars = kwargs.get('env_vars', {})
    self._share_across_processes = large_model or model_copies is not None
    self._model_copies = model_copies or 1

Implementation of the ModelHandler interface for TensorRT.

    Example Usage::

        pcoll | RunInference(
        TensorRTEngineHandlerNumPy(
            min_batch_size=1,
            max_batch_size=1,
            engine_path="my_uri"))

    **NOTE:** This API and its implementation are under development and
    do not provide backward compatibility guarantees.

Args:
        min_batch_size: minimum accepted batch size.
        max_batch_size: maximum accepted batch size.
        inference_fn: the inference function to use on RunInference calls.
        default: _default_tensorRT_inference_fn
        large_model: set to true if your model is large enough to run into
        memory pressure if you load multiple copies. Given a model that
        consumes N memory and a machine with W cores and M memory, you should
        set this to True if N*W > M.
        model_copies: The exact number of models that you would like loaded
        onto your machine. This can be useful if you exactly know your CPU or
        GPU capacity and want to maximize resource utilization.
        max_batch_duration_secs: the maximum amount of time to buffer
        a batch before emitting; used in streaming contexts.
        kwargs: Additional arguments like 'engine_path' and 'onnx_path' are
        currently supported. 'env_vars' can be used to set environment variables
        before loading the model.

    See https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/
    for details


def __init__(self, variable_holder=None, **kwargs):
    self._variable_holder = variable_holder or VariableHolder(share_variables=True)
    name = kwargs.pop('name', 'wrapped_function_graph')
    collections = kwargs.pop('collections', {})
    self.graph = func_graph.FuncGraph(name, collections=collections, **kwargs)
    self._wrapped_function = WrappedFunction(self.graph, self._variable_holder)
    self._functions = {}

Class for wrapping multiple TF 1.X functions in a single graph.

    Maintains a dictionary mapping names to wrapped functions. See
    `tf.compat.v1.wrap_function` to learn more about wrapping V1 functions.

    Functions wrapped using this class have access to variables and collections
    created in other wrapped functions, using the standard TF 1.X API (
    `tf.compat.v1.get_variable` or
    `tf.compat.v1.get_default_graph().get_collection(...)`)

    Outside a function, variables and collections may be accessed using the
    `variables` and `graph` properties.

Example:
    ```
    def add_v1(x):
        with tf.compat.v1.variable_scope('vars', reuse=tf.compat.v1.AUTO_REUSE):
        v = tf.compat.v1.get_variable('v', shape=[], dtype=tf.int32)
        return v + x

    def increment_var_v1(x):
        with tf.compat.v1.variable_scope('vars', reuse=tf.compat.v1.AUTO_REUSE):
        v = tf.compat.v1.get_variable('v', shape=[], dtype=tf.int32)
        return v.assign_add(x)

    g = WrappedGraph()
    add = g.wrap_function(add_v1, [tf.TensorSpec([], tf.int32)])
    increment_var = g.wrap_function(increment_var_v1,
                                    [tf.TensorSpec([], tf.int32)])

    assert len(g.variables) == 1
    assert g.variables[0].numpy() == 0
    increment_var(tf.constant(5))
    assert g.variables[0].numpy() == 5

    ```


def remove_padding_from_sc(value_in_checkpoint: tensor.Tensor, variable_shape: tuple[int, int]) -> tensor.Tensor:
    checkpoint_value_shape = value_in_checkpoint.shape.as_list()
    is_init_value_padded = all([i >= j for i, j in zip(checkpoint_value_shape, variable_shape)])
    if not is_init_value_padded:
        return value_in_checkpoint
    begin = [0] * len(checkpoint_value_shape)
    return array_ops.slice(value_in_checkpoint, begin=begin, size=variable_shape)

Removes padding, if any, from sparsecore checkpoint.

Args:
        value_in_checkpoint: input tensor value, usually from checkpoint.
        variable_shape: Expected shape of tensor after removing padding.

Returns:
        A slice of the input tensor to match the variable_shape if the
        variable shape is a valid slice if the input tensor.


def __ne__(self: EventSetOrNode, other: Any) -> EventSetOrNode:
    if isinstance(other, self.__class__):
        from temporian.core.operators.binary import not_equal
        return not_equal(input_1=self, input_2=other)
    if isinstance(other, T_SCALAR + (bool, str)):
        from temporian.core.operators.scalar import not_equal_scalar
        return not_equal_scalar(input=self, value=other)
    self._raise_error('ne', other, 'int,float,bool,str')
    assert False

Computes not equal (`self != other`) element-wise with another
    [`EventSet`][temporian.EventSet] or a scalar value.

    If an EventSet, each feature in `self` is compared element-wise to
    the feature in `other` in the same position. `self` and `other`
    must have the same sampling and the same number of features.

    If a scalar value, each item in each feature in `self` is compared to
    `other`.

    Note that it will always return True on NaNs (even if both are).

    Example with EventSet:
        ```python
        >>> a = tp.event_set(
        ...     timestamps=[1, 2, 3],
        ...     features={"f1": [0, 100, 200]}
        ... )
        >>> b = tp.event_set(
        ...     timestamps=[1, 2, 3],
        ...     features={"f2": [-10, 100, 5]},
        ...     same_sampling_as=a
        ... )

        >>> c = a != b
        >>> c
        indexes: []
        features: [('f1', bool_)]
        events:
            (3 events):
                timestamps: [1. 2. 3.]
                'f1': [ True False True]
        ...

        ```

    Example with scalar value:
        ```python
        >>> a = tp.event_set(
        ...     timestamps=[1, 2, 3],
        ...     features={"f1": [0, 100, 200], "f2": [-10, 100, 5]}
        ... )

        >>> b = a != 100
        >>> b
        indexes: []
        features: [('f1', bool_), ('f2', bool_)]
        events:
            (3 events):
                timestamps: [1. 2. 3.]
                'f1': [ True False True]
                'f2': [ True False True]
        ...

        ```

Args:
        other: EventSet or scalar value.

Returns:
        Result of the comparison.


def check_output(expected: List[str]):

    def _check_inner(actual: List[PCollection[str]]):
        formatted_actual = actual | beam.Flatten() | beam.Map(lambda row: str(beam.Row(**row._asdict())))
        assert_matches_stdout(formatted_actual, expected)
    return _check_inner

Helper function to check the output of a pipeline against expected values.

    This function takes a list of expected output strings and returns a
    callable that can be used within a Beam pipeline to assert that the
    actual output matches the expected output.

Args:
        expected: A list of strings representing the expected output elements.

Returns:
        A callable that takes a list of PCollections and asserts their combined
        elements match the expected output.


def self_adjoint_eigvals(tensor, name=None):
    e, _ = gen_linalg_ops.self_adjoint_eig_v2(tensor, compute_v=False, name=name)
    return e

Computes the eigenvalues of one or more self-adjoint matrices.

    Note: If your program backpropagates through this function, you should replace
    it with a call to tf.linalg.eigh (possibly ignoring the second output) to
    avoid computing the eigen decomposition twice. This is because the
    eigenvectors are used to compute the gradient w.r.t. the eigenvalues. See
    _SelfAdjointEigV2Grad in linalg_grad.py.

Args:
        tensor: `Tensor` of shape `[..., N, N]`.
        name: string, optional name of the operation.

Returns:
        e: Eigenvalues. Shape is `[..., N]`. The vector `e[..., :]` contains the `N`
        eigenvalues of `tensor[..., :, :]`.


def create_source_map(nodes, code, filepath):
    reparsed_nodes = parser.parse(code, preamble_len=0, single_node=False)
    for node in reparsed_nodes:
        resolve(node, code, filepath, node.lineno, node.col_offset)
    source_map = {}
    try:
        for before, after in ast_util.parallel_walk(nodes, reparsed_nodes):
            origin_info = anno.getanno(before, anno.Basic.ORIGIN, default=None)
            final_info = anno.getanno(after, anno.Basic.ORIGIN, default=None)
            if origin_info is None or final_info is None:
                continue
            line_loc = LineLocation(final_info.loc.filename, final_info.loc.lineno)
            existing_origin = source_map.get(line_loc)
            if existing_origin is not None:
                if existing_origin.loc.line_loc == origin_info.loc.line_loc:
                    if existing_origin.loc.lineno >= origin_info.loc.lineno:
                        continue
                if existing_origin.loc.col_offset <= origin_info.loc.col_offset:
                    continue
            source_map[line_loc] = origin_info
    except ValueError as err:
        new_msg = 'Inconsistent ASTs detected. This is a bug. Cause: \n'
        new_msg += str(err)
        new_msg += 'Diff:\n'
        for n, rn in zip(nodes, reparsed_nodes):
            nodes_str = pretty_printer.fmt(n, color=False, noanno=True)
            reparsed_nodes_str = pretty_printer.fmt(rn, color=False, noanno=True)
            diff = difflib.context_diff(nodes_str.split('\n'), reparsed_nodes_str.split('\n'), fromfile='Original nodes', tofile='Reparsed nodes', n=7)
            diff = '\n'.join(diff)
            new_msg += diff + '\n'
        raise ValueError(new_msg)
    return source_map

Creates a source map between an annotated AST and the code it compiles to.

    Note: this function assumes nodes nodes, code and filepath correspond to the
    same code.

Args:
        nodes: Iterable[ast.AST, ...], one or more AST modes.
        code: Text, the source code in which nodes are found.
        filepath: Text

Returns:
        Dict[LineLocation, OriginInfo], mapping locations in code to locations
        indicated by origin annotations in node.


def set_global_step(self, new_global_step, name=None):
    return gen_data_flow_ops.accumulator_set_global_step(self._accumulator_ref, math_ops.cast(ops.convert_to_tensor(new_global_step), _dtypes.int64), name=name)

Sets the global time step of the accumulator.

    The operation logs a warning if we attempt to set to a time step that is
    lower than the accumulator's own time step.

Args:
        new_global_step: Value of new time step. Can be a variable or a constant
        name: Optional name for the operation.

Returns:
        Operation that sets the accumulator's time step.


def _constant_value(ragged_factory, inner_factory, pylist, dtype, ragged_rank, inner_shape):
    if ragged_tensor.is_ragged(pylist):
        raise TypeError('pylist may not be a RaggedTensor or RaggedTensorValue.')
    if not isinstance(pylist, (list, tuple)) and np.ndim(pylist) == 0:
        if ragged_rank is not None and ragged_rank != 0:
            raise ValueError('Invalid pylist=%r: incompatible with ragged_rank=%d' % (pylist, ragged_rank))
        if inner_shape is not None and inner_shape:
            raise ValueError('Invalid pylist=%r: incompatible with dim(inner_shape)=%d' % (pylist, len(inner_shape)))
        return inner_factory(pylist, dtype, ())
    if ragged_rank is not None and ragged_rank < 0:
        raise ValueError('Invalid ragged_rank=%r: must be nonnegative' % ragged_rank)
    scalar_depth, max_depth = _find_scalar_and_max_depth(pylist)
    if scalar_depth is not None:
        if max_depth > scalar_depth:
            raise ValueError('Invalid pylist=%r: empty list nesting is greater than scalar value nesting' % pylist)
        if ragged_rank is not None and max_depth < ragged_rank:
            raise ValueError(f'Invalid pylist={pylist}, max depth smaller than ragged_rank={ragged_rank}')
    if inner_shape is not None and ragged_rank is not None:
        expected_depth = ragged_rank + len(inner_shape) + 1
        if scalar_depth is not None and expected_depth != scalar_depth or (scalar_depth is None and expected_depth < max_depth):
            raise ValueError('Invalid pylist=%r: incompatible with ragged_rank=%d and dim(inner_shape)=%d' % (pylist, ragged_rank, len(inner_shape)))
    if ragged_rank == 0 or (ragged_rank is None and (max_depth < 2 or (inner_shape is not None and max_depth - len(inner_shape) < 2))):
        return inner_factory(pylist, dtype, inner_shape)
    if inner_shape is None:
        if ragged_rank is None:
            inner_shape = ()
        else:
            inner_shape = _default_inner_shape_for_pylist(pylist, ragged_rank)
    if ragged_rank is None:
        if scalar_depth is None:
            ragged_rank = max(1, max_depth - 1)
        else:
            ragged_rank = max(1, scalar_depth - 1 - len(inner_shape))
    nested_splits = []
    values = pylist
    for dim in range(ragged_rank):
        nested_splits.append([0])
        concatenated_values = []
        for row in values:
            nested_splits[dim].append(nested_splits[dim][-1] + len(row))
            concatenated_values.extend(row)
        values = concatenated_values
    values = inner_factory(values, dtype=dtype, shape=(len(values),) + inner_shape, name='values')
    for row_splits in reversed(nested_splits):
        values = ragged_factory(values, row_splits)
    return values

Constructs a constant RaggedTensor or RaggedTensorValue.

Args:
        ragged_factory: A factory function with the signature:
        `ragged_factory(values, row_splits)`
        inner_factory: A factory function with the signature: `inner_factory(pylist,
        dtype, shape, name)`
        pylist: A nested `list`, `tuple` or `np.ndarray`.
        dtype: Data type for returned value.
        ragged_rank: Ragged rank for returned value.
        inner_shape: Inner value shape for returned value.

Returns:
        A value returned by `ragged_factory` or `inner_factory`.

Raises:
        ValueError: If the scalar values in `pylist` have inconsistent nesting
        depth; or if ragged_rank or inner_shape are incompatible with `pylist`.


def _estimate_initial_position(log_moneyness, total_variance):
    dtype = total_variance.dtype
    minvol_index = tf.argmin(total_variance, axis=1)
    m = tf.gather(log_moneyness, minvol_index, axis=1, batch_dims=1)
    sigma = 0.5 * tf.ones_like(minvol_index, dtype=dtype)
    rho = tf.zeros_like(minvol_index, dtype=dtype)
    y = total_variance
    x = tf.sqrt((log_moneyness - m[:, None]) ** 2 + sigma[:, None] ** 2)
    e_x = tf.math.reduce_mean(x, axis=1)
    e_y = tf.math.reduce_mean(y, axis=1)
    e_xy = tf.math.reduce_mean(x * y, axis=1)
    var_x = tf.math.reduce_variance(x, axis=1)
    b = (e_xy - e_x * e_y) / var_x
    a = e_y - b * e_x
    initial_position = tf.transpose([a, b, rho, m, sigma])
    return initial_position

Provides a heuristic initial guess for the SVI parameters.

    The values for `rho` and `sigma` are predetermined.
    `rho = 0` enforces the symmetry of the initial skew.
    `sigma = 0.5` enforces certain smoothness at the vertex of the skew.

    The value for `m` estimated as the position of the vertex of the skew:
    `m` is the log-moneyness corresponding to the smallest input variance.

    The values for `a`, `b` are computed using a simple linear regression, using
    the input data and the above estimates for `rho`, `sigma` and `m`.

Args:
        log_moneyness: A rank 2 real `Tensor` of shape [batch_size, num_strikes].
        The log-moneyness `k := log(K/F)` of the options.
        total_variance: A rank 2 real `Tensor` of shape [batch_size, num_strikes].
        The target total variance to be approximated by the SVI model.

Returns:
    A rank 2 real `Tensor` of shape [batch_size, 5], representing an initial
    guess for the SVI parameter optimization.


def _format_data_list_with_options(self, data_list):
    if self._options and self._options.experimental_replication_mode == InputReplicationMode.PER_REPLICA and (not self._options.experimental_fetch_to_device):
        return [data_list]
    else:
        return data_list

Change the data in to a list type if required.

    The OwnedMultiDeviceIterator returns the list data type,
    while the PER_REPLICA iterator (when used with prefetch disabled)
    returns without the enclosed list. This is to fix the inconsistency.

Args:
        data_list: data_list

Returns:
        list


def build_relative_position(query_size, key_size):
    q_ids = tf.range(query_size, dtype=tf.int32)
    k_ids = tf.range(key_size, dtype=tf.int32)
    rel_pos_ids = q_ids[:, None] - tf.tile(tf.reshape(k_ids, [1, -1]), [query_size, 1])
    rel_pos_ids = rel_pos_ids[:query_size, :]
    rel_pos_ids = tf.expand_dims(rel_pos_ids, axis=0)
    return tf.cast(rel_pos_ids, tf.int64)

Build relative position according to the query and key

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

Args:
        query_size (int): the length of query
        key_size (int): the length of key

Returns:
        `tf.Tensor`: A tensor with shape [1, query_size, key_size]


def fram_wave(waveform: np.array, hop_length: int=160, fft_window_size: int=400, center: bool=True):
    warnings.warn('The function `fram_wave` is deprecated and will be removed in version 4.31.0 of Transformers', FutureWarning)
    frames = []
    for i in range(0, waveform.shape[0] + 1, hop_length):
        if center:
            half_window = (fft_window_size - 1) // 2 + 1
            start = i - half_window if i > half_window else 0
            end = i + half_window if i < waveform.shape[0] - half_window else waveform.shape[0]
            frame = waveform[start:end]
            if start == 0:
                padd_width = (-i + half_window, 0)
                frame = np.pad(frame, pad_width=padd_width, mode='reflect')
            elif end == waveform.shape[0]:
                padd_width = (0, i - waveform.shape[0] + half_window)
                frame = np.pad(frame, pad_width=padd_width, mode='reflect')
        else:
            frame = waveform[i:i + fft_window_size]
            frame_width = frame.shape[0]
            if frame_width < waveform.shape[0]:
                frame = np.lib.pad(frame, pad_width=(0, fft_window_size - frame_width), mode='constant', constant_values=0)
        frames.append(frame)
    frames = np.stack(frames, 0)
    return frames

In order to compute the short time fourier transform, the waveform needs to be split in overlapping windowed
    segments called `frames`.

    The window length (window_length) defines how much of the signal is contained in each frame, while the hop length
    defines the step between the beginning of each new frame.

Args:
        waveform (`np.array` of shape `(sample_length,)`):
            The raw waveform which will be split into smaller chunks.
        hop_length (`int`, *optional*, defaults to 160):
            Step between each window of the waveform.
        fft_window_size (`int`, *optional*, defaults to 400):
            Defines the size of the window.
        center (`bool`, defaults to `True`):
            Whether or not to center each frame around the middle of the frame. Centering is done by reflecting the
            waveform on the left and on the right.

Returns:
        framed_waveform (`np.array` of shape `(waveform.shape // hop_length , fft_window_size)`):
            The framed waveforms that can be fed to `np.fft`.


def __init__(self, func, argnames, func_name=None, grad_func=None, python_grad_func=None, out_names=None, **kwargs):
    self._func = func
    self._argnames = argnames
    self._func_name = func_name
    assert grad_func is None or isinstance(grad_func, _OverloadedFunction)
    self._grad_func = grad_func
    self._python_grad_func = python_grad_func
    self._out_names = out_names
    self._extra_kwargs = kwargs
    self._overload = {}

Creates _DefinedFunction.

Args:
        func:  A python callable which constructs a tf function body.
        argnames: A list of strings for function argument names.
        func_name: The function name. Defaults to None, in which derives from
        'func'.
        grad_func: This function's gradient function, if not None. Defaults
        to None.
        python_grad_func: A python callable implementing the gradient of
        the function python-side.
        out_names: A list of strings for the function return value names.
        **kwargs: The keyword arguments. **kwargs is passed to every call
        site of this function.

Raises:
        ValueError: The function definition is invalid.


def __init__(self, checkpoint_dir: Text, save_secs: Optional[int]=None, save_steps: Optional[int]=None, saver: Optional[saver_lib.Saver]=None, checkpoint_basename: Text='model.ckpt', scaffold: Optional[monitored_session.Scaffold]=None, listeners: Optional[List[basic_session_run_hooks.CheckpointSaverListener]]=None):
    save_path = os.path.join(checkpoint_dir, checkpoint_basename)
    logging.info('Create AsyncCheckpointSaverHook saving to path\n%s', save_path)
    if listeners:
        logging.info(' with %d listener(s).', len(listeners))
    if saver is not None and scaffold is not None:
        raise ValueError('You cannot provide both saver and scaffold.')
    self._saver = saver
    self._save_thread = None
    self._write_graph_thread = None
    self._checkpoint_dir = checkpoint_dir
    self._save_path = save_path
    self._scaffold = scaffold
    self._timer = basic_session_run_hooks.SecondOrStepTimer(every_secs=save_secs, every_steps=save_steps)
    self._listeners = listeners or []
    self._steps_per_run = 1
    self._summary_writer = None
    self._global_step_tensor = None
    self._last_checkpoint_step = None
    global _END_TIME_OF_LAST_WRITE
    with _END_TIME_OF_LAST_WRITE_LOCK:
        if _END_TIME_OF_LAST_WRITE is None:
            _END_TIME_OF_LAST_WRITE = time.time()

Initializes a `CheckpointSaverHook`.

Args:
        checkpoint_dir: `str`, base directory for the checkpoint files.
        save_secs: `int`, save every N secs.
        save_steps: `int`, save every N steps.
        saver: `Saver` object, used for saving.
        checkpoint_basename: `str`, base name for the checkpoint files.
        scaffold: `Scaffold`, use to get saver object.
        listeners: List of `CheckpointSaverListener` subclass instances. Used for
        callbacks that run immediately before or after this hook saves the
        checkpoint.

Raises:
        ValueError: One of `save_steps` or `save_secs` should be set.
        ValueError: At most one of `saver` or `scaffold` should be set.


def _update(self, item, feed_item):
    self._api().update(profileId=self.profile_id, body=item).execute()

Updates a new item in CM.

Args:
        item: The CM object to update.
        feed_item: The feed item from the Bulkdozer feed representing the item to
        update.


def smart_cond(pred, true_fn=None, false_fn=None, name=None):
    if isinstance(pred, variables.Variable):
        return cond.cond(pred, true_fn=true_fn, false_fn=false_fn, name=name)
    return smart_module.smart_cond(pred, true_fn=true_fn, false_fn=false_fn, name=name)

Return either `true_fn()` if predicate `pred` is true else `false_fn()`.

    If `pred` is a bool or has a constant value, we return either `true_fn()`
    or `false_fn()`, otherwise we use `tf.cond` to dynamically route to both.

Args:
        pred: A scalar determining whether to return the result of `true_fn` or
        `false_fn`.
        true_fn: The callable to be performed if pred is true.
        false_fn: The callable to be performed if pred is false.
        name: Optional name prefix when using `tf.cond`.

Returns:
        Tensors returned by the call to either `true_fn` or `false_fn`.

Raises:
        TypeError: If `true_fn` or `false_fn` is not callable.


def _get_attrs_items(obj):
    attrs = getattr(obj.__class__, '__attrs_attrs__')
    attr_names = (a.name for a in attrs)
    return [(attr_name, getattr(obj, attr_name)) for attr_name in attr_names]

Returns a list of (name, value) pairs from an attrs instance.

    TODO(b/268078256): check if this comment is valid, and if so, ensure it's
    handled in the function below.
    The list will be sorted by name.

Args:
        obj: an object.

Returns:
        A list of (attr_name, attr_value) pairs, sorted by attr_name.


def get_next(self):
    raise NotImplementedError('Iterator.get_next()')

Returns the next element.

    >>> dataset = tf.data.Dataset.from_tensors(42)
    >>> iterator = iter(dataset)
    >>> print(iterator.get_next())
    tf.Tensor(42, shape=(), dtype=int32)

Returns:
        A (nested) structure of values matching `tf.data.Iterator.element_spec`.

Raises:
        `tf.errors.OutOfRangeError`: If the end of the iterator has been reached.


def register(self, candidate, name=None):
    if not name:
        name = candidate.__name__
    if name in self._registry:
        frame = self._registry[name][_LOCATION_TAG]
        raise KeyError("Registering two %s with name '%s'! (Previous registration was in %s %s:%d)" % (self._name, name, frame.name, frame.filename, frame.lineno))
    logging.vlog(1, 'Registering %s (%s) in %s.', name, candidate, self._name)
    stack = traceback.extract_stack(limit=3)
    stack_index = min(2, len(stack) - 1)
    if stack_index >= 0:
        location_tag = stack[stack_index]
    else:
        location_tag = ('UNKNOWN', 'UNKNOWN', 'UNKNOWN', 'UNKNOWN', 'UNKNOWN')
    self._registry[name] = {_TYPE_TAG: candidate, _LOCATION_TAG: location_tag}

Registers a Python object "candidate" for the given "name".

Args:
        candidate: The candidate object to add to the registry.
        name: An optional string specifying the registry key for the candidate.
            If None, candidate.__name__ will be used.

Raises:
        KeyError: If same name is used twice.


A function that is executed upon a test passing.

    Implementation is optional.

Args:
        record: records.TestResultRecord, a copy of the test record for
        this test, containing all information of the test execution
        including exception objects.


def freeze_saved_model(saved_model_dir, input_arrays, input_shapes, output_arrays, tag_set, signature_key):
    meta_graph = get_meta_graph_def(saved_model_dir, tag_set)
    signature_def = get_signature_def(meta_graph, signature_key)
    inputs, outputs = get_inputs_outputs(signature_def)
    collection_def = meta_graph.collection_def
    if constants.ASSETS_KEY in collection_def:
        raise ValueError('SavedModels with assets/ directory are not supported.')
    graph = ops.Graph()
    with session.Session(graph=graph) as sess:
        loader.load(sess, meta_graph.meta_info_def.tags, saved_model_dir)
        in_tensors = _get_tensors(graph, inputs, input_arrays)
        out_tensors = _get_tensors(graph, outputs, output_arrays)
        util.set_tensor_shapes(in_tensors, input_shapes)
        frozen_graph_def = util.freeze_graph(sess, in_tensors, out_tensors)
        return (frozen_graph_def, in_tensors, out_tensors, sess.graph)

Converts a SavedModel to a frozen graph.

Args:
        saved_model_dir: SavedModel directory to convert.
        input_arrays: List of input tensors to freeze graph with. Uses input arrays
        from SignatureDef when none are provided.
        input_shapes: Dict of strings representing input tensor names to list of
        integers representing input shapes (e.g., {"foo": : [1, 16, 16, 3]}).
        Automatically determined when input shapes is None (e.g., {"foo" : None}).
        output_arrays: List of output tensors to freeze graph with. Uses output
        arrays from SignatureDef when none are provided.
        tag_set: Set of tags identifying the MetaGraphDef within the SavedModel to
        analyze. All tags in the tag set must be present.
        signature_key: Key identifying SignatureDef containing inputs and outputs.

Returns:
        frozen_graph_def: Frozen GraphDef.
        in_tensors: List of input tensors for the graph.
        out_tensors: List of output tensors for the graph.
        graph: `Graph` object.

Raises:
        ValueError:
        SavedModel doesn't contain a MetaGraphDef identified by tag_set.
        signature_key is not in the MetaGraphDef.
        assets/ directory is in the MetaGraphDef.
        input_shapes does not match the length of input_arrays.
        input_arrays or output_arrays are not valid.


def words_string(fake: Faker, n: int) -> str:
    return ' '.join(fake.words(n))

Provide Faker words as a joined string.

Args:
        * fake: Faker instance
        * n: number of words

Returns:
        * string of n words joined by spaces


def __init__(self, rate, **kwargs):
    super(SpatialDropout1D, self).__init__(rate, **kwargs)
    self.input_spec = InputSpec(ndim=3)

Spatial 1D version of Dropout.

    This version performs the same function as Dropout, however, it drops
    entire 1D feature maps instead of individual elements. If adjacent frames
    within feature maps are strongly correlated (as is normally the case in
    early convolution layers) then regular dropout will not regularize the
    activations and will otherwise just result in an effective learning rate
    decrease. In this case, SpatialDropout1D will help promote independence
    between feature maps and should be used instead.

Args:
        rate: Float between 0 and 1. Fraction of the input units to drop.

    Call arguments:
        inputs: A 3D tensor.
        training: Python boolean indicating whether the layer should behave in
        training mode (adding dropout) or in inference mode (doing nothing).

    Input shape:
        3D tensor with shape:
        `(samples, timesteps, channels)`

    Output shape:
        Same as input.

    References:
        - [Efficient Object Localization Using Convolutional
        Networks](https://arxiv.org/abs/1411.4280)


def _parse_resource(self, uri: str, json_obj: Dict[str, Any]) -> Optional[_T]:
    json_parser = _json_parser.JsonParser(self.handler, self.resource_time_zone)
    resource_type = json_obj.get('resourceType')
    if resource_type is None:
        raise ValueError(f'JSON for URI {uri} does not have a resource type.')
    if resource_type == 'Bundle':
        json_value = _find_resource_in_bundle(uri, json_obj)
        if json_value is None:
            return None
        else:
            target = self.proto_cls()
            json_parser.merge_value(json_value, target)
            return target
    else:
        target = self.proto_cls()
        json_parser.merge_value(json_obj, target)
        return target

Parses a protocol buffer for the given JSON object.

Args:
        uri: The URI of the resource to parse.
        json_obj: The JSON object to parse into a proto.

Returns:
        The protocol buffer for the resource or `None` if it can not be found.