def unwrap(maybe_tf_decorator):
    decorators = []
    cur = maybe_tf_decorator
    while True:
        if isinstance(cur, TFDecorator):
            decorators.append(cur)
        elif _has_tf_decorator_attr(cur):
            decorators.append(getattr(cur, '_tf_decorator'))
        else:
            break
        if not hasattr(decorators[-1], 'decorated_target'):
            break
        cur = decorators[-1].decorated_target
    return (decorators, cur)

Unwraps an object into a list of TFDecorators and a final target.

Args:
        maybe_tf_decorator: Any callable object.

Returns:
        A tuple whose first element is an list of TFDecorator-derived objects that
        were applied to the final callable target, and whose second element is the
        final undecorated callable target. If the `maybe_tf_decorator` parameter is
        not decorated by any TFDecorators, the first tuple element will be an empty
        list. The `TFDecorator` list is ordered from outermost to innermost
        decorators.


def _finish_parsing(self, instrumentation_block):
    formatter = _InstrumentationBlockFormatter(instrumentation_block)
    return formatter.has_completed_result_block_format(self.DEFAULT_INSTRUMENTATION_ERROR_MESSAGE)

Finishes parsing the instrumentation result block for the final
    instrumentation run status.

Args:
        instrumentation_block: _InstrumentationBlock, the instrumentation
        result block for the instrumenation run. Potentially, thisi
        could actually be method block if the instrumentation outputi
        is malformed.

Returns:
        A boolean indicating whether the instrumentation run completed
        with all the tests passing.

Raises:
        signals.TestError: Error raised if the instrumentation failed to
        complete with either a pass or fail status.


def wrap_flash_attention(query, key, value, decoder_segment_ids, custom_mask=None, attn_logits_soft_cap=None, head_shards=1, q_seq_shards=1):
    if decoder_segment_ids is not None:
        assert query.shape[2] == decoder_segment_ids.q.shape[1], 'Sharding along sequence dimension not allowed in TPU kernel attention'
    if custom_mask is not None:
        mask = splash_attention_mask.NumpyMask(array=custom_mask)
    else:
        mask = splash_attention_mask.CausalMask(shape=(query.shape[2], query.shape[2]))
    multi_head_mask = splash_attention_mask.MultiHeadMask(masks=(mask,) * query.shape[1])
    splash_kernel = splash_attention_kernel.make_splash_mha(mask=multi_head_mask, head_shards=head_shards, q_seq_shards=q_seq_shards, attn_logits_soft_cap=attn_logits_soft_cap)
    return jax.vmap(splash_kernel)(query, key, value, segment_ids=decoder_segment_ids)

Applies a wrapped flash attention mechanism using the Splash kernel.
    This function prepares the appropriate attention mask (causal or custom),
    constructs a multi-head mask, and applies the Splash multi-head attention
    kernel to the provided query, key, and value tensors. It supports optional
    sharding and soft capping of attention logits.

Args:
        query: jax.Array. The query tensor of shape
            (batch, num_heads, seq_len, head_dim).
        key: jax.Array. The key tensor of shape
            (batch, num_heads, seq_len, head_dim).
        value: jax.Array. The value tensor of shape
            (batch, num_heads, seq_len, head_dim).
        decoder_segment_ids: Optional. Segment IDs for the decoder, used for
            sharding or masking.
        custom_mask: Optional[jax.Array]. A custom attention mask to apply. If
            None, a causal mask is used.
        attn_logits_soft_cap: Optional[float]. If provided, applies a soft cap
            to the attention logits.
        head_shards: int, default=1. Number of shards for the attention heads.
        q_seq_shards: int, default=1. Number of shards for the query sequence
            dimension.

Returns:
        jax.Array: The result of applying the Splash multi-head attention
            kernel to the inputs.

Raises:
        AssertionError: If sharding along the sequence dimension is attempted
            with decoder_segment_ids.


def _list_inputs_or_outputs(self, recursive, node_name, depth, control, op_type, do_outputs=False):
    if do_outputs:
        tracker = self._debug_dump.node_recipients
        type_str = 'Recipients of'
        short_type_str = 'recipients'
    else:
        tracker = self._debug_dump.node_inputs
        type_str = 'Inputs to'
        short_type_str = 'inputs'
    lines = []
    font_attr_segs = {}
    node_name, _ = debug_graphs.parse_node_or_tensor_name(node_name)
    if not self._debug_dump.node_exists(node_name):
        return cli_shared.error('There is no node named "%s" in the partition graphs' % node_name)
    if recursive:
        max_depth = depth
    else:
        max_depth = 1
    if control:
        include_ctrls_str = ', control %s included' % short_type_str
    else:
        include_ctrls_str = ''
    line = '%s node "%s"' % (type_str, node_name)
    font_attr_segs[0] = [(len(line) - 1 - len(node_name), len(line) - 1, 'bold')]
    lines.append(line + ' (Depth limit = %d%s):' % (max_depth, include_ctrls_str))
    command_template = 'lo -c -r %s' if do_outputs else 'li -c -r %s'
    self._dfs_from_node(lines, font_attr_segs, node_name, tracker, max_depth, 1, [], control, op_type, command_template=command_template)
    lines.append('')
    lines.append('Legend:')
    lines.append('  (d): recursion depth = d.')
    if control:
        lines.append('  (Ctrl): Control input.')
    if op_type:
        lines.append('  [Op]: Input node has op type Op.')
    return debugger_cli_common.RichTextLines(lines, font_attr_segs=font_attr_segs)

Helper function used by list_inputs and list_outputs.

    Format a list of lines to display the inputs or output recipients of a
    given node.

Args:
        recursive: Whether the listing is to be done recursively, as a boolean.
        node_name: The name of the node in question, as a str.
        depth: Maximum recursion depth, applies only if recursive == True, as an
        int.
        control: Whether control inputs or control recipients are included, as a
        boolean.
        op_type: Whether the op types of the nodes are to be included, as a
        boolean.
        do_outputs: Whether recipients, instead of input nodes are to be
        listed, as a boolean.

Returns:
        Input or recipient tree formatted as a RichTextLines object.


