def _check_tf1_flags(flags, unparsed):

    def _get_message_unparsed(flag, orig_flag, new_flag):
        if flag.startswith(orig_flag):
            return '\n  Use {0} instead of {1}'.format(new_flag, orig_flag)
        return ''
    if unparsed:
        output = ''
        for flag in unparsed:
            output += _get_message_unparsed(flag, '--input_file', '--graph_def_file')
            output += _get_message_unparsed(flag, '--savedmodel_directory', '--saved_model_dir')
            output += _get_message_unparsed(flag, '--std_value', '--std_dev_values')
            output += _get_message_unparsed(flag, '--batch_size', '--input_shapes')
            output += _get_message_unparsed(flag, '--dump_graphviz', '--dump_graphviz_dir')
        if output:
            raise ValueError(output)
    if flags.graph_def_file and (not flags.input_arrays or not flags.output_arrays):
        raise ValueError('--input_arrays and --output_arrays are required with --graph_def_file')
    if flags.input_shapes:
        if not flags.input_arrays:
            raise ValueError('--input_shapes must be used with --input_arrays')
        if flags.input_shapes.count(':') != flags.input_arrays.count(','):
            raise ValueError('--input_shapes and --input_arrays must have the same number of items')
    if flags.std_dev_values or flags.mean_values:
        if bool(flags.std_dev_values) != bool(flags.mean_values):
            raise ValueError('--std_dev_values and --mean_values must be used together')
        if flags.std_dev_values.count(',') != flags.mean_values.count(','):
            raise ValueError('--std_dev_values, --mean_values must have the same number of items')
    if (flags.default_ranges_min is None) != (flags.default_ranges_max is None):
        raise ValueError('--default_ranges_min and --default_ranges_max must be used together')
    if flags.dump_graphviz_video and (not flags.dump_graphviz_dir):
        raise ValueError('--dump_graphviz_video must be used with --dump_graphviz_dir')
    if flags.custom_opdefs and (not flags.experimental_new_converter):
        raise ValueError('--custom_opdefs must be used with --experimental_new_converter')
    if flags.custom_opdefs and (not flags.allow_custom_ops):
        raise ValueError('--custom_opdefs must be used with --allow_custom_ops')
    if flags.experimental_select_user_tf_ops and (not flags.experimental_new_converter):
        raise ValueError('--experimental_select_user_tf_ops must be used with --experimental_new_converter')

Checks the parsed and unparsed flags to ensure they are valid in 1.X.

    Raises an error if previously support unparsed flags are found. Raises an
    error for parsed flags that don't meet the required conditions.

Args:
        flags: argparse.Namespace object containing TFLite flags.
        unparsed: List of unparsed flags.

Raises:
        ValueError: Invalid flags.


def uses_star_kwargs_in_call(node):
    if sys.version_info[:2] >= (3, 5):
        for keyword in node.keywords:
            if keyword.arg is None:
                return True
    elif node.kwargs:
        return True
    return False

Check if an ast.Call node uses arbitrary-length **kwargs.

    This function works with the AST call node format of Python3.5+
    as well as the different AST format of earlier versions of Python.

Args:
        node: The ast.Call node to check arg values for.

Returns:
        True if the node uses starred variadic positional args or keyword args.
        False if it does not.


def _HandleDuplicates(self, new_aliases):
    name_to_alias = {}
    out = []
    for a in new_aliases:
        if a.name not in name_to_alias:
            name_to_alias[a.name] = a
            out.append(a)
            continue
        existing = name_to_alias[a.name]
        if self._EquivalentAliases(existing, a):
            continue
        existing_name = existing.type.name or existing.type.__class__.__name__
        a_name = a.type.name or a.type.__class__.__name__
        raise KeyError(f'Duplicate top level items: {existing_name!r}, {a_name!r}')
    return out

Handle duplicate module-level aliases.

    Aliases pointing to qualified names could be the result of importing the
    same entity through multiple import paths, which should not count as an
    error; instead we just deduplicate them.

Args:
        new_aliases: The list of new aliases to deduplicate

Returns:
        A deduplicated list of aliases.

Raises:
        KeyError: If there is a name clash.


def remove_undocumented(module_name, allowed_exception_list=None, doc_string_modules=None):
    current_symbols = set(dir(_sys.modules[module_name]))
    should_have = make_all(module_name, doc_string_modules)
    should_have += allowed_exception_list or []
    extra_symbols = current_symbols - set(should_have)
    target_module = _sys.modules[module_name]
    for extra_symbol in extra_symbols:
        if extra_symbol.startswith('_'):
            continue
        fully_qualified_name = module_name + '.' + extra_symbol
        _HIDDEN_ATTRIBUTES[fully_qualified_name] = (target_module, getattr(target_module, extra_symbol))
        delattr(target_module, extra_symbol)

Removes symbols in a module that are not referenced by a docstring.

Args:
        module_name: the name of the module (usually `__name__`).
        allowed_exception_list: a list of names that should not be removed.
        doc_string_modules: a list of modules from which to take the docstrings.
        If None, then a list containing only the module named `module_name` is used.

        Furthermore, if a symbol previously added with `add_to_global_allowlist`,
        then it will always be allowed. This is useful for internal tests.

Returns:
        None


def __init__(self, loss=None, predictions=None, metrics=None):
    if loss is not None:
        loss_dict = self._wrap_and_check_outputs(loss, self.LOSS_NAME)
        self._loss = self._prefix_output_keys(loss_dict, self.LOSS_NAME)
    if predictions is not None:
        pred_dict = self._wrap_and_check_outputs(predictions, self.PREDICTIONS_NAME)
        self._predictions = self._prefix_output_keys(pred_dict, self.PREDICTIONS_NAME)
    if metrics is not None:
        self._metrics = self._wrap_and_check_metrics(metrics)

Constructor for SupervisedOutput (ie, Train or Eval output).

Args:
        loss: dict of Tensors or single Tensor representing calculated loss.
        predictions: dict of Tensors or single Tensor representing model
        predictions.
        metrics: Dict of metric results keyed by name.
        The values of the dict can be one of the following:
        (1) instance of `Metric` class.
        (2) (metric_value, update_op) tuples, or a single tuple.
        metric_value must be a Tensor, and update_op must be a Tensor or Op.

Raises:
        ValueError: if any of the outputs' dict keys are not strings or tuples of
        strings or the values are not Tensors (or Operations in the case of
        update_op).


def fuse_awq_modules(model, quantization_config):
    if isinstance(quantization_config, dict):
        quantization_config = AwqConfig.from_dict(quantization_config)
    backend = quantization_config.backend
    modules_to_fuse = get_modules_to_fuse(model, quantization_config)
    modules_to_not_convert = getattr(quantization_config, 'modules_to_not_convert', None)
    if backend == AwqBackendPackingMethod.AUTOAWQ:
        from awq.modules.fused.attn import QuantAttentionFused
        from awq.modules.fused.mlp import QuantFusedMLP
        from awq.modules.fused.norm import FasterTransformerRMSNorm
    else:
        raise ValueError('Fusing is only supported for the AutoAWQ backend')
    fused_attention_modules = []
    for name, module in model.named_modules():
        if modules_to_not_convert is not None:
            if any((module_name_to_not_convert in name for module_name_to_not_convert in modules_to_not_convert)):
                continue
        _fuse_awq_layernorm(modules_to_fuse['layernorm'], module, FasterTransformerRMSNorm)
        if quantization_config.version != 'ipex':
            _fuse_awq_mlp(model, name, modules_to_fuse['mlp'], module, QuantFusedMLP)
        else:
            logger.info('The IPEX version AWQ does not support fuse mlp for now.')
        attention_has_been_fused = _fuse_awq_attention_layers(model, module, modules_to_fuse, name, QuantAttentionFused)
        if attention_has_been_fused:
            fused_attention_modules.append(name.split('.')[0])
    if len(fused_attention_modules) > 0:
        for module_name, module in model.named_modules():
            if any((module_name in fused_attention_modules for fused_attention_parent_module in fused_attention_modules)):
                if hasattr(module, 'config') and hasattr(module.config, '_attn_implementation'):
                    module.config._attn_implementation = 'custom'
    return model

Optionally fuse some modules in the model to speedup inference.

Args:
        model (`~PreTrainedModel`):
            The model to fuse - note this model should have been converted into AWQ format beforehand.
        quantization_config (`Union[AwqConfig, dict]`):
            The quantization configuration to use.


def post_process_pose_estimation(self, outputs: 'VitPoseEstimatorOutput', boxes: Union[List[List[List[float]]], np.ndarray], kernel_size: int=11, threshold: Optional[float]=None, target_sizes: Union[TensorType, List[Tuple]]=None):
    batch_size, num_keypoints, _, _ = outputs.heatmaps.shape
    if target_sizes is not None:
        if batch_size != len(target_sizes):
            raise ValueError('Make sure that you pass in as many target sizes as the batch dimension of the logits')
    centers = np.zeros((batch_size, 2), dtype=np.float32)
    scales = np.zeros((batch_size, 2), dtype=np.float32)
    flattened_boxes = list(itertools.chain(*boxes))
    for i in range(batch_size):
        if target_sizes is not None:
            image_width, image_height = (target_sizes[i][0], target_sizes[i][1])
            scale_factor = np.array([image_width, image_height, image_width, image_height])
            flattened_boxes[i] = flattened_boxes[i] * scale_factor
        width, height = (self.size['width'], self.size['height'])
        center, scale = box_to_center_and_scale(flattened_boxes[i], image_width=width, image_height=height)
        centers[i, :] = center
        scales[i, :] = scale
    preds, scores = self.keypoints_from_heatmaps(outputs.heatmaps.cpu().numpy(), centers, scales, kernel=kernel_size)
    all_boxes = np.zeros((batch_size, 4), dtype=np.float32)
    all_boxes[:, 0:2] = centers[:, 0:2]
    all_boxes[:, 2:4] = scales[:, 0:2]
    poses = torch.tensor(preds)
    scores = torch.tensor(scores)
    labels = torch.arange(0, num_keypoints)
    bboxes_xyxy = torch.tensor(coco_to_pascal_voc(all_boxes))
    results: List[List[Dict[str, torch.Tensor]]] = []
    pose_bbox_pairs = zip(poses, scores, bboxes_xyxy)
    for image_bboxes in boxes:
        image_results: List[Dict[str, torch.Tensor]] = []
        for _ in image_bboxes:
            pose, score, bbox_xyxy = next(pose_bbox_pairs)
            score = score.squeeze()
            keypoints_labels = labels
            if threshold is not None:
                keep = score > threshold
                pose = pose[keep]
                score = score[keep]
                keypoints_labels = keypoints_labels[keep]
            pose_result = {'keypoints': pose, 'scores': score, 'labels': keypoints_labels, 'bbox': bbox_xyxy}
            image_results.append(pose_result)
        results.append(image_results)
    return results

Transform the heatmaps into keypoint predictions and transform them back to the image.