def _preprocess(self, inputs: Sequence[torch.Tensor], freq: Sequence[int]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    input_ts, input_padding, inp_freq = ([], [], [])
    for i, ts in enumerate(inputs):
        input_len = ts.shape[0]
        padding = torch.zeros(input_len + self.horizon_len, dtype=ts.dtype, device=ts.device)
        if input_len < self.context_len:
            num_front_pad = self.context_len - input_len
            ts = torch.cat([torch.zeros(num_front_pad, dtype=ts.dtype, device=ts.device), ts], dim=0)
            padding = torch.cat([torch.ones(num_front_pad, dtype=ts.dtype, device=padding.device), padding], dim=0)
        elif input_len > self.context_len:
            ts = ts[-self.context_len:]
            padding = padding[-(self.context_len + self.horizon_len):]
        input_ts.append(ts)
        input_padding.append(padding)
        inp_freq.append(freq[i])
    return (torch.stack(input_ts, dim=0), torch.stack(input_padding, dim=0), torch.tensor(inp_freq, dtype=torch.int32).reshape(-1, 1))

Formats and pads raw inputs to feed into the model.

    This function both pads each time series to match the context length, and
    pads the inputs to meet the SPMD shape requirement.

Args:
        inputs: A list of 1d Tensors. Each Tensor is the context time series of
        a single forecast task.
        freq: list of frequencies

Returns:
    A tuple of:
    - the padded input time series to meet the model required context.
    - the padding indicator.
    - the number of padded examples for SPMD so that each core has the same
        number (a multiple of `batch_size`) of examples.


def get_absolute_positions(self, abs_pos_embeddings, has_cls_token, height, width):
    if has_cls_token:
        abs_pos_embeddings = abs_pos_embeddings[:, 1:]
    num_position = abs_pos_embeddings.shape[1]
    size = int(math.sqrt(num_position))
    if size * size != num_position:
        raise ValueError('Absolute position embeddings must be a square number.')
    if torch.jit.is_tracing() or (size != height or size != width):
        new_abs_pos_embeddings = nn.functional.interpolate(abs_pos_embeddings.reshape(1, size, size, -1).permute(0, 3, 1, 2), size=(height, width), mode='bicubic', align_corners=False)
        return new_abs_pos_embeddings.permute(0, 2, 3, 1)
    else:
        return abs_pos_embeddings.reshape(1, height, width, -1)

Calculate absolute positional embeddings. If needed, resize embeddings and remove cls_token dimension for the
    original embeddings.

Args:
        abs_pos_embeddings (`torch.Tensor`):
            Absolute positional embeddings with (1, num_position, num_channels).
        has_cls_token (`bool`):
            If true, has 1 embedding in abs_pos_embeddings for cls token.
        height (`int`):
            Height of input image tokens.
        width (`int`):
            Width of input image tokens.

Returns:
        Absolute positional embeddings after processing with shape (1, height, width, num_channels)


def add_link(self, name: str, url: str) -> None:

Adds a related link to current trial.

    Added links can be retrieved from the `Trial.related_links` property via
    `pg.poll_result`.

Args:
        name: Name for the related link.
        url: URL for this link.


def at(dataset, index):
    return structure.from_tensor_list(dataset.element_spec, gen_experimental_dataset_ops.get_element_at_index(dataset._variant_tensor, index, output_types=structure.get_flat_tensor_types(dataset.element_spec), output_shapes=structure.get_flat_tensor_shapes(dataset.element_spec)))

Returns the element at a specific index in a datasest.

    Currently, random access is supported for the following tf.data operations:

        - `tf.data.Dataset.from_tensor_slices`,
        - `tf.data.Dataset.from_tensors`,
        - `tf.data.Dataset.shuffle`,
        - `tf.data.Dataset.batch`,
        - `tf.data.Dataset.shard`,
        - `tf.data.Dataset.map`,
        - `tf.data.Dataset.range`,
        - `tf.data.Dataset.zip`,
        - `tf.data.Dataset.skip`,
        - `tf.data.Dataset.repeat`,
        - `tf.data.Dataset.list_files`,
        - `tf.data.Dataset.SSTableDataset`,
        - `tf.data.Dataset.concatenate`,
        - `tf.data.Dataset.enumerate`,
        - `tf.data.Dataset.parallel_map`,
        - `tf.data.Dataset.prefetch`,
        - `tf.data.Dataset.take`,
        - `tf.data.Dataset.cache` (in-memory only)

        Users can use the cache operation to enable random access for any dataset,
        even one comprised of transformations which are not on this list.
        E.g., to get the third element of a TFDS dataset:

        ```python
        ds = tfds.load("mnist", split="train").cache()
        elem = tf.data.Dataset.experimental.at(ds, 3)
        ```

Args:
        dataset: A `tf.data.Dataset` to determine whether it supports random access.
        index: The index at which to fetch the element.

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

Raises:
        UnimplementedError: If random access is not yet supported for a dataset.


def collect_per_output_metric_info(metrics, output_names, output_shapes, loss_fns, from_serialized=False, is_weighted=False):
    if not metrics:
        return [{} for _ in output_names]
    if isinstance(metrics, list):
        any_sub_list = any((isinstance(m, list) for m in metrics))
        if any_sub_list:
            if len(metrics) != len(output_names):
                raise ValueError('When passing a list of lists as `metrics`, it should have one entry per model output. The model has ' + str(len(output_names)) + ' outputs, but you passed metrics=' + str(metrics))
            nested_metrics = [generic_utils.to_list(m) for m in metrics]
        elif len(output_names) > 1:
            nested_metrics = []
            for _ in output_names:
                nested_metrics.append([metrics_module.clone_metric(m) for m in metrics])
        else:
            nested_metrics = [metrics]
    elif isinstance(metrics, collections.abc.Mapping):
        generic_utils.check_for_unexpected_keys('metrics', metrics, output_names)
        nested_metrics = []
        for name in output_names:
            output_metrics = generic_utils.to_list(metrics.get(name, []))
            nested_metrics.append(output_metrics)
    else:
        raise TypeError('Type of `metrics` argument not understood. Expected a list or dictionary, found: ' + str(metrics))
    per_output_metrics = []
    for i, metrics in enumerate(nested_metrics):
        metrics_dict = collections.OrderedDict()
        for metric in metrics:
            metric_name = get_metric_name(metric, is_weighted)
            metric_fn = get_metric_function(metric, output_shape=output_shapes[i], loss_fn=loss_fns[i])
            metric_fn._from_serialized = from_serialized
            if not isinstance(metric_fn, metrics_module.Metric):
                metric_fn = metrics_module.MeanMetricWrapper(metric_fn, name=metric_name)
                metric_fn._from_serialized = False
            metrics_dict[metric_name] = metric_fn
        per_output_metrics.append(metrics_dict)
    return per_output_metrics

Maps metric names and functions to model outputs.

Args:
        metrics: a list or a list of lists or a dict of metric functions.
        output_names: a list of the names (strings) of model outputs.
        output_shapes: a list of the shapes (strings) of model outputs.
        loss_fns: a list of the loss functions corresponding to the model outputs.
        from_serialized: whether the model the metrics are being sourced from is
        being initialized from a serialized format.
        is_weighted: Boolean indicating whether the given metrics are weighted.

Returns:
        A list (one entry per model output) of dicts.
        For instance, if the model has 2 outputs, and for the first output
        we want to compute "binary_accuracy" and "binary_crossentropy",
        and just "binary_accuracy" for the second output,
        the list would look like: `[{
            'acc': binary_accuracy(),
            'ce': binary_crossentropy(),
        }, {
            'acc': binary_accuracy(),
        }]`

Raises:
        TypeError: if an incorrect type is passed for the `metrics` argument.


def _fill_meta_graph_def(meta_graph_def: meta_graph_pb2.MetaGraphDef, saveable_view: _SaveableView, signature_functions: Dict[str, Callable[..., Any]], namespace_whitelist: List[str], save_custom_gradients: bool, create_saver: bool, enable_debug_stripper: bool, defaults=None) -> Tuple[_AssetInfo, ops.Graph]:
    resource_initializers = saveable_view.get_concrete_resource_initializers()
    exported_graph = ops.Graph()
    resource_initializer_ops = []
    with exported_graph.as_default():
        object_map, tensor_map, asset_info = saveable_view.map_resources()
        signatures = _generate_signatures(signature_functions, object_map, defaults)
    if save_custom_gradients:
        _trace_gradient_functions(exported_graph, saveable_view)
    with exported_graph.as_default():
        for resource_initializer_function in resource_initializers:
            asset_dependencies = []
            for capture in resource_initializer_function.graph.external_captures:
                asset_initializer = asset_info.asset_initializers_by_resource.get(capture, None)
                if asset_initializer is not None:
                    asset_dependencies.append(asset_initializer)
            with ops.control_dependencies(asset_dependencies):
                mapped_initializer = object_map[resource_initializer_function]
                resource_initializer_ops.append(mapped_initializer())
        resource_initializer_ops.extend(asset_info.asset_initializers_by_resource.values())
        with ops.control_dependencies(resource_initializer_ops):
            init_op = control_flow_ops.no_op()
        meta_graph_def.collection_def[constants.MAIN_OP_KEY].node_list.value.append(init_op.name)
        meta_graph_def.signature_def[constants.INIT_OP_SIGNATURE_KEY].CopyFrom(signature_def_utils.op_signature_def(init_op, constants.INIT_OP_SIGNATURE_KEY))

    def call_with_mapped_captures(function, args):
        if function in object_map:
            return object_map[function](*args)
        return saved_model_exported_concrete.ExportedConcreteFunction(function, tensor_map)(*args)
    for obj in object_map.values():
        obj._maybe_initialize_trackable()
    named_saveable_objects, registered_savers = save_util_v1.frozen_saveables_and_savers(graph_view=saveable_view.augmented_graph_view, object_map=object_map, to_graph=exported_graph, call_with_mapped_captures=call_with_mapped_captures)
    if create_saver:
        saver = functional_saver.MultiDeviceSaver.from_saveables(named_saveable_objects, registered_savers, call_with_mapped_captures)
        with exported_graph.as_default():
            saver_def = saver.to_proto()
            meta_graph_def.saver_def.CopyFrom(saver_def)
    _dependency_sorted_node_ids(saveable_view)
    graph_def, _ = exported_graph._as_graph_def(add_shapes=True, use_pybind11_proto=False)
    graph_def.library.registered_gradients.extend(saveable_view.gradient_defs)
    _verify_ops(graph_def, namespace_whitelist)
    meta_graph_def.graph_def.CopyFrom(graph_def)
    meta_graph_def.meta_info_def.tags.append(tag_constants.SERVING)
    if saveable_view.options.extra_tags:
        for tag in saveable_view.options.extra_tags:
            meta_graph_def.meta_info_def.tags.append(tag)
    meta_graph_def.meta_info_def.tensorflow_version = versions.__version__
    meta_graph_def.meta_info_def.tensorflow_git_version = versions.__git_version__
    meta_graph_def.meta_info_def.stripped_default_attrs = True
    meta_graph_def.asset_file_def.extend(asset_info.asset_defs)
    for signature_key, signature in signatures.items():
        meta_graph_def.signature_def[signature_key].CopyFrom(signature)
    meta_graph.strip_graph_default_valued_attrs(meta_graph_def)
    if sys.byteorder == 'big':
        utils_impl.swap_function_tensor_content(meta_graph_def, 'big', 'little')
    if enable_debug_stripper:
        _strip_debug_nodes(meta_graph_def)
    meta_graph_def.meta_info_def.stripped_op_list.MergeFrom(meta_graph.stripped_op_list_for_graph(meta_graph_def.graph_def))
    return (asset_info, exported_graph)

Generates a MetaGraph which calls `signature_functions`.

Args:
        meta_graph_def: The MetaGraphDef proto to fill.
        saveable_view: The _SaveableView being exported.
        signature_functions: A dictionary mapping signature keys to concrete
        functions containing signatures to add to the MetaGraph.
        namespace_whitelist: List of strings containing whitelisted op namespaces.
        save_custom_gradients: Whether to save custom gradients.
        create_saver: Whether to add SavedModel's native save and restore ops.
        enable_debug_stripper: Whether to strip the debug nodes from the graph.
        defaults: A dictionary mapping signature_key to dictionary of
        user_specified_name to Tensor representing default values.

Returns:
        A tuple of (_AssetInfo, Graph) containing the captured assets and
        exported Graph generated from tracing the saveable_view.


def type_parameter(self, unknown: _UnknownType, base_class: pytd.Class, item: pytd.TemplateItem) -> StrictType:
    assert is_unknown(unknown)
    name = unknown.name + '.' + base_class.name + '.' + item.type_param.name
    return StrictType(name)

This generates the type parameter when matching against a generic type.

    For example, when we match ~unknown1 against list[T], we need an additional
    type to model the T in "~unknown1[T]". This type would have the name
    "~unknown1.list.T".

Args:
        unknown: An unknown type. This is the type that's matched against
        base_class[T].
        base_class: The base class of the generic we're matching the unknown
        against. E.g. "list".
        item: The actual type parameter. ("T" in the examples above).

Returns:
        A type (pytd.Node) to represent this type parameter.


def __init__(self, embedding_shape, initializer, weight_collections=None, trainable=True, name=None, **kwargs):
    super(_EmbeddingColumnLayer, self).__init__(trainable=trainable, name=name, **kwargs)
    self._embedding_shape = embedding_shape
    self._initializer = initializer
    self._weight_collections = weight_collections

Constructor.

Args:
        embedding_shape: Shape of the embedding variable used for lookup.
        initializer: A variable initializer function to be used in embedding
        variable initialization.
        weight_collections: A list of collection names to which the Variable will
        be added. Note that, variables will also be added to collections
        `tf.GraphKeys.GLOBAL_VARIABLES` and `ops.GraphKeys.MODEL_VARIABLES`.
        trainable: If `True` also add the variable to the graph collection
        `GraphKeys.TRAINABLE_VARIABLES` (see `tf.Variable`).
        name: Name of the layer
        **kwargs: keyword named properties.


def __init__(self, config: CvtConfig, stage: int, **kwargs):
    super().__init__(**kwargs)
    self.config = config
    self.stage = stage
    if self.config.cls_token[self.stage]:
        self.cls_token = self.add_weight(shape=(1, 1, self.config.embed_dim[-1]), initializer=get_initializer(self.config.initializer_range), trainable=True, name='cvt.encoder.stages.2.cls_token')
    self.embedding = TFCvtEmbeddings(self.config, patch_size=config.patch_sizes[self.stage], num_channels=config.num_channels if self.stage == 0 else config.embed_dim[self.stage - 1], stride=config.patch_stride[self.stage], embed_dim=config.embed_dim[self.stage], padding=config.patch_padding[self.stage], dropout_rate=config.drop_rate[self.stage], name='embedding')
    drop_path_rates = tf.linspace(0.0, config.drop_path_rate[self.stage], config.depth[stage])
    drop_path_rates = [x.numpy().item() for x in drop_path_rates]
    self.layers = [TFCvtLayer(config, num_heads=config.num_heads[self.stage], embed_dim=config.embed_dim[self.stage], kernel_size=config.kernel_qkv[self.stage], stride_q=config.stride_q[self.stage], stride_kv=config.stride_kv[self.stage], padding_q=config.padding_q[self.stage], padding_kv=config.padding_kv[self.stage], qkv_projection_method=config.qkv_projection_method[self.stage], qkv_bias=config.qkv_bias[self.stage], attention_drop_rate=config.attention_drop_rate[self.stage], drop_rate=config.drop_rate[self.stage], mlp_ratio=config.mlp_ratio[self.stage], drop_path_rate=drop_path_rates[self.stage], with_cls_token=config.cls_token[self.stage], name=f'layers.{j}') for j in range(config.depth[self.stage])]

Cvt stage (encoder block). Each stage has 2 parts :
    - (1) A Convolutional Token Embedding layer
    - (2) A Convolutional Transformer Block (layer).
    The classification token is added only in the last stage.

Args:
        config ([`CvtConfig`]): Model configuration class.
        stage (`int`): Stage number.


def pmean(tensor, axis_name=None):
    if axis_name != _pmap_config.axis_name():
        raise ValueError('axis_name (%s) is not equal to that of the surrounding pmap (%s)' % (axis_name, _pmap_config.axis_name()))
    devices = _pmap_config.devices()
    if devices is None:
        raise ValueError("Can't retrieve the device list from the surrounding pmap")
    if tpu_devices(devices):
        raise ValueError('pmean for TPU is not supported yet.')
    else:
        return gen_collective_ops.collective_reduce(input=tensor, group_size=len(devices), group_key=_GROUP_KEY, instance_key=_get_instance_key(), merge_op='Add', final_op='Div', subdiv_offsets=(0,))

Mean all-reduction.

Args:
        tensor: A tensor.
        axis_name: The axis name to reduce. Must equal to that of the surrounding
        pmap.

Returns:
        The mean of the `tensor` replicas on each participating devices.


def __init__(self, num_evals, steps_per_run=1):
    self._num_evals = num_evals
    self._evals_completed = None
    self._steps_per_run_initial_value = steps_per_run

Constructs the run hook.

Args:
        num_evals: The number of evaluations to run for. if set to None, will
        iterate the dataset until all inputs are exhausted.
        steps_per_run: Number of steps executed per run call.


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

Obtains the month the timestamps in an
    [`EventSet`][temporian.EventSet]'s sampling are in.

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

    Output feature contains numbers between 1 and 12.

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

    Usage example:
        ```python
        >>> a = tp.event_set(
        ...    timestamps=["2023-02-04", "2023-02-20", "2023-03-01", "2023-05-07"],
        ...    name='special_events'
        ... )
        >>> b = a.calendar_month()
        >>> b
        indexes: ...
        features: [('calendar_month', int32)]
        events:
            (4 events):
                timestamps: [...]
                'calendar_month': [2 2 3 5]
        ...

        ```

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

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


def repeat(n: int, body: Callable[..., Union[core_types.TensorLike, Iterable]], inputs: Optional[List[core_types.TensorLike]]=None, infeed_queue: Optional[tpu_feed.InfeedQueue]=None, name: Any=None) -> List[core_types.TensorLike]:

    def _convert_to_list(xs):
        if not isinstance(xs, (list, tuple)):
            return [xs]
        else:
            return list(xs)

    def cond(i, *args):
        del args
        return i < n

    def body_wrapper(i, *args):
        return [i + 1] + _convert_to_list(body(*args))
    inputs = [0] if inputs is None else [0] + _convert_to_list(inputs)
    outputs = while_loop(cond, body_wrapper, inputs=inputs, infeed_queue=infeed_queue, name=name)
    outputs = _convert_to_list(outputs)
    if len(outputs) == 1:
        return outputs[0].op
    else:
        return outputs[1:]

Builds a training loop that executes a fixed number of iterations.

    The set of loop-carried tensors correspond to `inputs`.
    `body` must be a function that takes and returns the values of the
    loop-carried tensors.

Args:
        n: the number of loop iterations
        body: a Python function that builds the loop body.
        inputs: a list of initial values passed into the training loop or None
        (equivalent to an empty list).
        infeed_queue: if not None, the infeed queue from which to append a tuple of
        arguments as inputs to condition.
        name: (Deprecated) Does nothing.

Returns:
        The final values of the loop-carried tensors.

Raises:
        ValueError: if there is a type error.


def GetUpdates(self, source, url, since):
    proto = url.split(':')[0]
    if proto not in ('http', 'https'):
        raise error.ConfigurationError('Unsupported protocol %s' % proto)
    conn = source.conn
    conn.setopt(pycurl.OPT_FILETIME, 1)
    conn.setopt(pycurl.ENCODING, 'bzip2, gzip')
    if since is not None:
        conn.setopt(pycurl.TIMEVALUE, int(since))
        conn.setopt(pycurl.TIMECONDITION, pycurl.TIMECONDITION_IFMODSINCE)
    retry_count = 0
    resp_code = 500
    while retry_count < source.conf['retry_max']:
        try:
            source.log.debug('fetching %s', url)
            resp_code, headers, body_bytes = curl.CurlFetch(url, conn, self.log)
            self.log.debug('response code: %s', resp_code)
        finally:
            if resp_code < 400:
                if resp_code == 304:
                    return []
                if resp_code == 200:
                    break
            retry_count += 1
            self.log.warning('Failed connection: attempt #%s.', retry_count)
            if retry_count == source.conf['retry_max']:
                self.log.debug('max retries hit')
                raise error.SourceUnavailable('Max retries exceeded.')
            time.sleep(source.conf['retry_delay'])
    headers = headers.split('\r\n')
    last_modified = conn.getinfo(pycurl.INFO_FILETIME)
    self.log.debug('last modified: %s', last_modified)
    if last_modified == -1:
        for header in headers:
            if header.lower().startswith('last-modified'):
                self.log.debug('%s', header)
                http_ts_string = header[header.find(':') + 1:].strip()
                last_modified = self.FromHttpToTimestamp(http_ts_string)
                break
        else:
            http_ts_string = ''
    else:
        http_ts_string = self.FromTimestampToHttp(last_modified)
    self.log.debug('Last-modified is: %s', http_ts_string)
    try:
        body_bytes = bz2.decompress(body_bytes)
        self.log.debug('bzip encoding found')
    except IOError:
        self.log.debug('bzip encoding not found')
    response = StringIO(body_bytes.decode('utf-8'))
    data_map = self.GetMap(cache_info=response)
    if http_ts_string:
        http_ts = self.FromHttpToTimestamp(http_ts_string)
        self.log.debug('setting last modified to: %s', http_ts)
        data_map.SetModifyTimestamp(http_ts)
    return data_map

Get updates from a source.

Args:
        source: A data source
        url: url to the data we want
        since: a timestamp representing the last change (None to force-get)

Returns:
        A tuple containing the map of updates and a maximum timestamp

Raises:
        ValueError: an object in the source map is malformed
        ConfigurationError:


def _GetBcastSubshape(subscripts):
    start = subscripts.find(ellipsis)
    if start == -1:
        return (0, 0)
    remaining = len(subscripts) - (start + len(ellipsis))
    end = -remaining if remaining > 0 else None
    return (start, end)

Returns a tuple denoting the slice mapping to ellipsis.

    For a given subscript, returns a tuple (start, end) denoting the start
    axis index and the (negative) end axis index respectively. For any input
    Tensor `x` described by the subscript, `x[start:end]` would be the slice
    represented by the ellipsis. E.g. For `ab...cd` returns `[1, -2]`.

    If ellipsis is not present in `subscripts`, returns `(0, 0)`.

Args:
        subscripts: A string denoting the einsum subscript.


def Run(self, conf, args):
    try:
        options, args = self.parser.parse_args(args)
    except SystemExit as e:
        return e.code
    if options.maps:
        self.log.info('Setting configured maps to %s', options.maps)
        conf.maps = options.maps
    if not options.incremental:
        self.log.debug('performing FULL update of caches')
    else:
        self.log.debug('performing INCREMENTAL update of caches')
    if options.delay:
        self.log.info('Delaying %d seconds before executing', options.delay)
        time.sleep(options.delay)
    return self.UpdateMaps(conf, incremental=options.incremental, force_write=options.force_write, force_lock=options.force_lock)

Run the Update command.

    See Command.Run() for full documentation on the Run() method.

Args:
        conf: a nss_cache.config.Config object
        args: a list of arguments to be parsed by this command

Returns:
        0 on success, nonzero on error


def VisitUnionType(self, union):
    intersection = self.hierarchy.ExpandSuperClasses(str(union.type_list[0]))
    for t in union.type_list[1:]:
        intersection.intersection_update(self.hierarchy.ExpandSuperClasses(str(t)))
    new_type_list = tuple((pytd.NamedType(cls) for cls in intersection if not self.hierarchy.HasSubClassInSet(cls, intersection)))
    if not new_type_list:
        return union
    return pytd_utils.JoinTypes(new_type_list)

Given a union type, try to find a simplification by using superclasses.

    This is a lossy optimization that tries to map a list of types to a common
    base type. For example, int and bool are both base classes of int, so it
    would convert "Union[int, bool]" to "int".

Args:
        union: A union type.

Returns:
        A simplified type, if available.


def blackman(x):
    if any_symbolic_tensors((x,)):
        return Blackman().symbolic_call(x)
    return backend.numpy.blackman(x)

Blackman window function.
    The Blackman window is a taper formed by using a weighted cosine.

Args:
        x: Scalar or 1D Tensor. Window length.

Returns:
        A 1D tensor containing the Blackman window values.

Example:
    >>> x = keras.ops.convert_to_tensor(5)
    >>> keras.ops.blackman(x)
    array([-1.3877788e-17,  3.4000000e-01,  1.0000000e+00,  3.4000000e-01,
        -1.3877788e-17], dtype=float32)


def _get_parameter_info(param_name, documented_params, source_args_dict, param_type, optional):
    description = None
    shape = None
    shape_string = ''
    is_documented = True
    additional_info = None
    if param_name in documented_params:
        if param_type == '' and documented_params[param_name].get('type', None) is not None:
            param_type = documented_params[param_name]['type']
        optional = documented_params[param_name]['optional']
        shape = documented_params[param_name]['shape']
        shape_string = shape if shape else ''
        additional_info = documented_params[param_name]['additional_info'] or ''
        description = f'{documented_params[param_name]['description']}\n'
    elif param_name in source_args_dict:
        shape = source_args_dict[param_name]['shape']
        shape_string = ' ' + shape if shape else ''
        description = source_args_dict[param_name]['description']
        additional_info = None
    else:
        is_documented = False
    optional_string = ', *optional*' if optional else ''
    return (param_type, optional_string, shape_string, additional_info, description, is_documented)

Get parameter documentation details from the appropriate source.
    Tensor shape, optional status and description are taken from the custom docstring in priority if available.
    Type is taken from the function signature first, then from the custom docstring if missing from the signature

Args:
        param_name (`str`): Name of the parameter
        documented_params (`dict`): Dictionary of documented parameters (manually specified in the docstring)
        source_args_dict (`dict`): Default source args dictionary to use if not in documented_params
        param_type (`str`): Current parameter type (may be updated)
        optional (`bool`): Whether the parameter is optional (may be updated)


def yield_value(modality, iterable):
    if modality == Modality.CORE:
        yield from _tf_core_yield_value(iterable)
    elif modality == Modality.DATA:
        yield from _tf_data_yield_value(iterable)
    else:
        raise ValueError('Unknown modality used {} for nested structure'.format(modality))

Yield elements of `iterable` in a deterministic order.

Args:
        modality: enum value of supported modality [Modality.CORE or Modality.DATA]
        iterable: an iterable.

Yields:
        The iterable elements in a deterministic order.


def _calculate_regression_loss(answer, aggregate_mask, dist_per_cell, numeric_values, numeric_values_scale, input_mask_float, logits_aggregation, config):
    expected_result = _calculate_expected_result(dist_per_cell, numeric_values, numeric_values_scale, input_mask_float, logits_aggregation, config)
    answer_masked = tf.where(tf.math.is_nan(answer), tf.zeros_like(answer), answer)
    if config.use_normalized_answer_loss:
        normalizer = tf.stop_gradient(tf.math.maximum(tf.math.abs(expected_result), tf.math.abs(answer_masked)) + EPSILON_ZERO_DIVISION)
        normalized_answer_masked = answer_masked / normalizer
        normalized_expected_result = expected_result / normalizer
        per_example_answer_loss = tf.compat.v1.losses.huber_loss(normalized_answer_masked * aggregate_mask, normalized_expected_result * aggregate_mask, delta=tf.cast(1.0, tf.float32), reduction=tf.losses.Reduction.NONE)
    else:
        per_example_answer_loss = tf.compat.v1.losses.huber_loss(answer_masked * aggregate_mask, expected_result * aggregate_mask, delta=tf.cast(config.huber_loss_delta, tf.float32), reduction=tf.losses.Reduction.NONE)
    if config.answer_loss_cutoff is None:
        large_answer_loss_mask = tf.ones_like(per_example_answer_loss, dtype=tf.float32)
    else:
        large_answer_loss_mask = tf.where(per_example_answer_loss > config.answer_loss_cutoff, tf.zeros_like(per_example_answer_loss, dtype=tf.float32), tf.ones_like(per_example_answer_loss, dtype=tf.float32))
    per_example_answer_loss_scaled = config.answer_loss_importance * (per_example_answer_loss * aggregate_mask)
    return (per_example_answer_loss_scaled, large_answer_loss_mask)

Calculates the regression loss per example.

Args:
        answer (`tf.Tensor` of shape `(batch_size,)`):
            Answer for every example in the batch. Nan if there is no scalar answer.
        aggregate_mask (`tf.Tensor` of shape `(batch_size,)`):
            A mask set to 1 for examples that should use aggregation functions.
        dist_per_cell (`torch.distributions.Bernoulli`):
            Cell selection distribution for each cell.
        numeric_values (`tf.Tensor` of shape `(batch_size, seq_length)`):
            Numeric values of every token. Nan for tokens which are not numeric values.
        numeric_values_scale (`tf.Tensor` of shape `(batch_size, seq_length)`):
            Scale of the numeric values of every token.
        input_mask_float (`tf.Tensor` of shape `(batch_size, seq_length)`):
            Mask for the table, without question tokens and table headers.
        logits_aggregation (`tf.Tensor` of shape `(batch_size, num_aggregation_labels)`):
            Logits per aggregation operation.
        config ([`TapasConfig`]):
            Model configuration class with all the parameters of the model

Returns:
        per_example_answer_loss_scaled (`tf.Tensor` of shape `(batch_size,)`): Scales answer loss for each example in
        the batch. large_answer_loss_mask (`tf.Tensor` of shape `(batch_size,)`): A mask which is 1 for examples for
        which their answer loss is larger than the answer_loss_cutoff.


def assert_input_compatibility(input_spec, inputs, layer_name):
    if not input_spec:
        return
    input_spec = tree.flatten(input_spec)
    if isinstance(inputs, dict):
        names = [spec.name for spec in input_spec]
        if all(names):
            list_inputs = []
            for name in names:
                if name not in inputs:
                    raise ValueError(f'Missing data for input "{name}". You passed a data dictionary with keys {list(inputs.keys())}. Expected the following keys: {names}')
                list_inputs.append(inputs[name])
            inputs = list_inputs
    inputs = tree.flatten(inputs)
    if len(inputs) != len(input_spec):
        raise ValueError(f'Layer "{layer_name}" expects {len(input_spec)} input(s), but it received {len(inputs)} input tensors. Inputs received: {inputs}')
    for input_index, (x, spec) in enumerate(zip(inputs, input_spec)):
        if spec is None:
            continue
        if x is None and spec.optional:
            continue
        if not hasattr(x, 'shape'):
            raise ValueError(f"Inputs to a layer should be tensors. Got '{x}' (of type {type(x)}) as input for layer '{layer_name}'.")
        shape = backend.standardize_shape(x.shape)
        ndim = len(shape)
        if spec.ndim is not None and (not spec.allow_last_axis_squeeze):
            if ndim != spec.ndim:
                raise ValueError(f'Input {input_index} of layer "{layer_name}" is incompatible with the layer: expected ndim={spec.ndim}, found ndim={ndim}. Full shape received: {shape}')
        if spec.max_ndim is not None:
            if ndim is not None and ndim > spec.max_ndim:
                raise ValueError(f'Input {input_index} of layer "{layer_name}" is incompatible with the layer: expected max_ndim={spec.max_ndim}, found ndim={ndim}')
        if spec.min_ndim is not None:
            if ndim is not None and ndim < spec.min_ndim:
                raise ValueError(f'Input {input_index} of layer "{layer_name}" is incompatible with the layer: expected min_ndim={spec.min_ndim}, found ndim={ndim}. Full shape received: {shape}')
        if spec.dtype is not None:
            dtype = backend.standardize_dtype(x.dtype)
            if dtype != spec.dtype:
                raise ValueError(f'Input {input_index} of layer "{layer_name}" is incompatible with the layer: expected dtype={spec.dtype}, found dtype={dtype}')
        if spec.axes:
            for axis, value in spec.axes.items():
                if value is not None and shape[axis] not in {value, None}:
                    raise ValueError(f'Input {input_index} of layer "{layer_name}" is incompatible with the layer: expected axis {axis} of input shape to have value {value}, but received input with shape {shape}')
        if spec.shape is not None:
            spec_shape = spec.shape
            if spec.allow_last_axis_squeeze:
                if shape and shape[-1] == 1:
                    shape = shape[:-1]
                if spec_shape and spec_shape[-1] == 1:
                    spec_shape = spec_shape[:-1]
            for spec_dim, dim in zip(spec_shape, shape):
                if spec_dim is not None and dim is not None:
                    if spec_dim != dim:
                        raise ValueError(f'Input {input_index} of layer "{layer_name}" is incompatible with the layer: expected shape={spec.shape}, found shape={shape}')

Checks compatibility between the layer and provided inputs.

    This checks that the tensor(s) `inputs` verify the input assumptions
    of a layer (if any). If not, a clear and actional exception gets raised.

Args:
        input_spec: An InputSpec instance, list of InputSpec instances, a nested
            structure of InputSpec instances, or None.
        inputs: Input tensor, list of input tensors, or a nested structure of
            input tensors.
        layer_name: String, name of the layer (for error message formatting).

Raises:
        ValueError: in case of mismatch between
            the provided inputs and the expectations of the layer.


def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
    if not self.supports_gradient_checkpointing:
        raise ValueError(f'{self.__class__.__name__} does not support gradient checkpointing.')
    if gradient_checkpointing_kwargs is None:
        gradient_checkpointing_kwargs = {'use_reentrant': True}
    gradient_checkpointing_func = functools.partial(checkpoint, **gradient_checkpointing_kwargs)
    _is_using_old_format = 'value' in inspect.signature(self._set_gradient_checkpointing).parameters
    if not _is_using_old_format:
        self._set_gradient_checkpointing(enable=True, gradient_checkpointing_func=gradient_checkpointing_func)
    else:
        self.apply(partial(self._set_gradient_checkpointing, value=True))
        logger.warning('You are using an old version of the checkpointing format that is deprecated (We will also silently ignore `gradient_checkpointing_kwargs` in case you passed it).Please update to the new format on your modeling file. To use the new format, you need to completely remove the definition of the method `_set_gradient_checkpointing` in your model.')
    if getattr(self, '_hf_peft_config_loaded', False):
        self.enable_input_require_grads()

Activates gradient checkpointing for the current model.

    Note that in other frameworks this feature can be referred to as "activation checkpointing" or "checkpoint
    activations".

    We pass the `__call__` method of the modules instead of `forward` because `__call__` attaches all the hooks of
    the module. https://discuss.pytorch.org/t/any-different-between-model-input-and-model-forward-input/3690/2

Args:
        gradient_checkpointing_kwargs (dict, *optional*):
            Additional keyword arguments passed along to the `torch.utils.checkpoint.checkpoint` function.


def cholesky(self, name: str='cholesky') -> 'LinearOperator':
    if not self._can_use_cholesky():
        raise ValueError('Cannot take the Cholesky decomposition: Not a positive definite self adjoint matrix.')
    with self._name_scope(name):
        return self._linop_cholesky()

Returns a Cholesky factor as a `LinearOperator`.

    Given `A` representing this `LinearOperator`, if `A` is positive definite
    self-adjoint, return `L`, where `A = L L^T`, i.e. the cholesky
    decomposition.

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

Returns:
        `LinearOperator` which represents the lower triangular matrix
        in the Cholesky decomposition.

Raises:
        ValueError: When the `LinearOperator` is not hinted to be positive
        definite and self adjoint.


def RemoveSubtypeAnnotation(node, value):
    attr = GetNodeAnnotation(node, Annotation.SUBTYPE)
    if attr and value in attr:
        attr.remove(value)
        SetNodeAnnotation(node, Annotation.SUBTYPE, attr)

Removes an annotation value from the subtype annotations on the node.

Args:
        node: the node.
        value: annotation value to remove.


def add_import(self, symbol, source_module_name, source_name, dest_module_name, dest_name):
    if source_module_name.endswith('python.modules_with_exports'):
        source_module_name = symbol.__module__
    import_str = self.format_import(source_module_name, source_name, dest_name)
    full_api_name = dest_name
    if dest_module_name:
        full_api_name = dest_module_name + '.' + full_api_name
    symbol_id = -1 if not symbol else id(symbol)
    self._check_already_imported(symbol_id, full_api_name)
    if not dest_module_name and dest_name.startswith('_'):
        self._underscore_names_in_root.add(dest_name)
    priority = 0
    if symbol:
        if hasattr(symbol, '__module__'):
            priority = int(source_module_name == symbol.__module__)
        if hasattr(symbol, '__name__'):
            priority += int(source_name == symbol.__name__)
    self._module_imports[dest_module_name][full_api_name].add((import_str, priority))

Adds this import to module_imports.

Args:
        symbol: TensorFlow Python symbol.
        source_module_name: (string) Module to import from.
        source_name: (string) Name of the symbol to import.
        dest_module_name: (string) Module name to add import to.
        dest_name: (string) Import the symbol using this name.

Raises:
        SymbolExposedTwiceError: Raised when an import with the same
        dest_name has already been added to dest_module_name.


def forward(self, pixel_values: torch.Tensor, dataset_index: Optional[torch.Tensor]=None, head_mask: Optional[torch.Tensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None):
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
    embedding_output = self.embeddings(pixel_values)
    outputs = self.encoder(embedding_output, dataset_index=dataset_index, head_mask=head_mask, output_attentions=output_attentions, output_hidden_states=True, return_dict=return_dict)
    hidden_states = outputs.hidden_states if return_dict else outputs[1]
    feature_maps = ()
    for stage, hidden_state in zip(self.stage_names, hidden_states):
        if stage in self.out_features:
            hidden_state = self.layernorm(hidden_state)
            feature_maps += (hidden_state,)
    if not return_dict:
        if output_hidden_states:
            output = (feature_maps,) + outputs[1:]
        else:
            output = (feature_maps,) + outputs[2:]
        return output
    return BackboneOutput(feature_maps=feature_maps, hidden_states=outputs.hidden_states if output_hidden_states else None, attentions=outputs.attentions)

dataset_index (`torch.Tensor` of shape `(batch_size,)`):
    Index to use in the Mixture-of-Experts (MoE) blocks of the backbone.

    This corresponds to the dataset index used during training, e.g. index 0 refers to COCO.

Example:
    ```python
    >>> from transformers import VitPoseBackboneConfig, VitPoseBackbone
    >>> import torch

    >>> config = VitPoseBackboneConfig(out_indices=[-1])
    >>> model = VitPoseBackbone(config)

    >>> pixel_values = torch.randn(1, 3, 256, 192)
    >>> dataset_index = torch.tensor([1])
    >>> outputs = model(pixel_values, dataset_index)
    ```


def is_deterministic(x: Any) -> bool:
    return not contains(x, type=NonDeterministic)

Returns if the input value is deterministic.

    Example::

        @pg.symbolize
        def foo(x, y):
        pass

        assert pg.is_deterministic(1)
        assert pg.is_deterministic(foo(1, 2))
        assert not pg.is_deterministic(pg.oneof([1, 2]))
        assert not pg.is_deterministic(foo(pg.oneof([1, 2]), 3))

Args:
        x: Value to query against.

Returns:
        True if value itself is not NonDeterministic and its child and nested
        child fields do not contain NonDeterministic values.


def __init__(self, label, element=None, state=None):
    self.label = label
    self.state = state
    if element is not None:
        self.set_element(element)

Initialize a processing context object with an element and state.

    The element represents one value from a PCollection that will be accessed
    by a DoFn object during pipeline execution, and state is an arbitrary object
    where counters and other pipeline state information can be passed in.

    DoFnProcessContext objects are also used as inputs to PartitionFn instances.

Args:
        label: label of the PCollection whose element is being processed.
        element: element of a PCollection being processed using this context.
        state: a DoFnState object with state to be passed in to the DoFn object.


def _embedding_lookup_for_sparse_tensor(self, inp: sparse_tensor.SparseTensor, weight: Optional[sparse_tensor.SparseTensor], table: tf_variables.Variable, feature: tpu_embedding_v2_utils.FeatureConfig) -> tensor.Tensor:

    def sparse_to_dense_computation(inp, weight):
        if weight is None:
            weight = sparse_tensor.SparseTensor(inp.indices, array_ops.ones_like(inp.values, dtype=dtypes.float32), dense_shape=inp.dense_shape)
        inp = sparse_ops.sparse_tensor_to_dense(inp)
        weight = sparse_ops.sparse_tensor_to_dense(weight)
        return (inp, weight)
    inp, weight = tpu_replication.outside_compilation(sparse_to_dense_computation, inp=inp, weight=weight)
    embeddings = embedding_ops.embedding_lookup_v2(table, inp)
    weight = array_ops.expand_dims(weight, -1)
    embeddings *= weight
    if not feature.output_shape and feature.max_sequence_length > 0:
        embeddings = self._pad_or_truncate_with_sequence_length(embeddings, feature.max_sequence_length)
    else:
        embeddings = self._apply_combiner_to_embeddings(embeddings, weight, feature.table.combiner)
    return embeddings

Embedding lookup for sparse tensor based on its feature config.

Args:
        inp: a single SparseTensor input.
        weight: None or SparseTensor which has the same shape of the input.
        table: a table variable.
        feature: a feature config.

Returns:
        Embedding lookup result.


def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor, head_mask: Optional[torch.Tensor]=None, output_attentions: bool=False) -> Tuple[torch.Tensor, ...]:
    if output_attentions or head_mask is not None:
        logger.warning_once('DistilBertSdpaAttention is used but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True` or `head_mask`. Falling back to the manual attention implementation, but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.')
        return super().forward(query, key, value, mask, head_mask, output_attentions)
    batch_size, _, _ = query.size()
    dim_per_head = self.dim // self.n_heads

    def shape(x: torch.Tensor) -> torch.Tensor:
        """separate heads"""
        return x.view(batch_size, -1, self.n_heads, dim_per_head).transpose(1, 2)

    def unshape(x: torch.Tensor) -> torch.Tensor:
        """group heads"""
        return x.transpose(1, 2).contiguous().view(batch_size, -1, self.n_heads * dim_per_head)
    q = shape(self.q_lin(query))
    k = shape(self.k_lin(key))
    v = shape(self.v_lin(value))
    if self.require_contiguous_qkv and q.device.type == 'cuda' and (mask is not None):
        q = q.contiguous()
        k = k.contiguous()
        v = v.contiguous()
    attn_output = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=self.dropout_prob if self.training else 0.0, is_causal=False)
    attn_output = unshape(attn_output)
    attn_output = self.out_lin(attn_output)
    return (attn_output,)

Parameters:
    query: torch.tensor(bs, seq_length, dim)
    key: torch.tensor(bs, seq_length, dim)
    value: torch.tensor(bs, seq_length, dim)
    mask: torch.tensor(bs, seq_length)

Returns:
    weights: torch.tensor(bs, n_heads, seq_length, seq_length) Attention weights context: torch.tensor(bs,
    seq_length, dim) Contextualized layer. Optional: only if `output_attentions=True`


def __init__(self, *, target_profiles: CollectionType[str], structure_definition: message.Message, base_type: str, element_type: str, backbone_element_path: Optional[str], _child_defs: ChildDefinitions, _slices: Tuple[Slice, ...], _raw_url: str, cardinality: Cardinality=Cardinality.SCALAR, root_element_definition: Optional[message.Message]=None) -> None:
    super().__init__(structure_definition=structure_definition, backbone_element_path=backbone_element_path, base_type=base_type, element_type=element_type, _child_defs=_child_defs, _slices=_slices, _raw_url=_raw_url, root_element_definition=root_element_definition, cardinality=cardinality)
    object.__setattr__(self, 'target_profiles', target_profiles)

Represents a FHIR Reference.

    See more at: https://build.fhir.org/references.html

Attributes:
        target_profiles: The types of resources to which the reference may link. See
        more at:
        https://build.fhir.org/elementdefinition-definitions.html#ElementDefinition.type.targetProfile