Args:
        outputs (`VitPoseEstimatorOutput`):
            VitPoseForPoseEstimation model outputs.
        boxes (`List[List[List[float]]]` or `np.ndarray`):
            List or array of bounding boxes for each image. Each box should be a list of 4 floats representing the bounding
            box coordinates in COCO format (top_left_x, top_left_y, width, height).
        kernel_size (`int`, *optional*, defaults to 11):
            Gaussian kernel size (K) for modulation.
        threshold (`float`, *optional*, defaults to None):
            Score threshold to keep object detection predictions.
        target_sizes (`torch.Tensor` or `List[Tuple[int, int]]`, *optional*):
            Tensor of shape `(batch_size, 2)` or list of tuples (`Tuple[int, int]`) containing the target size
            `(height, width)` of each image in the batch. If unset, predictions will be resize with the default value.

Returns:
        `List[List[Dict]]`: A list of dictionaries, each dictionary containing the keypoints and boxes for an image
        in the batch as predicted by the model.


def run_suite_class(argv=None):
    cli_args = _parse_cli_args(argv)
    suite_class = _find_suite_class()
    if cli_args.list_tests:
        _print_test_names_for_suite(suite_class)
        sys.exit(0)
    test_configs = config_parser.load_test_config_file(cli_args.config, cli_args.test_bed)
    config_count = len(test_configs)
    if config_count != 1:
        logging.error('Expect exactly one test config, found %d', config_count)
    config = test_configs[0]
    runner = test_runner.TestRunner(log_dir=config.log_path, testbed_name=config.testbed_name)
    suite = suite_class(runner, config)
    test_selector = _parse_raw_test_selector(cli_args.tests)
    suite.set_test_selector(test_selector)
    suite_record = SuiteInfoRecord(suite_class_name=suite_class.__name__)
    console_level = logging.DEBUG if cli_args.verbose else logging.INFO
    ok = False
    with runner.mobly_logger(console_level=console_level) as log_path:
        try:
            suite.setup_suite(config.copy())
            try:
                suite_record.suite_begin()
                runner.run()
                ok = runner.results.is_all_pass
                print(ok)
            except signals.TestAbortAll:
                pass
        finally:
            suite.teardown_suite()
            suite_record.suite_end()
            suite_record.suite_run_display_name = suite.get_suite_run_display_name()
            suite_record.extras = suite.get_suite_info().copy()
            _dump_suite_info(suite_record, log_path)
    if not ok:
        sys.exit(1)

Executes tests in the test suite.

Args:
        argv: A list that is then parsed as CLI args. If None, defaults to sys.argv.


def register_keras_serializable(package='Custom', name=None):

    def decorator(arg):
        """Registers a class with the Keras serialization framework."""
        class_name = name if name is not None else arg.__name__
        registered_name = package + '>' + class_name
        if tf_inspect.isclass(arg) and (not hasattr(arg, 'get_config')):
            raise ValueError('Cannot register a class that does not have a get_config() method.')
        if registered_name in _GLOBAL_CUSTOM_OBJECTS:
            raise ValueError('%s has already been registered to %s' % (registered_name, _GLOBAL_CUSTOM_OBJECTS[registered_name]))
        if arg in _GLOBAL_CUSTOM_NAMES:
            raise ValueError('%s has already been registered to %s' % (arg, _GLOBAL_CUSTOM_NAMES[arg]))
        _GLOBAL_CUSTOM_OBJECTS[registered_name] = arg
        _GLOBAL_CUSTOM_NAMES[arg] = registered_name
        return arg
    return decorator

Registers an object with the Keras serialization framework.

    This decorator injects the decorated class or function into the Keras custom
    object dictionary, so that it can be serialized and deserialized without
    needing an entry in the user-provided custom object dict. It also injects a
    function that Keras will call to get the object's serializable string key.

    Note that to be serialized and deserialized, classes must implement the
    `get_config()` method. Functions do not have this requirement.

    The object will be registered under the key 'package>name' where `name`,
    defaults to the object name if not passed.

Args:
        package: The package that this class belongs to.
        name: The name to serialize this class under in this package. If None, the
        class' name will be used.

Returns:
        A decorator that registers the decorated class with the passed names.


def flip(x, axis=None):
    if any_symbolic_tensors((x,)):
        return Flip(axis=axis).symbolic_call(x)
    return backend.numpy.flip(x, axis=axis)

Reverse the order of elements in the tensor along the given axis.

    The shape of the tensor is preserved, but the elements are reordered.

Args:
        x: Input tensor.
        axis: Axis or axes along which to flip the tensor. The default,
            `axis=None`, will flip over all of the axes of the input tensor.

Returns:
        Output tensor with entries of `axis` reversed.


def process_queries(self, text: Union[TextInput, List[TextInput]], **kwargs: Unpack[ColPaliProcessorKwargs]) -> BatchFeature:
    return self.__call__(text=text, **kwargs)

Prepare for the model one or several texts. This method is a wrapper around the `__call__` method of the ColPaliProcessor's
    [`ColPaliProcessor.__call__`].

    This method forwards the `text` and `kwargs` arguments to the tokenizer.