def make_fn(shared_variable_store, device_id):
    variable_scope_access_index = {}
    assert isinstance(device_id, int)

    def create_new_variable(next_creator, **kwargs):
        """Create the variable using `next_creator` and store it."""
        canonical_name = _canonicalize_variable_name(kwargs.get('name'))
        v = next_creator(**kwargs)
        if canonical_name not in shared_variable_store:
            shared_variable_store[canonical_name] = []
        shared_variable_store[canonical_name].append(v)
        return v

    def reuse_variable(next_creator, **kwargs):
        """Re-use existing variable from store with same name (in order)."""
        del next_creator
        name = kwargs.get('name')
        canonical_name = _canonicalize_variable_name(name)
        try:
            variable_index = variable_scope_access_index.get(canonical_name, 0)
            v = shared_variable_store[canonical_name][variable_index]
            variable_scope_access_index[canonical_name] = variable_index + 1
            return v
        except (KeyError, IndexError):
            raise RuntimeError('Tried to create variable {} with mismatching name on device {}'.format(name, device_id))
    if device_id == 0:
        return create_new_variable
    else:
        return reuse_variable

Construct the variable creator function for device `device_id`.

    Constructs custom variable creator functions for the given device.
    On first device (device_id == 0), it creates the variable using the
    `next_creator`, and stores it in the provided `shared_variable_store`.
    On all other devices (device_id > 0), it tries to re-use the variable
    already created with the same name. If no such variable exists, it throws an
    error.
    Additionally, we de-uniquify variable names before checking for matches. This
    helps re-use variables which are intended to be the same but have different
    names due to variable uniquification happening upstream. Since this might
    mean we may have multiple variables with the same canonical name, we store
    them in a list per canonical name and return them in the same order as well.

Args:
        shared_variable_store: A dictionary that we will use to store variables
        created on the first device, and re-used by creators for other devices.
        device_id: Integer index of the device whose creator should be
        constructed.

Returns:
        An appropriate creator function based on device_id.


def dummy_inputs(self):
    input_ids = tf.constant(DUMMY_INPUTS, dtype=tf.int32)
    batch_size, seq_len = input_ids.shape
    VISION_DUMMY_INPUTS = tf.random.uniform(shape=(batch_size, self.config.vision_config.num_channels, self.config.vision_config.image_size, self.config.vision_config.image_size), dtype=tf.float32)
    pixel_values = tf.constant(VISION_DUMMY_INPUTS)
    dummy = {'pixel_values': pixel_values, 'input_ids': input_ids}
    return dummy

Dummy inputs to build the network.

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


def call(self, input_features: TFModelInputType | None=None, decoder_input_ids: np.ndarray | tf.Tensor | None=None, decoder_attention_mask: np.ndarray | tf.Tensor | None=None, decoder_position_ids: np.ndarray | tf.Tensor | None=None, head_mask: np.ndarray | tf.Tensor | None=None, decoder_head_mask: np.ndarray | tf.Tensor | None=None, cross_attn_head_mask: np.ndarray | tf.Tensor | None=None, encoder_outputs: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]]=None, past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]]=None, decoder_inputs_embeds: Optional[Tuple[Union[np.ndarray, tf.Tensor]]]=None, labels: np.ndarray | tf.Tensor | None=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, training: bool=False) -> Union[Tuple[tf.Tensor], TFSeq2SeqLMOutput]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    if labels is not None:
        if decoder_input_ids is None and decoder_inputs_embeds is None:
            decoder_input_ids = shift_tokens_right(labels, self.config.pad_token_id, self.config.decoder_start_token_id)
    outputs = self.model(input_features, decoder_input_ids=decoder_input_ids, encoder_outputs=encoder_outputs, decoder_attention_mask=decoder_attention_mask, decoder_position_ids=decoder_position_ids, head_mask=head_mask, decoder_head_mask=decoder_head_mask, cross_attn_head_mask=cross_attn_head_mask, past_key_values=past_key_values, decoder_inputs_embeds=decoder_inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, training=training)
    decoder_last_hidden_state = outputs[0]
    lm_logits = tf.matmul(decoder_last_hidden_state, self.get_output_embeddings().weights, transpose_b=True)
    loss = None if labels is None else self.hf_compute_loss(labels, lm_logits)
    if not return_dict:
        output = (lm_logits,) + outputs[1:]
        return (loss,) + output if loss is not None else output
    return TFSeq2SeqLMOutput(loss=loss, logits=lm_logits, past_key_values=outputs.past_key_values, decoder_hidden_states=outputs.decoder_hidden_states, decoder_attentions=outputs.decoder_attentions, cross_attentions=outputs.cross_attentions, encoder_last_hidden_state=outputs.encoder_last_hidden_state, encoder_hidden_states=outputs.encoder_hidden_states, encoder_attentions=outputs.encoder_attentions)

labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
    Labels for computing the language modeling loss. Indices should either be in `[0, ..., config.vocab_size]`
    or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is
    only computed for the tokens with labels in `[0, ..., config.vocab_size]`.

Returns:

Example:
    ```python
    >>> import tensorflow as tf
    >>> from transformers import AutoProcessor, TFWhisperForConditionalGeneration
    >>> from datasets import load_dataset

    >>> processor = AutoProcessor.from_pretrained("openai/whisper-tiny.en")
    >>> model = TFWhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en")

    >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")

    >>> inputs = processor(ds[0]["audio"]["array"], return_tensors="tf")
    >>> input_features = inputs.input_features

    >>> generated_ids = model.generate(input_features=input_features)

    >>> transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    >>> transcription
    ' Mr. Quilter is the apostle of the middle classes, and we are glad to welcome his gospel.'
    ```


def call(self, input_ids: Optional[tf.Tensor]=None, token_type_ids: Optional[tf.Tensor]=None, inputs_embeds: Optional[tf.Tensor]=None, training: bool=False) -> tf.Tensor:
    assert not (input_ids is None and inputs_embeds is None)
    if input_ids is not None:
        check_embeddings_within_bounds(input_ids, self.config.vocab_size)
        inputs_embeds = tf.gather(params=self.weight, indices=input_ids)
    input_shape = shape_list(inputs_embeds)[:-1]
    if token_type_ids is None:
        token_type_ids = tf.fill(dims=input_shape, value=0)
    token_type_embeds = tf.gather(params=self.token_type_embeddings, indices=token_type_ids)
    final_embeddings = inputs_embeds + token_type_embeds
    final_embeddings = self.LayerNorm(inputs=final_embeddings)
    final_embeddings = self.dropout(inputs=final_embeddings, training=training)
    return final_embeddings

Applies embedding based on inputs tensor.

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