Args:
        text (`str`, `List[str]`, `List[List[str]]`):
            The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
            (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
            `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
        return_tensors (`str` or [`~utils.TensorType`], *optional*):
            If set, will return tensors of a particular framework. Acceptable values are:

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

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

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


def on_epoch_begin(self, epoch, logs=None):
    logs = self._process_logs(logs)
    for callback in self.callbacks:
        callback.on_epoch_begin(epoch, logs)

Calls the `on_epoch_begin` methods of its callbacks.

    This function should only be called during TRAIN mode.

Args:
        epoch: Integer, index of epoch.
        logs: Dict. Currently no data is passed to this argument for this method
        but that may change in the future.


def categorical_crossentropy(target, output, from_logits=False, axis=-1):
    target = tensor_conversion.convert_to_tensor_v2_with_dispatch(target)
    output = tensor_conversion.convert_to_tensor_v2_with_dispatch(output)
    target.shape.assert_is_compatible_with(output.shape)
    if hasattr(output, '_keras_logits'):
        output = output._keras_logits
        if from_logits:
            warnings.warn('"`categorical_crossentropy` received `from_logits=True`, but the `output` argument was produced by a sigmoid or softmax activation and thus does not represent logits. Was this intended?"')
        from_logits = True
    if from_logits:
        return nn.softmax_cross_entropy_with_logits_v2(labels=target, logits=output, axis=axis)
    if (not isinstance(output, (ops.EagerTensor, variables_module.Variable)) and output.op.type == 'Softmax') and (not hasattr(output, '_keras_history')):
        assert len(output.op.inputs) == 1
        output = output.op.inputs[0]
        return nn.softmax_cross_entropy_with_logits_v2(labels=target, logits=output, axis=axis)
    output = output / math_ops.reduce_sum(output, axis, True)
    epsilon_ = _constant_to_tensor(epsilon(), output.dtype.base_dtype)
    output = clip_ops.clip_by_value(output, epsilon_, 1.0 - epsilon_)
    return -math_ops.reduce_sum(target * math_ops.log(output), axis)

Categorical crossentropy between an output tensor and a target tensor.

Args:
        target: A tensor of the same shape as `output`.
        output: A tensor resulting from a softmax
            (unless `from_logits` is True, in which
            case `output` is expected to be the logits).
        from_logits: Boolean, whether `output` is the
            result of a softmax, or is a tensor of logits.
        axis: Int specifying the channels axis. `axis=-1` corresponds to data
            format `channels_last`, and `axis=1` corresponds to data format
            `channels_first`.

Returns:
        Output tensor.

Raises:
        ValueError: if `axis` is neither -1 nor one of the axes of `output`.

Example:
    >>> a = tf.constant([1., 0., 0., 0., 1., 0., 0., 0., 1.], shape=[3,3])
    >>> print(a)
    tf.Tensor(
        [[1. 0. 0.]
        [0. 1. 0.]
        [0. 0. 1.]], shape=(3, 3), dtype=float32)
    >>> b = tf.constant([.9, .05, .05, .05, .89, .06, .05, .01, .94], shape=[3,3])
    >>> print(b)
    tf.Tensor(
        [[0.9  0.05 0.05]
        [0.05 0.89 0.06]
        [0.05 0.01 0.94]], shape=(3, 3), dtype=float32)
    >>> loss = tf.keras.backend.categorical_crossentropy(a, b)
    >>> print(np.around(loss, 5))
    [0.10536 0.11653 0.06188]
    >>> loss = tf.keras.backend.categorical_crossentropy(a, a)
    >>> print(np.around(loss, 5))
    [0. 0. 0.]


def NeedsSeparatingHyphenHyphen(self, flag='help'):
    element = self.GetLastHealthyElement()
    component = element.component
    spec = inspectutils.GetFullArgSpec(component)
    return spec.varkw is not None or flag in spec.args or flag in spec.kwonlyargs

Returns whether a the trace need '--' before '--help'.

    '--' is needed when the component takes keyword arguments, when the value of
    flag matches one of the argument of the component, or the component takes in
    keyword-only arguments(e.g. argument with default value).

Args:
        flag: the flag available for the trace

Returns:
        True for needed '--', False otherwise.


def _gather(self, per_replica_value, destinations, axis, options=None):
    if isinstance(per_replica_value, indexed_slices.IndexedSlices):
        raise NotImplementedError('gather/all_gather does not support IndexedSlices')
    if options is None:
        options = collective_util.Options()
    per_replica_value = _make_tensor_into_per_replica(per_replica_value)
    validate_destinations(destinations)
    if self._num_between_graph_workers == 1 and len(per_replica_value.values) == 1 and _devices_match(per_replica_value, destinations, self._canonicalize_devices):
        with ops.device(per_replica_value.values[0].device):
            v = array_ops.identity(per_replica_value.values[0])
        return distribute_utils.regroup((v,), wrap_class=value_lib.Mirrored)
    return self._gather_implementation(per_replica_value, destinations, axis, options)

Gather `per_replica_value` to `destinations`.

Args:
        per_replica_value: a `tf.distribute.DistributedValues`, or a `tf.Tensor`
        like object.
        destinations: a `tf.distribute.DistributedValues`, a `tf.Variable`, a
        `tf.Tensor` alike object, or a device string. It specifies the devices
        to gather to. To perform an all-gather, pass the same to `value` and
        `destinations`. Note that if it's a `tf.Variable`, the value is gathered
        to the devices of that variable, and this method doesn't update the
        variable.
        axis: specifies the dimension to gather along within each replica's
        tensor.
        options: a `tf.distribute.experimental.CommunicationOptions`. See
        `tf.distribute.experimental.CommunicationOptions` for details.

Returns:
        A `tf.Tensor` or `tf.distribute.DistributedValues`

Raises:
        ValueError: if per_replica_value can't be converted to a
        `tf.distribute.DistributedValues` or if destinations is not a string,
        `tf.Variable` or `tf.distribute.DistributedValues`.


def _combine_named_parameters(**kwargs) -> list[OrderedDict[str, Any]]:
    sort_by_key = lambda k: k[0]
    combinations: list[list[tuple[str, Any]]] = []
    for key, values in sorted(kwargs.items(), key=sort_by_key):
        if not isinstance(values, list):
            values = [values]
        combinations.append([(key, value) for value in values])
    return [OrderedDict(result) for result in itertools.product(*combinations)]

Generate combinations based on its keyword arguments.

    Two sets of returned combinations can be concatenated using +.  Their product
    can be computed using `times()`.

Args:
        **kwargs: keyword arguments of form `option=[possibilities, ...]` or
        `option=the_only_possibility`.

Returns:
        a list of dictionaries for each combination. Keys in the dictionaries are
        the keyword argument names.  Each key has one value - one of the
        corresponding keyword argument values.


def unpad_image(tensor, original_size):
    if not isinstance(original_size, (list, tuple)):
        if not isinstance(original_size, (torch.Tensor, np.ndarray)):
            raise TypeError(f'image_size invalid type: {type(original_size)} not valid, should be either list, tuple, np.ndarray or tensor')
        original_size = original_size.tolist()
    original_height, original_width = original_size
    current_height, current_width = tensor.shape[1:]
    original_aspect_ratio = original_width / original_height
    current_aspect_ratio = current_width / current_height
    if original_aspect_ratio > current_aspect_ratio:
        scale_factor = current_width / original_width
        new_height = int(round(original_height * scale_factor, 7))
        padding = (current_height - new_height) // 2
        unpadded_tensor = tensor[:, padding:current_height - padding, :]
    else:
        scale_factor = current_height / original_height
        new_width = int(round(original_width * scale_factor, 7))
        padding = (current_width - new_width) // 2
        unpadded_tensor = tensor[:, :, padding:current_width - padding]
    return unpadded_tensor

Unpads a PyTorch tensor of a padded and resized image.

Args:
        tensor (`torch.Tensor`):
            The image tensor, assumed to be of shape (num_channels, height, width).
        original_size (`tuple`):
            The original size of the image (height, width).

Returns:
        `torch.Tensor`: The unpadded image tensor.


def create_analyzer_cli(dump):
    analyzer = analyzer_cli.DebugAnalyzer(dump, _cli_config_from_temp_file())
    registry = debugger_cli_common.CommandHandlerRegistry()
    registry.register_command_handler('list_tensors', analyzer.list_tensors, analyzer.get_help('list_tensors'), prefix_aliases=['lt'])
    registry.register_command_handler('node_info', analyzer.node_info, analyzer.get_help('node_info'), prefix_aliases=['ni'])
    registry.register_command_handler('list_inputs', analyzer.list_inputs, analyzer.get_help('list_inputs'), prefix_aliases=['li'])
    registry.register_command_handler('list_outputs', analyzer.list_outputs, analyzer.get_help('list_outputs'), prefix_aliases=['lo'])
    registry.register_command_handler('print_tensor', analyzer.print_tensor, analyzer.get_help('print_tensor'), prefix_aliases=['pt'])
    registry.register_command_handler('print_source', analyzer.print_source, analyzer.get_help('print_source'), prefix_aliases=['ps'])
    registry.register_command_handler('list_source', analyzer.list_source, analyzer.get_help('list_source'), prefix_aliases=['ls'])
    registry.register_command_handler('eval', analyzer.evaluate_expression, analyzer.get_help('eval'), prefix_aliases=['ev'])
    return (analyzer, registry)

Create an analyzer CLI.

Args:
        dump: A `DebugDumpDir` object to base the analyzer CLI on.

Returns:
        1) A `DebugAnalyzer` object created based on `dump`.
        2) A `CommandHandlerRegistry` that is based on the `DebugAnalyzer` object
        and has the common tfdbg commands, e.g., lt, ni, li, lo, registered.


def add_chunk(self, chunk: Union[message.Message, bytes], field_tags: util.FieldTypes, index=None) -> None:
    if self._parent_splitter is not None:
        self._parent_splitter.add_chunk(chunk, self._fields_in_parent + field_tags, index)
    else:
        assert self._chunks is not None
        assert self._chunked_message is not None
        field = self._chunked_message.chunked_fields.add(field_tag=util.get_field_tag(self._proto, field_tags))
        new_chunk_index = len(self._chunks)
        field.message.chunk_index = new_chunk_index
        self._add_chunk_order.append(id(chunk))
        if index is None:
            self._chunks.append(chunk)
        else:
            self._chunks.insert(index, chunk)
            self._fix_chunk_order = True

Adds a new chunk and updates the ChunkedMessage proto.

Args:
        chunk: Proto message or bytes.
        field_tags: Field information about the placement of the chunked data
        within self._proto.
        index: Optional index at which to insert the chunk. The chunk ordering is
        important for merging.


def square(x):
    if any_symbolic_tensors((x,)):
        return Square().symbolic_call(x)
    return backend.numpy.square(x)

Return the element-wise square of the input.

Args:
        x: Input tensor.

Returns:
        Output tensor, the square of `x`.


def report_delete(config, auth, account, report_id=None, name=None):
    report = report_get(config, auth, account, report_id, name)
    if report:
        account_id, advertiser_ids = parse_account(config, auth, account)
        is_superuser, profile_id = get_profile_for_api(config, auth, account_id)
        kwargs = {'profileId': profile_id, 'accountId': account_id} if is_superuser else {'profileId': profile_id}
        kwargs['reportId'] = report['id']
        API_DCM(config, auth, internal=is_superuser).reports().delete(**kwargs).execute()
    elif config.verbose:
        print('DCM DELETE: No Report')

Deletes a DCM report based on name or ID.

    Bulletproofing:
    https://developers.google.com/doubleclick-advertisers/v3.2/reports/delete

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

Returns:
        * None


def flush(self, force: bool=False) -> bool | Any:
    return self._messages.flush(force=force)

Flushes the underlying log message queue.

Args:
        * force: If True, force queue to flush

Returns:
        * True, if flushed with no errors
        * False, if not flushed
        * Error value from logger, if flushed with errors


def _check_instance_type(type_constraint, instance, var_name=None, verbose=False):
    hint_type = "argument: '%s'" % var_name if var_name is not None else 'return type'
    try:
        check_constraint(type_constraint, instance)
    except SimpleTypeHintError:
        if verbose:
            verbose_instance = '%s, ' % instance
        else:
            verbose_instance = ''
        raise TypeCheckError('Type-hint for %s violated. Expected an instance of %s, instead found %san instance of %s.' % (hint_type, type_constraint, verbose_instance, type(instance)))
    except CompositeTypeHintError as e:
        raise TypeCheckError('Type-hint for %s violated: %s' % (hint_type, e))

A helper function to report type-hint constraint violations.

Args:
        type_constraint: An instance of a 'TypeConstraint' or a built-in Python
        type.
        instance: The candidate object which will be checked by to satisfy
        'type_constraint'.
        var_name: If 'instance' is an argument, then the actual name for the
        parameter in the original function definition.

Raises:
        TypeCheckError: If 'instance' fails to meet the type-constraint of
        'type_constraint'.


def build_from_items(self, items: list[_ItemType] | None) -> imports_map.ImportsMap | None:
    if not items:
        return None
    imports_multimap = self._build_multimap(items)
    assert imports_multimap is not None
    return self._finalize(imports_multimap)

Create a file mapping from a list of (short path, path) tuples.

    Builds a dict of short_path to full name
        (e.g. "path/to/file.py" =>
            "$GENDIR/rulename~~pytype-gen/path_to_file.py~~pytype"

Args:
        items: A list of (short_path, full_path) tuples.

Returns:
        Dict of .py short_path to list of .pytd path or None if no items


def convertDateStrToDateTimeStr(date, time='00:00:00'):
    if not date == None:
        date = '%sT%sZ' % (date, time)
    return date

Convert Date string (YYYY-MM-DD) to a datetime string by adding the desired time (YYYY-MM-DDTHH:mm:SSZ)

Args:
            date: the date as a string to be converted
            time: the time as a string to be added to the date

Returns:
            A string representation of a datetime in the following
            format YYYY-MM-DDTHH:mm:SSZ


def verify_multiple_breaks(self, ds_fn, num_outputs, num_breaks=10, sparse_tensors=False, verify_exhausted=True, assert_items_equal=False):
    self.verify_run_with_breaks(ds_fn, self.gen_break_points(num_outputs, num_breaks), num_outputs, sparse_tensors=sparse_tensors, verify_exhausted=verify_exhausted, assert_items_equal=assert_items_equal)

Attempts to save/restore at multiple break points.

Args:
        ds_fn: 0-argument function that returns a Dataset.
        num_outputs: Total number of outputs expected from this Dataset.
        num_breaks: The number of break points. These are uniformly spread in [0,
        num_outputs] both inclusive.
        sparse_tensors: Whether dataset is built from SparseTensor(s).
        verify_exhausted: Whether to verify that the iterator has been exhausted
        after producing `num_outputs` elements.
        assert_items_equal: Tests the output has the expected elements regardless
        of order.

Raises:
        AssertionError if any test fails.


def _strip_debug_nodes(meta_graph_def: meta_graph_pb2.MetaGraphDef) -> None:

    def erase_regular_node_attributes(node: node_def_pb2.NodeDef) -> None:
        """Erases regular node attributes."""
        attributes_to_remove = [attribute for attribute in node.attr.keys() if not attribute.startswith('_')]
        for attribute in attributes_to_remove:
            node.attr.pop(attribute)

    def prune_all_non_t_attributes(node: node_def_pb2.NodeDef) -> None:
        """Prunes all attributes that are not `T`."""
        if 'T' in node.attr:
            t_value = node.attr['T']
            node.ClearField('attr')
            node.attr['T'].CopyFrom(t_value)
        else:
            node.ClearField('attr')

    def is_control_input(name: str) -> str:
        """Returns whether or not the input is a control input."""
        return name and name[0] == '^'

    def as_control_dep(name: str) -> str:
        """Returns the input as a control dependency."""
        return '^' + name.split(':')[0]

    def maybe_do_strip(node: node_def_pb2.NodeDef) -> None:
        """Strips the graph from Assert and CheckNumerics ops.

    For Assert ops, this function also rewrites all of the inputs to the nodes
    that were transformed by making them into control dependencies. It also
    removes all of the regular node attributes, that is all node attributes
    that do not start with `_`.

    For CheckNumerics ops, this function turns the op into an Identity op,
    which will be pruned later (according to the original implementation in
    grappler's `debug_stripper.cc`. Then, since Identity ops only take one
    input, it leaves the first input as is while transforming the other ones
    into control dependencies.

    Args:
      node: The node to potentally strip.
    """
        if node.op == 'Assert' or node.op == 'PrintV2':
            node.op = 'NoOp'
            erase_regular_node_attributes(node)
            new_inputs = []
            for inp in node.input:
                if not is_control_input(inp):
                    new_inputs.append(as_control_dep(inp))
                else:
                    new_inputs.append(inp)
            node.ClearField('input')
            node.input.extend(new_inputs)
        elif node.op == 'CheckNumerics' or node.op == 'Print':
            node.op = 'Identity'
            prune_all_non_t_attributes(node)
            for i in range(1, len(node.input)):
                if not is_control_input(node.input[i]):
                    node.input[i] = as_control_dep(node.input[i])
    for node in meta_graph_def.graph_def.node:
        maybe_do_strip(node)
    for func in meta_graph_def.graph_def.library.function:
        for node in func.node_def:
            maybe_do_strip(node)

An experimental function to remove debug nodes from the final graph.

    This function removes all Assert and CheckNumerics nodes from the meta_graph.
    It strips the operators in both the nodes and in all of the function defs,
    with the Assert ops being replaced by `NoOp`s and the CheckNumerics ops being
    transformed into `Identity` ops. In addition to this, it creates control
    inputs for the nodes that are not relevant for the op. For more information
    about control inputs please see go/how-tensors-flow#control-dependencies.

Args:
    meta_graph_def: The meta_graph that will be exported.


def partition_or_replicate_on_host(tensor, dims):
    if dims is None:
        return itertools.repeat(tensor)
    dims = np.array(dims)
    output = [tensor]
    shape_list = np.array(tensor.shape.as_list())
    quotients, remainders = np.divmod(shape_list, dims)
    for axis, (quotient, remainder, dim, original_size) in enumerate(zip(quotients, remainders, dims, shape_list)):
        if dim <= 1:
            continue
        if remainder > 0:
            ceil_ratio = quotient + 1
            num_full_slots, left_over = np.divmod(original_size, ceil_ratio)
            num_or_size_splits = [ceil_ratio] * num_full_slots + [left_over]
            if len(num_or_size_splits) < dim:
                num_or_size_splits += [0] * (dim - len(num_or_size_splits))
            new_output = []
            for x in output:
                new_output.append(array_ops.split(x, num_or_size_splits=num_or_size_splits, axis=axis))
            output = new_output
        else:
            output = [array_ops.split(x, int(dim), axis=axis) for x in output]
        output = nest.flatten(output)
    return output

Partitions or replicates the input tensor.

    The ops inside this function are placed on the host side.

Args:
    tensor: The input tensor which will be partitioned or replicated.
    dims: A list of integer describes how to partition the input tensor.

Returns:
    An iterator of `Tensor`s or a list of partitioned tensors.


def get_parent_dir_for_name(module_name):
    name_split = module_name.split('.')
    if not name_split:
        return None
    try:
        spec = importlib.util.find_spec(name_split[0])
    except ValueError:
        return None
    if not spec or not spec.origin:
        return None
    base_path = os.path.dirname(spec.origin)
    return os.path.join(base_path, *name_split[1:-1])

Get parent directory for module with the given name.

Args:
        module_name: Module name, e.g. tf_keras.api._v2.keras.

Returns:
        Path to the parent directory if module is found and None otherwise.
        Given example above, it should return:
        /root_path/tf_keras/api/_v2.


def _get_accumulator(tensor):
    assert isinstance(tensor.graph, func_graph_module.FuncGraph)

    def get_func_graph_output(t):
        """Returns t or Identity(t) whichever exists in graph outputs else None."""
        for output in tensor.graph.outputs:
            if output is t:
                return t
        identity_op = t.consumers()[0]
        if identity_op.type == 'Identity' and any((identity_op.outputs[0] is t for t in tensor.graph.outputs)):
            return identity_op.outputs[0]
        return None
    for consumer in tensor.consumers():
        if consumer.type != 'TensorListPushBack':
            continue
        accum_input_idx = -1
        for accum_input_idx, inp in enumerate(tensor.graph.inputs):
            if inp is consumer.inputs[0]:
                break
        else:
            continue
        output = get_func_graph_output(consumer.outputs[0])
        if output is None:
            continue
        for accum_output_idx, out in enumerate(tensor.graph.outputs):
            if out is output:
                if accum_input_idx == accum_output_idx:
                    return output
                break
    return None

Returns TensorList if any containing accumulated values of tensor.

    We try to find a pattern of the form:

        input_tl   tensor
        \        /
        (TensorListPushBack)
            |
        output_tl

    which satisfies the following conditions:

    1. input_tl must be in tensor.graph.inputs.
    2. output_tl or Identity(output_tl) must be in tensor.graph.outputs.
    3. tensor.graph.input_index(input_tl) == tensor.graph.output_index(output_t).

    output_tl or Identity(output_tl) (whichever is in tensor.graph.outputs) is
    returned if such a pattern is found else None is returned.

Args:
        tensor: The Tensor to be accumulated.

Returns:
        A variant tensor in the same graph as `tensor` or None if no accumulator is
        found.


def post_process_panoptic_sample(out_logits: np.ndarray, masks: np.ndarray, boxes: np.ndarray, processed_size: Tuple[int, int], target_size: Tuple[int, int], is_thing_map: Dict, threshold=0.85) -> Dict:
    scores, labels = score_labels_from_class_probabilities(out_logits)
    keep = (labels != out_logits.shape[-1] - 1) & (scores > threshold)
    cur_scores = scores[keep]
    cur_classes = labels[keep]
    cur_boxes = center_to_corners_format(boxes[keep])
    if len(cur_boxes) != len(cur_classes):
        raise ValueError('Not as many boxes as there are classes')
    cur_masks = masks[keep]
    cur_masks = resize(cur_masks[:, None], processed_size, resample=PILImageResampling.BILINEAR)
    cur_masks = safe_squeeze(cur_masks, 1)
    b, h, w = cur_masks.shape
    cur_masks = cur_masks.reshape(b, -1)
    stuff_equiv_classes = defaultdict(list)
    for k, label in enumerate(cur_classes):
        if not is_thing_map[label]:
            stuff_equiv_classes[label].append(k)
    seg_img = get_segmentation_image(cur_masks, processed_size, target_size, stuff_equiv_classes, deduplicate=True)
    area = get_mask_area(cur_masks, processed_size, n_classes=len(cur_scores))
    if cur_classes.size() > 0:
        filtered_small = np.array([a <= 4 for a in area], dtype=bool)
        while filtered_small.any():
            cur_masks = cur_masks[~filtered_small]
            cur_scores = cur_scores[~filtered_small]
            cur_classes = cur_classes[~filtered_small]
            seg_img = get_segmentation_image(cur_masks, (h, w), target_size, stuff_equiv_classes, deduplicate=True)
            area = get_mask_area(seg_img, target_size, n_classes=len(cur_scores))
            filtered_small = np.array([a <= 4 for a in area], dtype=bool)
    else:
        cur_classes = np.ones((1, 1), dtype=np.int64)
    segments_info = [{'id': i, 'isthing': is_thing_map[cat], 'category_id': int(cat), 'area': a} for i, (cat, a) in enumerate(zip(cur_classes, area))]
    del cur_classes
    with io.BytesIO() as out:
        PIL.Image.fromarray(seg_img).save(out, format='PNG')
        predictions = {'png_string': out.getvalue(), 'segments_info': segments_info}
    return predictions

Converts the output of [`DetrForSegmentation`] into panoptic segmentation predictions for a single sample.

Args:
        out_logits (`torch.Tensor`):
            The logits for this sample.
        masks (`torch.Tensor`):
            The predicted segmentation masks for this sample.
        boxes (`torch.Tensor`):
            The prediced bounding boxes for this sample. The boxes are in the normalized format `(center_x, center_y,
            width, height)` and values between `[0, 1]`, relative to the size the image (disregarding padding).
        processed_size (`Tuple[int, int]`):
            The processed size of the image `(height, width)`, as returned by the preprocessing step i.e. the size
            after data augmentation but before batching.
        target_size (`Tuple[int, int]`):
            The target size of the image, `(height, width)` corresponding to the requested final size of the
            prediction.
        is_thing_map (`Dict`):
            A dictionary mapping class indices to a boolean value indicating whether the class is a thing or not.
        threshold (`float`, *optional*, defaults to 0.85):
            The threshold used to binarize the segmentation masks.


def _gen_save_and_restore_functions(checkpoint_factory_map: object_identity.ObjectIdentityDictionary) -> object_identity.ObjectIdentityDictionary:
    saveable_fn_map = object_identity.ObjectIdentityDictionary()
    for obj, factory_data_list in checkpoint_factory_map.items():
        if resource_variable_ops.is_resource_variable(obj) or not factory_data_list:
            continue
        if factory_data_list[0].name == trackable_utils.SERIALIZE_TO_TENSORS_NAME:
            assert len(factory_data_list) == 1
            saveable_fn_map[obj] = {trackable_utils.SERIALIZE_TO_TENSORS_NAME: tracing_utils.trace_save_and_restore(obj)}
        else:
            saveable_fn_map[obj] = trace_saveable_util.trace_save_restore_function_map(obj, factory_data_list)
    return saveable_fn_map

Generates global and individual save/restore concrete functions.

    The global functions records the ops to save and restore the entire object to
    a file prefix, while the individual functions save and restore value tensors
    for resources.

    This function is intended to run on the output of
    `save_util_v1.get_checkpoint_factories_and_keys(object_names)`,
    which returns the generated a map of `_CheckpointFactoryData`.

Args:
        checkpoint_factory_map: A dictionary mapping trackable objects to
        a list of `_CheckpointFactoryData`.

Returns:
        Tuple of (
        saveable_fn_map: Maps obj -> factory name -> (concrete save, restore)
        )


def _ParseTensorName(tensor_name):
    components = tensor_name.split(':')
    if len(components) == 2:
        try:
            output_index = int(components[1])
        except ValueError:
            raise ValueError(f'Cannot convert {tensor_name!r} to a tensor name. Second component of the name following the `:` should be an int. Got {components[1]}.')
        return (components[0], output_index)
    elif len(components) == 1:
        return (components[0], 0)
    else:
        raise ValueError(f"Cannot convert '{tensor_name}' to a tensor name. Tensor names should not contain more than 1 `:`. Obtained {len(components) - 1}")

Parses a tensor name into an operation name and output index.

    This function will canonicalize tensor names as follows:

    * "foo:0"       -> ("foo", 0)
    * "foo:7"       -> ("foo", 7)
    * "foo"         -> ("foo", 0)
    * "foo:bar:baz" -> ValueError

Args:
        tensor_name: The name of a tensor.

Returns:
        A tuple containing the operation name, and the output index.

Raises:
        ValueError: If `tensor_name' cannot be interpreted as the name of a tensor.


def assertNumpyObjectTensorsRecursivelyEqual(self, a, b, msg):
    if isinstance(a, np.ndarray) and a.dtype == object:
        self.assertEqual(a.dtype, b.dtype, msg)
        self.assertEqual(a.shape, b.shape, msg)
        self.assertLen(a, len(b), msg)
        for a_val, b_val in zip(a, b):
            self.assertNumpyObjectTensorsRecursivelyEqual(a_val, b_val, msg)
    else:
        self.assertAllEqual(a, b, msg)

Check that two numpy arrays are equal.

    For arrays with dtype=object, check values recursively to see if a and b
    are equal.  (c.f. `np.array_equal`, which checks dtype=object values using
    object identity.)

Args:
        a: A numpy array.
        b: A numpy array.
        msg: Message to display if a != b.


def histogram(namespace: Union[Type, str], name: str, bucket_type: 'BucketType', logger: Optional['MetricLogger']=None) -> 'Metrics.DelegatingHistogram':
    namespace = UserMetrics.get_namespace(namespace)
    return Metrics.DelegatingHistogram(MetricName(namespace, name), bucket_type, logger)

Obtains or creates a Histogram metric.

Args:
        namespace: A class or string that gives the namespace to a metric
        name: A string that gives a unique name to a metric
        bucket_type: A type of bucket used in a histogram. A subclass of
        apache_beam.utils.histogram.BucketType
        logger: MetricLogger for logging locally aggregated metric

Returns:
        A Histogram object.


def map_subgraph(init_tensor, sources, disallowed_placeholders, visited_ops, op_outputs, add_sources):
    ops_to_visit = [_as_operation(init_tensor)]
    extra_sources = object_identity.ObjectIdentitySet()
    while ops_to_visit:
        op = ops_to_visit.pop()
        if op in visited_ops:
            continue
        visited_ops.add(op)
        should_raise = False
        if disallowed_placeholders is not None and op in disallowed_placeholders:
            should_raise = True
        elif op.type == 'Placeholder':
            if disallowed_placeholders is None and (not add_sources):
                should_raise = True
            extra_sources.update(op.outputs)
        if should_raise:
            raise UnliftableError('Unable to lift tensor %s because it depends transitively on placeholder %s via at least one path, e.g.: %s' % (repr(init_tensor), repr(op), show_path(op, init_tensor, sources)))
        for inp in graph_inputs(op):
            op_outputs[inp].add(op)
            if inp not in visited_ops and inp not in (sources or extra_sources):
                ops_to_visit.append(inp)
    return extra_sources

Walk a Graph and capture the subgraph between init_tensor and sources.

    Note: This function mutates visited_ops and op_outputs.

Args:
        init_tensor:  A Tensor or Operation where the subgraph terminates.
        sources:  A set of Tensors where subgraph extraction should stop.
        disallowed_placeholders: An optional set of ops which may not appear in the
        lifted graph. Defaults to all placeholders.
        visited_ops: A set of operations which were visited in a prior pass.
        op_outputs: A defaultdict containing the outputs of an op which are to be
        copied into the new subgraph.
        add_sources: A boolean indicating whether placeholders which are not in
        sources should be allowed.

Returns:
        The set of placeholders upon which init_tensor depends and are not in
        sources.

Raises:
        UnliftableError: if init_tensor depends on a placeholder which is not in
        sources and add_sources is False.


def get_const_dim_count(node_def):
    const_value = values_from_const(node_def)
    return const_value.ndim

Get the number of dimensions for a Const node.

Args:
        node_def: Const NodeDef.

Returns:
        Number of dimensions for the Const node.


def _prune_nodes_from_input_and_recipient_maps(self, nodes_to_prune):
    for node in nodes_to_prune:
        del self._node_inputs[node]
        del self._node_ctrl_inputs[node]
        del self._node_recipients[node]
        del self._node_ctrl_recipients[node]

Prune nodes out of input and recipient maps.

Args:
        nodes_to_prune: (`list` of `str`) Names of the nodes to be pruned.


def forward(self, input_ids: Optional[torch.Tensor]=None, past_key_values: Optional[Tuple[Tuple[torch.Tensor]]]=None, attention_mask: Optional[torch.Tensor]=None, token_type_ids: Optional[torch.Tensor]=None, position_ids: Optional[torch.Tensor]=None, head_mask: Optional[torch.Tensor]=None, inputs_embeds: Optional[torch.Tensor]=None, labels: Optional[torch.Tensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, **kwargs: Any) -> Union[Tuple, SequenceClassifierOutputWithPast]:
    if 'pixel_values' in kwargs:
        warnings.warn('The `pixel_values` argument is deprecated and will be removed in v4.47, use `input_ids` instead.', FutureWarning)
        if input_ids is not None:
            raise ValueError('You cannot pass both `pixel_values` and `input_ids`. Please make sure to only pass `input_ids`.')
        input_ids = kwargs.pop('pixel_values')
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    transformer_outputs = self.transformer(input_ids, past_key_values=past_key_values, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    hidden_states = transformer_outputs[0]
    pooled_hidden_states = hidden_states.mean(dim=1)
    logits = self.score(pooled_hidden_states)
    loss = None
    if labels is not None:
        if self.config.problem_type is None:
            if self.num_labels == 1:
                self.config.problem_type = 'regression'
            elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
                self.config.problem_type = 'single_label_classification'
            else:
                self.config.problem_type = 'multi_label_classification'
        if self.config.problem_type == 'regression':
            loss_fct = MSELoss()
            if self.num_labels == 1:
                loss = loss_fct(logits.squeeze(), labels.squeeze())
            else:
                loss = loss_fct(logits, labels)
        elif self.config.problem_type == 'single_label_classification':
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
        elif self.config.problem_type == 'multi_label_classification':
            loss_fct = BCEWithLogitsLoss()
            loss = loss_fct(logits, labels)
    if not return_dict:
        output = (logits,) + transformer_outputs[1:]
        return (loss,) + output if loss is not None else output
    return SequenceClassifierOutputWithPast(loss=loss, logits=logits, past_key_values=transformer_outputs.past_key_values, hidden_states=transformer_outputs.hidden_states, attentions=transformer_outputs.attentions)

input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
    `input_ids_length` = `sequence_length` if `past_key_values` is `None` else
    `past_key_values[0][0].shape[-2]` (`sequence_length` of input past key value states). Indices of input
    sequence tokens in the vocabulary.

    If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as
    `input_ids`.

    Indices can be obtained using [`AutoImageProcessor`]. See [`ImageGPTImageProcessor.__call__`] for details.
    labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
    Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
    config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
    `config.num_labels > 1` a classification loss is computed (Cross-Entropy).

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

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

    >>> image_processor = AutoImageProcessor.from_pretrained("openai/imagegpt-small")
    >>> model = ImageGPTForImageClassification.from_pretrained("openai/imagegpt-small")

    >>> inputs = image_processor(images=image, return_tensors="pt")
    >>> outputs = model(**inputs)
    >>> logits = outputs.logits
    ```


def _caching_device(rnn_cell):
    if context.executing_eagerly():
        return None
    if not getattr(rnn_cell, '_enable_caching_device', False):
        return None
    if control_flow_util.IsInWhileLoop(ops.get_default_graph()):
        logging.warning('Variable read device caching has been disabled because the RNN is in tf.while_loop loop context, which will cause reading stalled value in forward path. This could slow down the training due to duplicated variable reads. Please consider updating your code to remove tf.while_loop if possible.')
        return None
    if rnn_cell._dtype_policy.compute_dtype != rnn_cell._dtype_policy.variable_dtype:
        logging.warning("Variable read device caching has been disabled since it doesn't work with the mixed precision API. This is likely to cause a slowdown for RNN training due to duplicated read of variable for each timestep, which will be significant in a multi remote worker setting. Please consider disabling mixed precision API if the performance has been affected.")
        return None
    return lambda op: op.device

Returns the caching device for the RNN variable.

    This is useful for distributed training, when variable is not located as same
    device as the training worker. By enabling the device cache, this allows
    worker to read the variable once and cache locally, rather than read it every
    time step from remote when it is needed.

    Note that this is assuming the variable that cell needs for each time step is
    having the same value in the forward path, and only gets updated in the
    backprop. It is true for all the default cells (SimpleRNN, GRU, LSTM). If the
    cell body relies on any variable that gets updated every time step, then
    caching device will cause it to read the stall value.

Args:
        rnn_cell: the rnn cell instance.


def on_run_end(self, request):
    self._is_run_start = False
    if request.performed_action == framework.OnRunStartAction.DEBUG_RUN:
        partition_graphs = None
        if request.run_metadata and request.run_metadata.partition_graphs:
            partition_graphs = request.run_metadata.partition_graphs
        elif request.client_graph_def:
            partition_graphs = [request.client_graph_def]
        if request.tf_error and (not os.path.isdir(self._dump_root)):
            raise request.tf_error
        debug_dump = debug_data.DebugDumpDir(self._dump_root, partition_graphs=partition_graphs)
        debug_dump.set_python_graph(self._sess.graph)
        passed_filter = None
        passed_filter_exclude_node_names = None
        if self._active_tensor_filter:
            if not debug_dump.find(self._tensor_filters[self._active_tensor_filter], first_n=1, exclude_node_names=self._active_filter_exclude_node_names):
                self._remove_dump_root()
                return framework.OnRunEndResponse()
            else:
                passed_filter = self._active_tensor_filter
                passed_filter_exclude_node_names = self._active_filter_exclude_node_names
                self._active_tensor_filter = None
                self._active_filter_exclude_node_names = None
        self._prep_debug_cli_for_run_end(debug_dump, request.tf_error, passed_filter, passed_filter_exclude_node_names)
        self._run_start_response = self._launch_cli()
        self._remove_dump_root()
    elif request.performed_action == framework.OnRunStartAction.PROFILE_RUN:
        self._prep_profile_cli_for_run_end(self._sess.graph, request.run_metadata)
        self._run_start_response = self._launch_cli()
    else:
        self._run_start_response = None
    return framework.OnRunEndResponse()

Overrides on-run-end callback.

    Actions taken:
        1) Load the debug dump.
        2) Bring up the Analyzer CLI.

Args:
        request: An instance of OnSessionInitRequest.

Returns:
        An instance of OnSessionInitResponse.


def _add_collection_def(meta_graph_def, key, export_scope=None):
    meta_graph.add_collection_def(meta_graph_def, key, export_scope=export_scope)

Adds a collection to MetaGraphDef protocol buffer.

Args:
        meta_graph_def: MetaGraphDef protocol buffer.
        key: One of the GraphKeys or user-defined string.
        export_scope: Optional `string`. Name scope to remove.


def _evalDecodeJpeg(self, image_name, parallelism, num_iters, crop_during_decode=None, crop_window=None, tile=None):
    ops.reset_default_graph()
    image_file_path = resource_loader.get_path_to_datafile(os.path.join('core', 'lib', 'jpeg', 'testdata', image_name))
    if not os.path.exists(image_file_path):
        image_file_path = resource_loader.get_path_to_datafile(os.path.join('..', '..', 'core', 'lib', 'jpeg', 'testdata', image_name))
    if tile is None:
        image_content = variable_scope.get_variable('image_%s' % image_name, initializer=io_ops.read_file(image_file_path))
    else:
        single_image = image_ops.decode_jpeg(io_ops.read_file(image_file_path), channels=3, name='single_image')
        tiled_image = array_ops.tile(single_image, tile)
        image_content = variable_scope.get_variable('tiled_image_%s' % image_name, initializer=image_ops.encode_jpeg(tiled_image))
    with session.Session() as sess:
        self.evaluate(variables.global_variables_initializer())
        images = []
        for _ in range(parallelism):
            if crop_window is None:
                image = image_ops.decode_jpeg(image_content, channels=3)
            elif crop_during_decode:
                image = image_ops.decode_and_crop_jpeg(image_content, crop_window, channels=3)
            else:
                image = image_ops.decode_jpeg(image_content, channels=3)
                image = image_ops.crop_to_bounding_box(image, offset_height=crop_window[0], offset_width=crop_window[1], target_height=crop_window[2], target_width=crop_window[3])
            images.append(image)
        r = control_flow_ops.group(*images)
        for _ in range(3):
            self.evaluate(r)
        start_time = time.time()
        for _ in range(num_iters):
            self.evaluate(r)
        end_time = time.time()
    return end_time - start_time

Evaluate DecodeJpegOp for the given image.

    TODO(tanmingxing): add decoding+cropping as well.

Args:
        image_name: a string of image file name (without suffix).
        parallelism: the number of concurrent decode_jpeg ops to be run.
        num_iters: number of iterations for evaluation.
        crop_during_decode: If true, use fused DecodeAndCropJpeg instead of
        separate decode and crop ops. It is ignored if crop_window is None.
        crop_window: if not None, crop the decoded image. Depending on
        crop_during_decode, cropping could happen during or after decoding.
        tile: if not None, tile the image to composite a larger fake image.

Returns:
        The duration of the run in seconds.


def __init__(self, config):
    super().__init__()
    in_channels = config.bottleneck_features
    self.transformer_encoder = nn.ModuleList([ZoeDepthTransformerEncoderLayer(config) for _ in range(config.num_patch_transformer_layers)])
    self.embedding_convPxP = nn.Conv2d(in_channels, config.patch_transformer_hidden_size, kernel_size=1, stride=1, padding=0)

ViT-like transformer block

Args:
        config (`ZoeDepthConfig`):
            Model configuration class defining the model architecture.


def id_in_cluster(cluster_spec, task_type, task_id):
    _validate_cluster_spec(cluster_spec, task_type, task_id)
    cluster_spec = normalize_cluster_spec(cluster_spec).as_dict()
    if task_type == 'chief':
        return 0
    if task_type == 'worker':
        return task_id + len(cluster_spec.get('chief', []))
    if task_type == 'evaluator':
        return task_id
    raise ValueError('There is no id for task_type %r' % task_type)

Returns a unique id for the task in the `task_type`'s cluster.

    It returns an id ranging from [0, `worker_count(task_type, task_id)`).

    Note: this function assumes that "evaluate" job is in its own cluster or its
    own partition of a cluster.

Args:
        cluster_spec: a dict, `ClusterDef` or `ClusterSpec` object to be validated.
        task_type: string indicating the type of the task.
        task_id: the id of the `task_type` in this cluster.

Returns:
        an int indicating the unique id.

    Throws:
        ValueError: if `task_type` is not "chief", "worker" or "evaluator".


def index_update(x, idx, y):
    return _index_update_helper(tf_np.ndarray._with_index_update, x, idx, y)

Pure equivalent of `x[idx] = y`.

    Returns the value of x that would result from the NumPy-style indexed
    assignment `x[idx] = y`. Because it's a pure function, `x` itself won't be
    changed.

Args:
        x: an array with the values to be updated.
        idx: a Numpy-style index, consisting of `None`, integers, slice objects,
        ellipses, ndarrays with integer dtypes, or a tuple of the above.
        y: the array of updates. `y` must be broadcastable to the shape of the array
        that would be returned by `x[idx]`.

Returns:
        The updated version of `x`.


Applies a rotation to a vector. Written out by hand to avoid transfer to avoid AMP downcasting.

Args:
        r: [*, 3, 3] rotation matrices
        t: [*, 3] coordinate tensors

Returns:
        [*, 3] rotated coordinates


def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, position_embeddings: Optional[torch.Tensor]=None, reference_points=None, spatial_shapes=None, spatial_shapes_list=None, level_start_index=None, output_attentions: bool=False):
    residual = hidden_states
    hidden_states, attn_weights = self.self_attn(hidden_states=hidden_states, attention_mask=attention_mask, encoder_hidden_states=hidden_states, encoder_attention_mask=attention_mask, position_embeddings=position_embeddings, reference_points=reference_points, spatial_shapes=spatial_shapes, spatial_shapes_list=spatial_shapes_list, level_start_index=level_start_index, output_attentions=output_attentions)
    hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
    hidden_states = residual + hidden_states
    hidden_states = self.self_attn_layer_norm(hidden_states)
    residual = hidden_states
    hidden_states = self.activation_fn(self.fc1(hidden_states))
    hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
    hidden_states = self.fc2(hidden_states)
    hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
    hidden_states = residual + hidden_states
    hidden_states = self.final_layer_norm(hidden_states)
    if self.training:
        if torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any():
            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
            hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
    return (hidden_states, attn_weights)

Args:
    hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
        Input to the layer.
    attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):
        Attention mask.
    position_embeddings (`torch.FloatTensor`, *optional*):
        Position embeddings, to be added to `hidden_states`.
    reference_points (`torch.FloatTensor`, *optional*):
        Reference points.
    spatial_shapes (`torch.LongTensor`, *optional*):
        Spatial shapes of the backbone feature maps.
    spatial_shapes_list (`List[Tuple[int, int]]`, *optional*):
        Spatial shapes of the backbone feature maps (but as list for export compatibility).
    level_start_index (`torch.LongTensor`, *optional*):
        Level start index.
    output_attentions (`bool`, *optional*):
        Whether or not to return the attentions tensors of all attention layers. See `attentions` under
        returned tensors for more detail.


def metrics(expected_box_encodings, expected_scores, actual_box_encodings, actual_scores):
    squashed_expected_scores = tf.math.divide(1.0, 1.0 + tf.math.exp(-expected_scores))
    squashed_actual_scores = tf.math.divide(1.0, 1.0 + tf.math.exp(-actual_scores))
    kld_metric = kl_divergence.symmetric_kl_divergence(expected_scores, actual_scores)
    high_scoring_indices = tf.math.logical_or(tf.math.greater(squashed_expected_scores, 0.1), tf.math.greater(squashed_actual_scores, 0.1))
    high_scoring_actual_boxes = tf.where(condition=tf.broadcast_to(input=high_scoring_indices, shape=tf.shape(actual_box_encodings)), x=actual_box_encodings, y=expected_box_encodings)
    box_diff = high_scoring_actual_boxes - expected_box_encodings
    box_squared_diff = tf.math.pow(box_diff, 2)
    box_mse = tf.divide(tf.math.reduce_sum(box_squared_diff), tf.math.maximum(tf.math.count_nonzero(high_scoring_indices, dtype=tf.float32), 1.0))
    ok = tf.logical_and(kld_metric < 0.1, box_mse < 0.01)
    return [kld_metric, box_mse, ok]

Calculate metrics from expected and actual blazeface outputs.

Args:
        expected_box_encodings: box encodings from model
        expected_scores: classifications from model
        actual_box_encodings: golden box encodings
        actual_scores: golden classifications

Returns:
        two-item list with classification error and localization error


def _decorator(func):
    func.__doc__ = '\n    Assert the condition `x {sym} y` holds element-wise.\n\n    This Op checks that `x[i] {sym} y[i]` holds for every pair of (possibly\n    broadcast) elements of `x` and `y`. If both `x` and `y` are empty, this is\n    trivially satisfied.\n\n    If `x` {sym} `y` does not hold, `message`, as well as the first `summarize`\n    entries of `x` and `y` are printed, and `InvalidArgumentError` is raised.\n\n    When using inside `tf.function`, this API takes effects during execution.\n    It\'s recommended to use this API with `tf.control_dependencies` to\n    ensure the correct execution order.\n\n    In the following example, without `tf.control_dependencies`, errors may\n    not be raised at all.\n    Check `tf.control_dependencies` for more details.\n\n    >>> def check_size(x):\n    ...   with tf.control_dependencies([\n    ...       tf.debugging.{opname}(tf.size(x), {test_var},\n    ...                       message=\'Bad tensor size\')]):\n    ...     return x\n\n    >>> check_size(tf.ones([2, 3], tf.float32))\n    Traceback (most recent call last):\n       ...\n    InvalidArgumentError: ...\n\n    Args:\n      x:  Numeric `Tensor`.\n      y:  Numeric `Tensor`, same dtype as and broadcastable to `x`.\n      message: A string to prefix to the default message. (optional)\n      summarize: Print this many entries of each tensor. (optional)\n      name: A name for this operation (optional).  Defaults to "{opname}".\n\n    Returns:\n      Op that raises `InvalidArgumentError` if `x {sym} y` is False. This can\n        be used with `tf.control_dependencies` inside of `tf.function`s to\n        block followup computation until the check has executed.\n      @compatibility(eager)\n      returns None\n      @end_compatibility\n\n    Raises:\n      InvalidArgumentError: if the check can be performed immediately and\n        `x == y` is False. The check can be performed immediately during eager\n        execution or if `x` and `y` are statically known.\n    '.format(sym=sym, opname=opname, test_var=test_var)
    return func

Decorator that adds docstring to the function for symbol `sym`.

Args:
        func: Function for a TensorFlow op

Returns:
        A version of `func` with documentation attached.


def generate(self, input_ids: torch.LongTensor, pixel_values: torch.FloatTensor, attention_mask: Optional[torch.LongTensor]=None, interpolate_pos_encoding: bool=False, **generate_kwargs) -> torch.LongTensor:
    vision_outputs = self.vision_model(pixel_values=pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)
    image_embeds = vision_outputs[0]
    image_attention_mask = torch.ones(image_embeds.size()[:-1], dtype=torch.long, device=image_embeds.device)
    if isinstance(input_ids, list):
        input_ids = torch.LongTensor(input_ids)
    question_outputs = self.text_encoder(input_ids=input_ids, attention_mask=attention_mask, encoder_hidden_states=image_embeds, encoder_attention_mask=image_attention_mask, return_dict=False)
    question_embeds = question_outputs[0]
    question_attention_mask = torch.ones(question_embeds.size()[:-1], dtype=torch.long, device=question_embeds.device)
    bos_ids = torch.full((question_embeds.size(0), 1), fill_value=self.decoder_start_token_id, device=question_embeds.device)
    outputs = self.text_decoder.generate(input_ids=bos_ids, eos_token_id=self.config.text_config.sep_token_id, pad_token_id=self.config.text_config.pad_token_id, encoder_hidden_states=question_embeds, encoder_attention_mask=question_attention_mask, **generate_kwargs)
    return outputs

Overrides *generate* function to be able to use the model as a conditional generator

    Parameters:
        input_ids (*torch.LongTensor* of shape *(batch_size, sequence_length)*):
            The sequence used as a prompt for the generation.
        pixel_values (*torch.FloatTensor* of shape *(batch_size, num_channels, image_height, image_width)*:
            Input image to be processed
        attention_mask (*torch.LongTensor* of shape *(batch_size, sequence_length)*, *optional*):
            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`. `1` for
            tokens that are NOT MASKED, `0` for MASKED tokens.
        **generate_kwargs:
            Additional arguments passed to the *generate* function of the decoder

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

    >>> model = BlipForQuestionAnswering.from_pretrained("Salesforce/blip-vqa-base")
    >>> processor = AutoProcessor.from_pretrained("Salesforce/blip-vqa-base")

    >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
    >>> image = Image.open(requests.get(url, stream=True).raw)
    >>> text = "How many cats are in the picture?"

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

    >>> outputs = model.generate(**inputs)
    >>> print(processor.decode(outputs[0], skip_special_tokens=True))
    2
    ```


def direct_transformers_import(path: str, file='__init__.py') -> ModuleType:
    name = 'transformers'
    location = os.path.join(path, file)
    spec = importlib.util.spec_from_file_location(name, location, submodule_search_locations=[path])
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    module = sys.modules[name]
    return module

Imports transformers directly

Args:
        path (`str`): The path to the source file
        file (`str`, *optional*): The file to join with the path. Defaults to "__init__.py".

Returns:
        `ModuleType`: The resulting imported module


def _maybe_load_initial_epoch_from_ckpt(self, initial_epoch):
    if self._training_state is not None:
        return self._training_state.maybe_load_initial_epoch_from_ckpt(initial_epoch, mode=ModeKeys.TRAIN)
    return initial_epoch

Maybe load initial epoch from ckpt considering possible worker recovery.

    Refer to tensorflow/python/keras/distribute/worker_training_state.py
    for more information.

Args:
        initial_epoch: The original initial_epoch user passes in in `fit()`.

Returns:
        If the training is recovering from previous failure under multi-worker
        training setting, return the epoch the training is supposed to continue
        at. Otherwise, return the `initial_epoch` the user passes in.


def eval(self, session=None):
    return self._variable.eval(session=session)

In a session, computes and returns the value of this variable.

    This is not a graph construction method, it does not add ops to the graph.

    This convenience method requires a session where the graph
    containing this variable has been launched. If no session is
    passed, the default session is used.  See `tf.compat.v1.Session` for more
    information on launching a graph and on sessions.

    ```python
    v = tf.Variable([1, 2])
    init = tf.compat.v1.global_variables_initializer()

    with tf.compat.v1.Session() as sess:
        sess.run(init)
        # Usage passing the session explicitly.
        print(v.eval(sess))
        # Usage with the default session.  The 'with' block
        # above makes 'sess' the default session.
        print(v.eval())
    ```

Args:
        session: The session to use to evaluate this variable. If none, the
        default session is used.

Returns:
        A numpy `ndarray` with a copy of the value of this variable.


def get_resize_output_image_size(input_image: ImageInput, size: Union[int, Tuple[int, int], List[int], Tuple[int]], patch_size: Union[int, Tuple[int, int], List[int], Tuple[int]], input_data_format: Optional[Union[str, ChannelDimension]]=None) -> tuple:
    max_height, max_width = size if isinstance(size, (tuple, list)) else (size, size)
    patch_height, patch_width = patch_size if isinstance(patch_size, (tuple, list)) else (patch_size, patch_size)
    height, width = get_image_size(input_image, input_data_format)
    ratio = max(height / max_height, width / max_width)
    if ratio > 1:
        height = int(math.floor(height / ratio))
        width = int(math.floor(width / ratio))
    num_height_tokens, num_width_tokens = _num_image_tokens((height, width), (patch_height, patch_width))
    return (num_height_tokens * patch_height, num_width_tokens * patch_width)

Find the target (height, width) dimension of the output image after resizing given the input image and the desired
    size.

Args:
        input_image (`ImageInput`):
            The image to resize.
        size (`int` or `Tuple[int, int]`):
            Max image size an input image can be. Must be a dictionary with the key "longest_edge".
        patch_size (`int` or `Tuple[int, int]`):
            The patch_size as `(height, width)` to use for resizing the image. If patch_size is an integer, `(patch_size, patch_size)`
            will be used
        input_data_format (`ChannelDimension`, *optional*):
            The channel dimension format of the input image. If unset, will use the inferred format from the input.

Returns:
        `tuple`: The target (height, width) dimension of the output image after resizing.


def clean_args(self, args, keys):
    for k in keys:
        if getattr(args, k) is None:
            delattr(args, k)

Clean None values out of the arg namespace.

    This lets us check for a config file arg based on whether the None default
    was overwritten.

Args:
        args: an argparse.Namespace.
        keys: Keys to clean if None


def create_script(target_test):
    script = f'\nimport os\nimport subprocess\n\nresult = subprocess.run(\n    ["python3", "-m", "pytest", "-v", "-rfEp", f"{target_test}"],\n    capture_output = True,\n    text=True,\n)\nprint(result.stdout)\n\nif f"PASSED {target_test}" in result.stdout:\n    print("test passed")\n    exit(0)\nelif len(result.stderr) > 0:\n    if "ERROR: file or directory not found: " in result.stderr:\n        print("test file or directory not found in this commit")\n        exit(0)\n    elif "ERROR: not found: " in result.stderr:\n        print("test not found in this commit")\n        exit(0)\n    else:\n        print(f"pytest failed to run: {{result.stderr}}")\n        exit(-1)\nelif f"FAILED {target_test}" in result.stdout:\n    print("test failed")\n    exit(2)\n\nexit(0)\n'
    with open('target_script.py', 'w') as fp:
        fp.write(script.strip())

Create a python script to be run by `git bisect run` to determine if `target_test` passes or fails.
    If a test is not found in a commit, the script with exit code `0` (i.e. `Success`).

Args:
        target_test (`str`): The test to check.

Returns:
        `str`: The script to be run by `git bisect run`.


def __call__(self, *args, **kwargs):
    if not self._verified:
        model = self._get_func()
        concrete_func = model.get_concrete_function(*args, **kwargs)
        converter = lite.TFLiteConverterV2.from_concrete_functions([concrete_func], model)
        if self._converter_target_spec is not None:
            converter.target_spec = self._converter_target_spec
        if self._converter_allow_custom_ops is not None:
            converter.allow_custom_ops = self._converter_allow_custom_ops
        try:
            converter.convert()
        except convert.ConverterError as err:
            self._decode_error(err)
        finally:
            self._verified = True
    return self._get_func()(*args, **kwargs)

Calls decorated function object.

    Also verifies if the function is compatible with TFLite.

Returns:
        A execution result of the decorated function.


def with_step(self, step):
    self._options['step'] = step
    return self

Which profile step to use for profiling.

    The 'step' here refers to the step defined by `Profiler.add_step()` API.

Args:
        step: When multiple steps of profiles are available, select which step's
        profile to use. If -1, use average of all available steps.

Returns:
        self


def compute_ngram_keys(self, ngrams: torch.LongTensor) -> torch.LongTensor:
    if len(ngrams.shape) != 3:
        raise ValueError(f'Ngrams should be of shape (batch_size, num_ngrams, ngram_len), but is {ngrams.shape}')
    if ngrams.shape[2] != self.ngram_len:
        raise ValueError(f'Ngrams should be of shape (batch_size, num_ngrams, ngram_len), where ngram_len is {self.ngram_len}, but is {ngrams.shape}')
    batch_size, _, _ = ngrams.shape
    hash_result = torch.ones(batch_size, device=self.device, dtype=torch.long)
    hash_result = torch.vmap(self.accumulate_hash, in_dims=(None, 1), out_dims=1)(hash_result, ngrams)
    keys = self.keys[None, None, :, None]
    hash_result = torch.vmap(self.accumulate_hash, in_dims=(None, 2), out_dims=2)(hash_result, keys)
    return hash_result

Computes random keys for each ngram and depth.

Args:
        ngrams (`torch.LongTensor`):
            Ngrams (batch_size, num_ngrams, ngram_len).

Returns:
        ngram keys (batch_size, num_ngrams, depth).


def is_directory_v2(path):
    try:
        return _pywrap_file_io.IsDirectory(compat.path_to_bytes(path))
    except errors.OpError:
        return False

Returns whether the path is a directory or not.

Args:
        path: string, path to a potential directory

Returns:
        True, if the path is a directory; False otherwise


def get_public_api(api_mapping_files: Sequence[str], file_prefixes_to_strip: Sequence[str], packages_to_ignore: Sequence[str], output_package: str, module_prefix: str) -> PublicAPI:
    ea = exported_api.ExportedApi()
    for f in api_mapping_files:
        ea.read(f)
    v1_entrypoints_by_module = collections.defaultdict(set)
    v2_entrypoints_by_module = collections.defaultdict(set)

    def add_exported_symbols(api_names: list[str], s: exported_api.ExportedSymbol, entrypoints_by_module: Mapping[str, set[_Entrypoint]]):
        for api_name in api_names:
            index_of_last_dot = api_name.rfind('.')
            index_of_first_dot = api_name.find('.')
            module = output_package
            if index_of_first_dot + 1 < index_of_last_dot:
                module += f'.{api_name[index_of_first_dot + 1:index_of_last_dot]}'
            name = api_name[index_of_last_dot + 1:]
            entrypoints_by_module[module].add(_Entrypoint(module, name, s))
    for s in ea.symbols:
        if _should_skip_file(s.file_name, file_prefixes_to_strip, packages_to_ignore, module_prefix):
            continue
        add_exported_symbols(s.v1_apis, s, v1_entrypoints_by_module)
        add_exported_symbols(s.v2_apis, s, v2_entrypoints_by_module)
    v1_generated_imports_by_module = collections.defaultdict(set)
    v2_generated_imports_by_module = collections.defaultdict(set)

    def add_generated_imports(entrypoints_by_module: Mapping[str, set[_Entrypoint]], generated_imports_by_module: Mapping[str, set[str]]):
        for module in entrypoints_by_module:
            i = module.rfind('.')
            if i == -1:
                continue
            while i != -1:
                parent = module[:i]
                generated_imports_by_module[parent].add(module)
                module = parent
                i = module.rfind('.')
    add_generated_imports(v1_entrypoints_by_module, v1_generated_imports_by_module)
    add_generated_imports(v2_entrypoints_by_module, v2_generated_imports_by_module)
    docs_by_module = {}
    for d in ea.docs:
        for m in d.modules:
            if m in docs_by_module:
                raise DocExportedTwiceError(f'Docstring at {d.file_name}:{d.line_no} is registered for {m}, which already has a registered docstring.')
            docs_by_module[m] = d.docstring
    return PublicAPI(v1_entrypoints_by_module=v1_entrypoints_by_module, v2_entrypoints_by_module=v2_entrypoints_by_module, v1_generated_imports_by_module=v1_generated_imports_by_module, v2_generated_imports_by_module=v2_generated_imports_by_module, docs_by_module=docs_by_module)

Generates the structure of the public API from the given files.

Args:
        api_mapping_files: List of files containing the exported API mappings and
        docstrings.
        file_prefixes_to_strip: A list of prefixes to strip from files when
        determining the packages to ignore.
        packages_to_ignore: A list of python packages that should be ignored when
        searching for tf_exports.
        output_package: The package to use for the imports.
        module_prefix: A prefix to add to the non-generated imports.

Raises:
        DocExportedTwiceError: Two docstrings are registered for the same module.

Returns:
        The public API structure.


def configure_worker(worker):
    ip_address = worker['ipAddress']
    url = (_VERSION_SWITCHER_ENDPOINT + '/{}?restartType={}').format(ip_address, version, restart_type)
    req = urllib.request.Request(url, data=b'')
    try:
        urllib.request.urlopen(req)
    except urllib.error.HTTPError as e:
        status_code = e.code
        if status_code == 404:
            raise Exception('Tensorflow version {} is not available on Cloud TPU, try a previous nightly version or refer to https://cloud.google.com/tpu/docs/release-notes for the latest official version.'.format(version))
        else:
            raise Exception('Failed to configure worker {}'.format(ip_address))

Configure individual TPU worker.

Args:
        worker: A dict with the field ipAddress where the configure request will
        be sent.


def _process_update(self, item, feed_item):
    lp = self.landing_page_dao.get(feed_item, required=True)
    feed_item[FieldMap.CAMPAIGN_LANDING_PAGE_ID] = lp['id']
    feed_item[FieldMap.CAMPAIGN_LANDING_PAGE_NAME] = lp['name']
    item['startDate'] = StringExtensions.convertDateTimeStrToDateStr(feed_item.get(FieldMap.CAMPAIGN_START_DATE, None))
    item['endDate'] = StringExtensions.convertDateTimeStrToDateStr(feed_item.get(FieldMap.CAMPAIGN_END_DATE, None))
    item['name'] = feed_item.get(FieldMap.CAMPAIGN_NAME, None)
    item['defaultLandingPageId'] = lp['id']

Updates a campaign based on the values from the feed.

Args:
        item: Object representing the campaign to be updated, this object is
        updated directly.
        feed_item: Feed item representing campaign values from the Bulkdozer feed.


def message_to_extension(msg: message.Message, extension_cls: Type[_T]) -> _T:
    extension = extension_cls()
    add_message_to_extension(msg, extension)
    return extension

Converts an Extension profile into a generic Extension type.

Args:
        msg: The Message to convert.
        extension_cls: The type of FHIR Extension to convert to.

Returns:
        A an instance of extension_cls.


def __init__(self, config: PatchTSMixerConfig):
    super().__init__(config)
    self.use_return_dict = config.use_return_dict
    self.patcher = nn.Linear(config.patch_length, config.d_model)
    if config.use_positional_encoding:
        self.positional_encoder = PatchTSMixerPositionalEncoding(config=config)
    else:
        self.positional_encoder = None
    self.mlp_mixer_encoder = PatchTSMixerBlock(config=config)
    if config.post_init:
        self.post_init()

Encoder for PatchTSMixer which inputs patched time-series and outputs patched embeddings.

Args:
        config (`PatchTSMixerConfig`):
            Configuration.


def _pad_images(self, images: 'torch.Tensor'):
    height, width = get_image_size(images, channel_dim=ChannelDimension.FIRST)
    pad_height = int(np.sqrt(height / 2) * 3)
    pad_width = int(np.sqrt(width / 2) * 3)
    return F.pad(images, padding=(pad_width, pad_height), padding_mode='reflect')

Args:
    image (`torch.Tensor`):
        Image to pad.


def floor(x):
    if any_symbolic_tensors((x,)):
        return Floor().symbolic_call(x)
    return backend.numpy.floor(x)

Return the floor of the input, element-wise.

    The floor of the scalar `x` is the largest integer `i`, such that `i <= x`.

Args:
        x: Input tensor.

Returns:
        Output tensor, element-wise floor of `x`.


def default_compute_objective(metrics: dict[str, float]) -> float:
    metrics = copy.deepcopy(metrics)
    loss = metrics.pop('eval_loss', None)
    _ = metrics.pop('epoch', None)
    speed_metrics = [m for m in metrics.keys() if m.endswith('_runtime') or m.endswith('_per_second') or m.endswith('_compilation_time')]
    for sm in speed_metrics:
        _ = metrics.pop(sm, None)
    return loss if len(metrics) == 0 else sum(metrics.values())

The default objective to maximize/minimize when doing an hyperparameter search. It is the evaluation loss if no
    metrics are provided to the [`Trainer`], the sum of all metrics otherwise.

Args:
        metrics (`Dict[str, float]`): The metrics returned by the evaluate method.

Returns:
        `float`: The objective to minimize or maximize


def _last_path_token(builder: expressions.Builder) -> str:
    if isinstance(builder.node, _evaluation.RootMessageNode):
        return ''
    return builder.node.to_path_token()

Returns `builder`'s last path token less the resource type.

    For example:
        * "Foo" returns "" (empty string)
        * "Foo.bar" returns "bar"
        * "Foo.bar.bats" returns "bats"

Args:
        builder: The `builder` whose relative path to return.


def get_static_batch_size(layer):
    batch_input_shape, _ = get_input_shape_and_dtype(layer)
    if batch_input_shape is not None:
        return tensor_shape.Dimension(batch_input_shape[0]).value
    return None

Gets the static batch size of a Layer.

Args:
        layer: a `Layer` instance.

Returns:
        The static batch size of a Layer.


def parse_device_list(device_list_str, key=None):
    try:
        clean_lines = str(device_list_str, 'utf-8').strip().split('\n')
    except UnicodeDecodeError:
        logging.warning('unicode decode error, origin str: %s', device_list_str)
        raise
    results = []
    for line in clean_lines:
        tokens = line.strip().split('\t')
        if len(tokens) == 2 and (key is None or tokens[1] == key):
            results.append(tokens[0])
    return results

Parses a byte string representing a list of devices.

    The string is generated by calling either adb or fastboot. The tokens in
    each string is tab-separated.

Args:
        device_list_str: Output of adb or fastboot.
        key: The token that signifies a device in device_list_str. Only devices
        with the specified key in device_list_str are parsed, such as 'device' or
        'fastbootd'. If not specified, all devices listed are parsed.

Returns:
        A list of android device serial numbers.


def __init__(self, label=None, prefix='', with_timestamp=False, with_window=False, level=None):
    super().__init__(label)
    self.prefix = prefix
    self.with_timestamp = with_timestamp
    self.with_window = with_window
    self.level = level

PTransform for printing the elements of a PCollection.

Args:
        label (str): (optional) A custom label for the transform.
        prefix (str): (optional) A prefix string to prepend to each logged element.
        with_timestamp (bool): (optional) Whether to include element's timestamp.
        with_window (bool): (optional) Whether to include element's window.
        level: (optional) The logging level for the output (e.g. `logging.DEBUG`,
        `logging.INFO`, `logging.WARNING`, `logging.ERROR`). If not specified,
        the log is printed to stdout.


def field_content_length(msg: message.Message, field: Union[descriptor.FieldDescriptor, str]) -> int:
    if isinstance(field, str):
        field = _field_descriptor_for_name(msg, field)
    if field_is_repeated(field):
        return len(getattr(msg, field.name))
    return 1 if msg.HasField(field.name) else 0

Returns the size of the field.

Args:
        msg: The Message whose fields to examine.
        field: The FieldDescriptor or name of the field to examine.

Returns:
        The number of elements at the provided field. If field describes a singular
        protobuf field, this will return 1. If the field is not set, returns 0.


def _symmetric_projection(self, n):
    q = self._orthogonal_matrix(n)
    mask = math_ops.cast(random_ops.random_normal([n], seed=self.seed) > 0, self.dtype)
    if self.seed:
        self.seed += 1
    c = math_ops.multiply(q, mask)
    return math_ops.matmul(c, array_ops.matrix_transpose(c))

Compute a n x n symmetric projection matrix.

Args:
        n: Dimension.

Returns:
        A n x n symmetric projection matrix, i.e. a matrix P s.t. P=P*P, P=P^T.


def __init__(self, *args, **kwargs):
    """Initializes the instance of _MockAdbProxy."""
    super().__init__(*args, **kwargs)
    self.default_host_port = MOCK_HOST_PORT
    self.mock_shell_func = mock.Mock()
    self.mock_forward_func = mock.Mock()

Mock class of adb proxy which covers all the calls used by snippet clients.

    To enable testing snippet clients, this class extends the functionality of
    base class from the following aspects:
    * Records the arguments of all the calls to the shell method and forward
        method.
    * Handles the adb calls to stop the snippet server in the shell function
        properly.

Attributes:
        default_host_port: int, return this port if `self.forward` got 'tcp:0' as
        host port.
        mock_shell_func: mock.Mock, used for recording the calls to the shell
        method.
        mock_forward_func: mock.Mock, used for recording the calls to the forward
        method.


def transform(self, y):
    if y.ndim == 1:
        return y.reshape(-1, 1)
    return y

Makes 1D y 2D.

Args:
        y : np.ndarray
            Target y to be transformed.

Returns:
        np.ndarray
            A numpy array, of dimension at least 2.


def load_batch(fpath, label_key='labels'):
    with open(fpath, 'rb') as f:
        d = cPickle.load(f, encoding='bytes')
        d_decoded = {}
        for k, v in d.items():
            d_decoded[k.decode('utf8')] = v
        d = d_decoded
    data = d['data']
    labels = d[label_key]
    data = data.reshape(data.shape[0], 3, 32, 32)
    return (data, labels)

Internal utility for parsing CIFAR data.

Args:
        fpath: path the file to parse.
        label_key: key for label data in the retrieve
            dictionary.

Returns:
        A tuple `(data, labels)`.


def __call__(self, shardable_tensors: Sequence[sharding_util.ShardableTensor]) -> Sequence[sharding_util.Shard]:
    tensors_by_task = {}
    for shardable_tensor in shardable_tensors:
        tensor = shardable_tensor.tensor
        checkpoint_key = shardable_tensor.checkpoint_key
        slice_spec = shardable_tensor.slice_spec
        tensors_by_task.setdefault(checkpoint_key, {})[slice_spec] = tensor
    return [tensors_by_task]

Callback to split tensors into shards based on their device spec task.

Args:
        shardable_tensors: A list of ShardableTensors.

Returns:
        List of shard dicts containing tensors.
        [ {checkpoint key: {slice_spec: tensor} } ]


def _compute_euclidean_distance(cls, inputs, clusters):
    output = []
    for inp in inputs:
        with ops.colocate_with(inp, ignore_existing=True):
            squared_distance = math_ops.reduce_sum(math_ops.square(inp), 1, keepdims=True) - 2 * math_ops.matmul(inp, clusters, transpose_b=True) + array_ops.transpose(math_ops.reduce_sum(math_ops.square(clusters), 1, keepdims=True))
            output.append(squared_distance)
    return output

Computes Euclidean distance between each input and each cluster center.

Args:
        inputs: list of input Tensors.
        clusters: cluster Tensor.

Returns:
        list of Tensors, where each element corresponds to each element in inputs.
        The value is the distance of each row to all the cluster centers.


def _create_uninitialized_mirrored_tpu_variables(**kwargs):
    if kwargs.get('initial_value', None) is None:
        return _create_mirrored_tpu_variables(**kwargs)
    value_list = []
    initial_value = None
    for i, d in enumerate(devices):
        with ops.device(d):
            if i == 0:
                initial_value = kwargs.get('initial_value', None)
                with maybe_init_scope():
                    if initial_value is not None:
                        if callable(initial_value):
                            initial_value = initial_value()
                        initial_value = ops.convert_to_tensor(initial_value, dtype=kwargs.get('dtype', None))
            if i > 0:
                var0name = value_list[0].name.split(':')[0]
                kwargs['name'] = '%s/replica_%d/' % (var0name, i)
            kwargs['initial_value'] = initial_value
            if kwargs.get('dtype', None) is None:
                kwargs['dtype'] = kwargs['initial_value'].dtype
            if kwargs.get('shape', None) is None:
                kwargs['shape'] = kwargs['initial_value'].shape
            with context.device_policy(context.DEVICE_PLACEMENT_SILENT):
                v = uninitialized_variable_creator(**kwargs)
            assert not isinstance(v, tpu_values.TPUMirroredVariable)
            value_list.append(v)
    return value_list

Returns a list of `tf.Variable`s.

    The list contains `number_replicas` `tf.Variable`s and can be used to
    initialize a `TPUMirroredVariable`.

Args:
        **kwargs: the keyword arguments for creating a variable