def w8a8_block_fp8_matmul_triton(A: torch.Tensor, B: torch.Tensor, As: torch.Tensor, Bs: torch.Tensor, block_size: List[int], output_dtype: torch.dtype=torch.float32) -> torch.Tensor:
    assert len(block_size) == 2
    block_n, block_k = (block_size[0], block_size[1])
    assert A.shape[-1] == B.shape[-1]
    assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous()
    assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1]
    M = A.numel() // A.shape[-1]
    assert B.ndim == 2 and B.is_contiguous() and (Bs.ndim == 2)
    N, K = B.shape
    assert triton.cdiv(N, block_n) == Bs.shape[0]
    assert triton.cdiv(K, block_k) == Bs.shape[1]
    C_shape = A.shape[:-1] + (N,)
    C = A.new_empty(C_shape, dtype=output_dtype)
    BLOCK_SIZE_M = 128
    if M < BLOCK_SIZE_M:
        BLOCK_SIZE_M = triton.next_power_of_2(M)
        BLOCK_SIZE_M = max(BLOCK_SIZE_M, 16)
    BLOCK_SIZE_K = block_k
    assert block_k % BLOCK_SIZE_K == 0
    BLOCK_SIZE_N = block_n

    def grid(META):
        return (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']),)
    _w8a8_block_fp8_matmul[grid](A, B, C, As, Bs, M, N, K, block_n, block_k, A.stride(-2), A.stride(-1), B.stride(1), B.stride(0), C.stride(-2), C.stride(-1), As.stride(-2), As.stride(-1), Bs.stride(1), Bs.stride(0), BLOCK_SIZE_M=BLOCK_SIZE_M, BLOCK_SIZE_N=BLOCK_SIZE_N, BLOCK_SIZE_K=BLOCK_SIZE_K, GROUP_SIZE_M=8)
    return C

This function performs matrix multiplication with block-wise
    quantization.
    It takes two input tensors `A` and `B` with scales `As` and `Bs`.
    The output is returned in the specified `output_dtype`.

Args:
        A: The input tensor, e.g., activation.
        B: The input tensor, e.g., weight.
        As: The per-token-group quantization scale for `A`.
        Bs: The per-block quantization scale for `B`.
        block_size: The block size for per-block quantization. It should
        be 2-dim, e.g., [128, 128].
        output_dytpe: The dtype of the returned tensor.

Returns:
        torch.Tensor: The result of matmul.


def assign(self, value, use_locking=False, name=None, read_value=True):
    assign = state_ops.assign(self._variable, value, use_locking=use_locking, name=name)
    if read_value:
        return assign
    return assign.op

Assigns a new value to the variable.

    This is essentially a shortcut for `assign(self, value)`.

Args:
        value: A `Tensor`. The new value for this variable.
        use_locking: If `True`, use locking during the assignment.
        name: The name of the operation to be created
        read_value: if True, will return something which evaluates to the new
        value of the variable; if False will return the assign op.

Returns:
        A `Tensor` that will hold the new value of this variable after
        the assignment has completed.


def copy(self) -> 'TraceableStack[T]':
    return TraceableStack(self._stack)

Return a copy of self referencing the same objects but in a new list.

    This method is implemented to support thread-local stacks.

Returns:
        TraceableStack with a new list that holds existing objects.


def get_log_line_timestamp(delta=None):
    return _get_timestamp('%m-%d %H:%M:%S.%f', delta)

Returns a timestamp in the format used by log lines.

    Default is current time. If a delta is set, the return value will be
    the current time offset by delta seconds.

Args:
        delta: Number of seconds to offset from current time; can be negative.

Returns:
        A timestamp in log line format with an offset.


def select(self, inputs: List[Any], global_state: pg.geno.AttributeDict, step: int) -> List[Any]:

Select a list of outputs from the inputs.

    A selector has two use cases:

    * Used as parents selector, which selects individuals from the population
        as parents for recombination. It will be called before the recombination
        step within the :meth:`pyglove.evolution.Evolution.propose` method.

    * Used as a population updater, which selects individuals from previous
        population as a new population. It will be called everytime the
        population is updated, triggered by the
        :meth:`pyglove.evolution.Evolution.feedback` method.

Args:
        inputs: a list of objects as input.
        global_state: An `AttributeDict` object as the global state container,
        which is readable/writable during the operation.
        step: Number of examples historically proposed, which can be used for
        determining a cross over schedule.


def __init__(self, backend=None):
    self._backend = backend or backend_module.backend()

A class that can be used to switch from one backend to another.

Example:
    ```python
    backend = DynamicBackend("tensorflow")
    y = backend.square(tf.constant(...))
    backend.set_backend("jax")
    y = backend.square(jax.numpy.array(...))
    ```

Args:
        backend: Initial backend to use (string).


def output_stream(self, instruction_id, transform_id):
    raise NotImplementedError(type(self))

Returns an output stream writing elements to transform_id.

Args:
        instruction_id: which instruction this stream belongs to
        transform_id: the transform_id of the returned stream


def assert_self_adjoint(self, name='assert_self_adjoint'):
    with self._name_scope(name):
        return self._assert_self_adjoint()

Returns an `Op` that asserts this operator is self-adjoint.

    Here we check that this operator is *exactly* equal to its hermitian
    transpose.

Args:
        name:  A string name to prepend to created ops.

Returns:
        An `Assert` `Op`, that, when run, will raise an `InvalidArgumentError` if
        the operator is not self-adjoint.


def _safe_scalar_div(numerator, denominator, name):
    numerator.get_shape().with_rank_at_most(1)
    denominator.get_shape().with_rank_at_most(1)
    return math_ops.div_no_nan(numerator, denominator, name=name)

Divides two values, returning 0 if the denominator is 0.

Args:
        numerator: A scalar `float64` `Tensor`.
        denominator: A scalar `float64` `Tensor`.
        name: Name for the returned op.

Returns:
        0 if `denominator` == 0, else `numerator` / `denominator`


def convert_to_bq_name(name: str) -> str:
    return BQ_REGEX.sub('_', name).lower()

Tranform the given string into a valid BigQuery name -
    convert non-alphanumeric characters to an underscore (_)
    and lowercase the result for consistency.

Args:
        * name: original name

Returns:
        * Transformed valid name


def normalize_batch_in_training(x, gamma, beta, reduction_axes, epsilon=0.001):
    if ndim(x) == 4 and list(reduction_axes) in [[0, 1, 2], [0, 2, 3]]:
        if not _has_nchw_support() and list(reduction_axes) == [0, 2, 3]:
            return _broadcast_normalize_batch_in_training(x, gamma, beta, reduction_axes, epsilon=epsilon)
        return _fused_normalize_batch_in_training(x, gamma, beta, reduction_axes, epsilon=epsilon)
    elif sorted(reduction_axes) == list(range(ndim(x)))[:-1]:
        return _regular_normalize_batch_in_training(x, gamma, beta, reduction_axes, epsilon=epsilon)
    else:
        return _broadcast_normalize_batch_in_training(x, gamma, beta, reduction_axes, epsilon=epsilon)

Computes mean and std for batch then apply batch_normalization on batch.

Args:
        x: Input tensor or variable.
        gamma: Tensor by which to scale the input.
        beta: Tensor with which to center the input.
        reduction_axes: iterable of integers,
            axes over which to normalize.
        epsilon: Fuzz factor.

Returns:
        A tuple length of 3, `(normalized_tensor, mean, variance)`.


def size(self, url):
    return self.metadata(url).size_in_bytes

Fetches file size for a URL.

Returns:
        int size of path according to the FileSystem.

Raises:
        ``BeamIOError``: if url doesn't exist.


def _partitioner(shape, dtype):
    if not isinstance(shape, tensor_shape.TensorShape):
        raise ValueError(f'shape is not a TensorShape: {shape}')
    if not shape.is_fully_defined():
        raise ValueError(f'shape is not fully defined: {shape}')
    dtype = dtypes.as_dtype(dtype)
    if dtype.base_dtype == dtypes.string:
        element_size = bytes_per_string_element
    else:
        element_size = dtype.size
    partitions = [1] * shape.ndims
    bytes_per_slice = 1.0 * (shape.num_elements() / shape.dims[axis].value) * element_size
    slices_per_shard = max(1, math.floor(max_shard_bytes / bytes_per_slice))
    axis_shards = int(math.ceil(1.0 * shape.dims[axis].value / slices_per_shard))
    if max_shards:
        axis_shards = min(max_shards, axis_shards)
    partitions[axis] = axis_shards
    return partitions

Partitioner that partitions shards to have max_shard_bytes total size.

Args:
        shape: A `TensorShape`.
        dtype: A `DType`.

Returns:
        A tuple representing how much to slice each axis in shape.

Raises:
        ValueError: If shape is not a fully defined `TensorShape` or dtype is not
        a `DType`.


def forecast_masking(inputs: torch.Tensor, num_forecast_mask_patches: Union[list, int], unmasked_channel_indices: Optional[list]=None, mask_value: int=0):
    if isinstance(num_forecast_mask_patches, int):
        num_forecast_mask_patches = [num_forecast_mask_patches]
    forecast_mask_ratios = [1 for _ in num_forecast_mask_patches]
    batch_size, num_channels, sequence_length, num_features = inputs.shape
    mask = torch.zeros(batch_size, num_channels, sequence_length, device=inputs.device)
    t_list = []
    total_length = 0
    total_ratio = sum(forecast_mask_ratios)
    for patch_length, ratio in zip(num_forecast_mask_patches, forecast_mask_ratios):
        if patch_length <= 0 or patch_length >= sequence_length:
            raise ValueError(f'num_forecast_mask_patches {patch_length} should be greater than 0 and less than total patches.')
        temp_len = int(batch_size * ratio / total_ratio)
        t_list.append([patch_length, ratio, temp_len])
        total_length += temp_len
    t_list = sorted(t_list, key=lambda x: x[2])
    if total_length < batch_size:
        t_list[0][2] = t_list[0][2] + (batch_size - total_length)
    elif total_length > batch_size:
        t_list[-1][2] = t_list[-1][2] + (total_length - batch_size)
    batch1 = 0
    for patch_len, _, temp_len in t_list:
        batch2 = batch1 + temp_len
        mask[batch1:batch2, :, -patch_len:] = 1
        batch1 = batch2
    perm = torch.randperm(mask.shape[0])
    mask = mask[perm]
    mask = mask.unsqueeze(-1).repeat(1, 1, 1, num_features)
    if unmasked_channel_indices is not None:
        mask[:, unmasked_channel_indices, :, :] = 0
    inputs_mask = inputs.masked_fill(mask.bool(), mask_value)
    return (inputs_mask, mask[..., 0])

Forecast masking that masks the last K patches where K is from the num_forecast_mask_patches.
    If num_forecast_mask_patches is a list, samples in the batch will be randomly masked by numbers defined in the list.

    Parameters:
        inputs (`torch.Tensor`):
            Input of shape `(bs, num_channels, num_patch, patch_length)`
        num_forecast_mask_patches (`list`):
            Number of patches to be masked at the end of each batch sample. e.g. 4 or [3, 5].
        unmasked_channel_indices (`list`, *optional*):
            Indices of channels that are not masked.
        mask_value (`int`, *optional*, defaults to 0):
            Values in the masked patches will be filled by `mask_value`.

Returns:
        `tuple(torch.Tensor)`: inputs_mask, masked input, same shape as inputs Tensor and Mask tensor of shape `(bs,
        num_channels , num_patch)` or `(bs, tsg1, tsg2, num_channels, num_patch)`


def normalize_tuple(value, n, name):
    if isinstance(value, int):
        return (value,) * n
    else:
        try:
            value_tuple = tuple(value)
        except TypeError:
            raise ValueError(f'Argument `{name}` must be a tuple of {str(n)} integers. Received: {str(value)}')
        if len(value_tuple) != n:
            raise ValueError(f'Argument `{name}` must be a tuple of {str(n)} integers. Received: {str(value)}')
        for single_value in value_tuple:
            try:
                int(single_value)
            except (ValueError, TypeError):
                raise ValueError(f'Argument `{name}` must be a tuple of {str(n)} integers. Received: {str(value)} including element {str(single_value)} of type {str(type(single_value))}')
        return value_tuple

Transforms a single integer or iterable of integers into an integer tuple.

Args:
        value: The value to validate and convert. Could an int, or any iterable
        of ints.
        n: The size of the tuple to be returned.
        name: The name of the argument being validated, e.g. "strides" or
        "kernel_size". This is only used to format error messages.

Returns:
        A tuple of n integers.

Raises:
        ValueError: If something else than an int/long or iterable thereof was
        passed.


def size(self, path: str) -> int:
    raise NotImplementedError

Get size in bytes of a file on the FileSystem.

Args:
        path: string filepath of file.

    Returns: int size of file according to the FileSystem.

Raises:
        ``BeamIOError``: if path doesn't exist.


def pad_to_square(self, image: np.ndarray, background_color: Union[int, Tuple[int, int, int]]=0, data_format: Optional[Union[str, ChannelDimension]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None) -> np.array:
    height, width = get_image_size(image, input_data_format)
    num_channels = image.shape[0] if input_data_format == ChannelDimension.FIRST else image.shape[-1]
    if height == width:
        image = to_channel_dimension_format(image, data_format, input_data_format) if data_format is not None else image
        return image
    max_dim = max(height, width)
    if isinstance(background_color, int):
        background_color = [background_color]
    elif len(background_color) != num_channels:
        raise ValueError(f'background_color must have no more than {num_channels} elements to match the number of channels')
    if input_data_format == ChannelDimension.FIRST:
        result = np.zeros((num_channels, max_dim, max_dim), dtype=image.dtype)
        for i, color in enumerate(background_color):
            result[i, :, :] = color
        if width > height:
            start = (max_dim - height) // 2
            result[:, start:start + height, :] = image
        else:
            start = (max_dim - width) // 2
            result[:, :, start:start + width] = image
    else:
        result = np.zeros((max_dim, max_dim, num_channels), dtype=image.dtype)
        for i, color in enumerate(background_color):
            result[:, :, i] = color
        if width > height:
            start = (max_dim - height) // 2
            result[start:start + height, :, :] = image
        else:
            start = (max_dim - width) // 2
            result[:, start:start + width, :] = image
    image = to_channel_dimension_format(result, data_format, input_data_format) if data_format is not None else result
    return image

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

Args:
        image (`np.ndarray`):
            The image to pad.
        background_color (`int` or `Tuple[int, int, int]`, *optional*, defaults to 0):
            The color to use for the padding. Can be an integer for single channel or a
            tuple of integers representing for multi-channel images. If passed as integer
            in mutli-channel mode, it will default to `0` in subsequent channels.
        data_format (`str` or `ChannelDimension`, *optional*):
            The channel dimension format for the output image. Can be one of:
                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
            If unset, will use same as the input image.
        input_data_format (`str` or `ChannelDimension`, *optional*):
            The channel dimension format for the input image. Can be one of:
                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
            If unset, will use the inferred format of the input image.

Returns:
        `np.ndarray`: The padded image.


def empty_function(function: _evaluation.EmptyFunction, operand_result: Optional[_sql_data_types.Select], params_result: Collection[_sql_data_types.StandardSqlExpression]) -> _sql_data_types.Select:
    del params_result
    if operand_result is None:
        raise ValueError('empty() cannot be called without an operand.')
    sql_alias = 'empty_'
    sql_data_type = _sql_data_types.Boolean
    if not _fhir_path_data_types.returns_collection(function.parent_node.return_type):
        return dataclasses.replace(operand_result, select_part=operand_result.select_part.is_null(_sql_alias=sql_alias))
    else:
        return _sql_data_types.Select(select_part=_sql_data_types.RawExpression('CASE WHEN COUNT(*) = 0 THEN TRUE ELSE FALSE END', _sql_data_type=sql_data_type, _sql_alias=sql_alias), from_part=str(operand_result.to_subquery()), where_part=f'{operand_result.sql_alias} IS NOT NULL', sql_dialect=_sql_data_types.SqlDialect.SPARK)

Generates Spark SQL representing the FHIRPath empty() function.

    Returns `TRUE` if the input collection is empty and `FALSE` otherwise.

    This is the opposite of `_ExistsFunction`. If the operand is empty, then the
    result is `TRUE`.

    The returned SQL expression is a table of cardinality 1, whose value is of
    `BOOL` type. By default, `_EmptyFunction` will return `FALSE` if given no
    operand.

Args:
        function: The FHIRPath AST `EmptyFunction` node
        operand_result: The expression which is being evaluated
        params_result: The parameter passed in to function

Returns:
        A compiled Spark SQL expression.

Raises:
        ValueError: When the function is called without an operand


def write_mutations(self, throttler, rpc_stats_callback, throttle_delay=1):
    while throttler.throttle_request(time.time() * 1000):
        _LOGGER.info('Delaying request for %ds due to previous failures', throttle_delay)
        time.sleep(throttle_delay)
        rpc_stats_callback(throttled_secs=throttle_delay)
    if self._batch is None:
        self._batch = self._client.batch()
        self._batch.begin()
        for element in self._batch_elements:
            self.add_to_batch(element)
    resource = resource_identifiers.DatastoreNamespace(self._project, '')
    labels = {monitoring_infos.SERVICE_LABEL: 'Datastore', monitoring_infos.METHOD_LABEL: 'BatchDatastoreWrite', monitoring_infos.RESOURCE_LABEL: resource, monitoring_infos.DATASTORE_NAMESPACE_LABEL: '', monitoring_infos.DATASTORE_PROJECT_ID_LABEL: self._project, monitoring_infos.STATUS_LABEL: 'ok'}
    service_call_metric = ServiceCallMetric(request_count_urn=monitoring_infos.API_REQUEST_COUNT_URN, base_labels=labels)
    try:
        start_time = time.time()
        self._batch.commit()
        end_time = time.time()
        service_call_metric.call('ok')
        rpc_stats_callback(successes=1)
        throttler.successful_request(start_time * 1000)
        commit_time_ms = int((end_time - start_time) * 1000)
        return commit_time_ms
    except (ClientError, GoogleAPICallError) as e:
        self._batch = None
        service_call_metric.call(e.code.value)
        rpc_stats_callback(errors=1)
        raise
    except HttpError as e:
        service_call_metric.call(e)
        rpc_stats_callback(errors=1)
        raise

Writes a batch of mutations to Cloud Datastore.

    If a commit fails, it will be retried up to 5 times. All mutations in the
    batch will be committed again, even if the commit was partially
    successful. If the retry limit is exceeded, the last exception from
    Cloud Datastore will be raised.

    Assumes that the Datastore client library does not perform any retries on
    commits. It has not been determined how such retries would interact with
    the retries and throttler used here.
    See ``google.cloud.datastore_v1.gapic.datastore_client_config`` for
    retry config.

Args:
        rpc_stats_callback: a function to call with arguments `successes` and
        `failures` and `throttled_secs`; this is called to record successful
        and failed RPCs to Datastore and time spent waiting for throttling.
        throttler: (``apache_beam.io.gcp.datastore.v1new.adaptive_throttler.
        AdaptiveThrottler``)
        Throttler instance used to select requests to be throttled.
        throttle_delay: (:class:`float`) time in seconds to sleep when
        throttled.

Returns:
        (int) The latency of the successful RPC in milliseconds.


def _calculate_expected_result(dist_per_cell, numeric_values, numeric_values_scale, input_mask_float, logits_aggregation, config):
    if config.use_gumbel_for_cells:
        gumbel_dist = torch.distributions.RelaxedBernoulli(temperature=config.temperature, logits=dist_per_cell.logits * config.temperature)
        scaled_probability_per_cell = gumbel_dist.sample()
    else:
        scaled_probability_per_cell = dist_per_cell.probs
    scaled_probability_per_cell = scaled_probability_per_cell / numeric_values_scale * input_mask_float
    count_result = torch.sum(scaled_probability_per_cell, dim=1)
    numeric_values_masked = torch.where(torch.isnan(numeric_values), torch.zeros_like(numeric_values), numeric_values)
    sum_result = torch.sum(scaled_probability_per_cell * numeric_values_masked, dim=1)
    avg_approximation = config.average_approximation_function
    if avg_approximation == AverageApproximationFunction.RATIO:
        average_result = sum_result / (count_result + EPSILON_ZERO_DIVISION)
    elif avg_approximation == AverageApproximationFunction.FIRST_ORDER:
        ex = torch.sum(scaled_probability_per_cell, dim=1, keepdim=True) - scaled_probability_per_cell + 1
        average_result = torch.sum(numeric_values_masked * scaled_probability_per_cell / ex, dim=1)
    elif avg_approximation == AverageApproximationFunction.SECOND_ORDER:
        ex = torch.sum(scaled_probability_per_cell, dim=1, keepdim=True) - scaled_probability_per_cell + 1
        pointwise_var = scaled_probability_per_cell * (1 - scaled_probability_per_cell)
        var = torch.sum(pointwise_var, dim=1, keepdim=True) - pointwise_var
        multiplier = (var / torch.square(ex) + 1) / ex
        average_result = torch.sum(numeric_values_masked * scaled_probability_per_cell * multiplier, dim=1)
    else:
        raise ValueError(f'Invalid average_approximation_function: {config.average_approximation_function}')
    if config.use_gumbel_for_aggregation:
        gumbel_dist = torch.distributions.RelaxedOneHotCategorical(config.aggregation_temperature, logits=logits_aggregation[:, 1:])
        aggregation_op_only_probs = gumbel_dist.sample()
    else:
        aggregation_op_only_probs = nn.functional.softmax(logits_aggregation[:, 1:] / config.aggregation_temperature, dim=-1)
    all_results = torch.cat([torch.unsqueeze(sum_result, dim=1), torch.unsqueeze(average_result, dim=1), torch.unsqueeze(count_result, dim=1)], dim=1)
    expected_result = torch.sum(all_results * aggregation_op_only_probs, dim=1)
    return expected_result

Calculates the expected result given cell and aggregation probabilities.

Args:
        dist_per_cell (`torch.distributions.Bernoulli`):
            Cell selection distribution for each cell.
        numeric_values (`torch.FloatTensor` of shape `(batch_size, seq_length)`):
            Numeric values of every token. Nan for tokens which are not numeric values.
        numeric_values_scale (`torch.FloatTensor` of shape `(batch_size, seq_length)`):
            Scale of the numeric values of every token.
        input_mask_float (`torch.FloatTensor` of shape `(batch_size, seq_length)`):
            Mask for the table, without question tokens and table headers.
        logits_aggregation (`torch.FloatTensor` of shape `(batch_size, num_aggregation_labels)`):
            Logits per aggregation operation.
        config ([`TapasConfig`]):
            Model configuration class with all the hyperparameters of the model

Returns:
        expected_result (`torch.FloatTensor` of shape `(batch_size,)`): The expected result per example.


def _apply_sequence_to_tensor_op(cls, op_fn, tensor_wrappers):
    raise NotImplementedError()

Applies given sequence-to-tensor op.

    This method is used for implementing ops that take a sequence of tensors and
    return a new tensor, such as tf.concat and tf.stack. Implementing wrappers
    should apply `op_fn` to the backing tensor(s) and return an new wrapper
    instance with the combined backing tensor.

Args:
    op_fn: Callable that applies sequence-to-tensor op to the given sequence
        of Tensors. E.g. applies tf.concat.
    tensor_wrappers: a sequence of tensor wrappers to be transformed. All
        elements have the type of the implementing TensorWrapper class.

Returns:
        A TensorWrapper instance with combined backing tensor(s).


def __init__(self, initial_learning_rate, decay_steps, num_periods=0.5, alpha=0.0, beta=0.001, name=None):
    super(LinearCosineDecay, self).__init__()
    self.initial_learning_rate = initial_learning_rate
    self.decay_steps = decay_steps
    self.num_periods = num_periods
    self.alpha = alpha
    self.beta = beta
    self.name = name

Applies linear cosine decay to the learning rate.

Args:
        initial_learning_rate: A scalar `float32` or `float64` Tensor or a Python
        number. The initial learning rate.
        decay_steps: A scalar `int32` or `int64` `Tensor` or a Python number.
        Number of steps to decay over.
        num_periods: Number of periods in the cosine part of the decay.
        See computation above.
        alpha: See computation above.
        beta: See computation above.
        name: String.  Optional name of the operation.  Defaults to
        'LinearCosineDecay'.


def convert(value, dtype=None, name=None, as_ref=False, preferred_dtype=None, accepted_result_types=(core.Symbol,)):
    if dtype is not None:
        dtype = dtypes.as_dtype(dtype)
    if preferred_dtype is not None:
        preferred_dtype = dtypes.as_dtype(preferred_dtype)
    overload = getattr(value, '__tf_tensor__', None)
    if overload is not None:
        return overload(dtype, name)
    for base_type, conversion_func in get(type(value)):
        ret = None
        if dtype is None and preferred_dtype is not None:
            try:
                ret = conversion_func(value, dtype=preferred_dtype, name=name, as_ref=as_ref)
            except (TypeError, ValueError):
                pass
            else:
                if ret is not NotImplemented and ret.dtype.base_dtype != preferred_dtype.base_dtype:
                    raise RuntimeError(_add_error_prefix(f'Conversion function {conversion_func!r} for type {base_type} returned incompatible dtype: requested = {preferred_dtype.base_dtype.name}, actual = {ret.dtype.base_dtype.name}', name=name))
        if ret is None:
            ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)
        if ret is NotImplemented:
            continue
        if isinstance(ret, core.Tensor):
            to_tensor = getattr(ret, '__tf_tensor__', None)
            ret = to_tensor() if to_tensor is not None else ret
        if not isinstance(ret, accepted_result_types):
            raise RuntimeError(_add_error_prefix(f'Conversion function {conversion_func!r} for type {base_type} returned non-Tensor: {ret!r}', name=name))
        if dtype and (not dtype.is_compatible_with(ret.dtype)):
            raise RuntimeError(_add_error_prefix(f'Conversion function {conversion_func} for type {base_type} returned incompatible dtype: requested = {dtype.name}, actual = {ret.dtype.name}', name=name))
        return ret
    raise TypeError(_add_error_prefix(f'Cannot convert {value!r} with type {type(value)} to Tensor: no conversion function registered.', name=name))

Converts `value` to a `Tensor` using registered conversion functions.

Args:
        value: An object whose type has a registered `Tensor` conversion function.
        dtype: Optional element type for the returned tensor. If missing, the type
        is inferred from the type of `value`.
        name: Optional name to use if a new `Tensor` is created.
        as_ref: Optional boolean specifying if the returned value should be a
        reference-type `Tensor` (e.g. Variable). Pass-through to the registered
        conversion function. Defaults to `False`.
        preferred_dtype: Optional element type for the returned tensor.
        Used when dtype is None. In some cases, a caller may not have a dtype
        in mind when converting to a tensor, so `preferred_dtype` can be used
        as a soft preference. If the conversion to `preferred_dtype` is not
        possible, this argument has no effect.
        accepted_result_types: Optional collection of types as an allow-list
        for the returned value. If a conversion function returns an object
        which is not an instance of some type in this collection, that value
        will not be returned.

Returns:
        A `Tensor` converted from `value`.

Raises:
        ValueError: If `value` is a `Tensor` and conversion is requested
        to a `Tensor` with an incompatible `dtype`.
        TypeError: If no conversion function is registered for an element in
        `values`.
        RuntimeError: If a registered conversion function returns an invalid
        value.


def add_decomposed_rel_pos(self, attn: torch.Tensor, query: torch.Tensor, rel_pos_h: torch.Tensor, rel_pos_w: torch.Tensor, q_size: Tuple[int, int], k_size: Tuple[int, int]) -> torch.Tensor:
    query_height, query_width = q_size
    key_height, key_width = k_size
    relative_position_height = self.get_rel_pos(query_height, key_height, rel_pos_h)
    relative_position_width = self.get_rel_pos(query_width, key_width, rel_pos_w)
    batch_size, _, dim = query.shape
    reshaped_query = query.reshape(batch_size, query_height, query_width, dim)
    rel_h = torch.einsum('bhwc,hkc->bhwk', reshaped_query, relative_position_height)
    rel_w = torch.einsum('bhwc,wkc->bhwk', reshaped_query, relative_position_width)
    attn = attn.reshape(batch_size, query_height, query_width, key_height, key_width)
    attn = attn + rel_h[:, :, :, :, None] + rel_w[:, :, :, None, :]
    attn = attn.reshape(batch_size, query_height * query_width, key_height * key_width)
    return attn

Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`.
    https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py

Args:
        attn (`torch.Tensor`):
            attention map.
        query (`torch.Tensor`):
            query q in the attention layer with shape (batch_size, query_height * query_width, channel).
        rel_pos_h (`torch.Tensor`):
            relative position embeddings (Lh, channel) for height axis.
        rel_pos_w (`torch.Tensor`):
            relative position embeddings (Lw, channel) for width axis.
        q_size (tuple):
            spatial sequence size of query q with (query_height, query_width).
        k_size (tuple):
            spatial sequence size of key k with (key_height, key_width).

Returns:
        attn (`torch.Tensor`):
            attention map with added relative positional embeddings.


def SplitLine(self, line, width):
    lines = []
    chunk = ''
    w = 0
    keep = False
    for normal, control in self.SplitIntoNormalAndControl(line):
        keep = True
        while True:
            n = width - w
            w += len(normal)
            if w <= width:
                break
            lines.append(chunk + normal[:n])
            chunk = ''
            keep = False
            w = 0
            normal = normal[n:]
        chunk += normal + control
    if chunk or keep:
        lines.append(chunk)
    return lines

Splits line into width length chunks.

Args:
        line: The line to split.
        width: The width of each chunk except the last which could be smaller than
        width.

Returns:
        A list of chunks, all but the last with display width == width.


def If(cond, inputs, then_branch, else_branch, name=None):
    if isinstance(then_branch, function._DefinedFunction):
        tlist = [_.type for _ in then_branch.definition.signature.output_arg]
        return gen_functional_ops._if(cond, inputs, tlist, then_branch, else_branch, name=name)
    then_out = then_branch.structured_outputs
    else_out = else_branch.structured_outputs
    nest.assert_same_structure(then_out, else_out, expand_composites=True)
    tlist = nest.flatten(then_branch.output_dtypes)
    ret = gen_functional_ops._if(cond, inputs, tlist, then_branch, else_branch, name=name)
    return nest.pack_sequence_as(then_out, ret, expand_composites=True)

output = Cond(inputs) ?

    then_branch(inputs) : else_branch(inputs).

Args:
        cond: A `Tensor`. A scalar. If the scalar is not a boolean, the scalar is
        converted to a boolean according to the following rule: if the scalar is a
        numerical value, non-zero means True and zero means False; if the scalar
        is a string, non-empty means True and empty means False.
        inputs: A list of input tensors.
        then_branch: A function takes 'inputs' and returns a list of tensors, whose
        types are the same as what else_branch returns.
        else_branch: A function takes 'inputs' and returns a list of tensors. whose
        types are the same as what then_branch returns.
        name: A name for the operation (optional).

Returns:
        A list of tensors returned by either then_branch(inputs)
        or else_branch(inputs).


def path_to_bytes(path):
    if hasattr(path, '__fspath__'):
        path = path.__fspath__()
    return as_bytes(path)

Converts input which is a `PathLike` object to `bytes`.

    Converts from any python constant representation of a `PathLike` object
    or `str` to bytes.

Args:
        path: An object that can be converted to path representation.

Returns:
        A `bytes` object.

    Usage:
        In case a simplified `bytes` version of the path is needed from an
        `os.PathLike` object.