def _create_pseudo_names(tensors, prefix):

    def one_index(ele):
        if isinstance(ele, int):
            return ele + 1
        return ele
    flat_paths = list(nest.yield_flat_paths(tensors))
    flat_paths = nest.map_structure(one_index, flat_paths)
    names = []
    for path in flat_paths:
        if not path:
            name = prefix + '1'
        else:
            name = '_'.join((str(p) for p in path))
            if isinstance(path[0], int):
                name = prefix + name
        names.append(name)
    return names

Creates pseudo {input | output} names for subclassed Models.

    Warning: this function should only be used to define default
    names for `Metics` and `SavedModel`. No other use cases should
    rely on a `Model`'s input or output names.

    Example with dict:

    `{'a': [x1, x2], 'b': x3}` becomes:
    `['a_1', 'a_2', 'b']`

    Example with list:

    `[x, y]` becomes:
    `['output_1', 'output_2']`

Args:
        tensors: `Model`'s outputs or inputs.
        prefix: 'output_' for outputs, 'input_' for inputs.

Returns:
        Flattened list of pseudo names.


def run_tac(model_path, targets, output_path):
    if not model_path:
        raise ValueError('Invalid model_path.')
    if not targets:
        raise ValueError('Targets are not specified.')
    if not output_path:
        raise ValueError('Invalid output_path.')
    return _pywrap_tac_wrapper.run_tac(model_path, targets, output_path)

Run target aware conversion for the given tflite model file.

Args:
        model_path: Path to the tflite model file.
        targets: A list of string of the desired targets. E.g., ['GPU', 'CPU'].
        output_path: The output path.

Returns:
        Whether the optimization succeeded.

Raises:
        ValueError:
        Invalid model_path.
        Targets are not specified.
        Invalid output_path.


def from_list(index, queues):
    if not queues or not isinstance(queues, list) or (not all((isinstance(x, QueueBase) for x in queues))):
        raise TypeError('A list of queues expected')
    dtypes = queues[0].dtypes
    if not all((dtypes == q.dtypes for q in queues[1:])):
        raise TypeError('Queues do not have matching component dtypes.')
    names = queues[0].names
    if not all((names == q.names for q in queues[1:])):
        raise TypeError('Queues do not have matching component names.')
    queue_shapes = [q.shapes for q in queues]
    reduced_shapes = [functools.reduce(_shape_common, s) for s in zip(*queue_shapes)]
    queue_refs = array_ops_stack.stack([x.queue_ref for x in queues])
    selected_queue = array_ops.gather(queue_refs, index)
    return QueueBase(dtypes=dtypes, shapes=reduced_shapes, names=names, queue_ref=selected_queue)

Create a queue using the queue reference from `queues[index]`.

Args:
        index: An integer scalar tensor that determines the input that gets
        selected.
        queues: A list of `QueueBase` objects.

Returns:
        A `QueueBase` object.

Raises:
        TypeError: When `queues` is not a list of `QueueBase` objects,
        or when the data types of `queues` are not all the same.


def AddBackpropLoopCounter(self, count, outer_grad_state):
    in_separate_functions = count.graph is not ops.get_default_graph()
    if in_separate_functions:
        count = array_ops.identity(count)
    else:
        one = constant_op.constant(1, name='b_count')
    self.Enter()
    self.AddName(count.name)
    enter_count = _Enter(count, self._name, is_constant=False, parallel_iterations=self._parallel_iterations, name='b_count')
    self.loop_enters.append(enter_count)
    merge_count = merge([enter_count, enter_count])[0]
    self._pivot_for_pred = merge_count
    if in_separate_functions:
        one = constant_op.constant(1, name='b_count')
    pred = math_ops.greater_equal(merge_count, one)
    self._pivot = loop_cond(pred, name='b_count')
    switch_count = switch(merge_count, self._pivot)
    index = math_ops.subtract(switch_count[1], one)
    self._pivot_for_body = index
    next_count = _NextIteration(index)
    merge_count.op._update_input(1, next_count)
    final_zero = exit(switch_count[0], name='b_count')
    self.loop_exits.append(final_zero)
    if outer_grad_state is not None:
        outer_grad_state.grad_sync._add_control_input(final_zero.op)
    self.ExitResult([final_zero])
    self.Exit()
    return next_count

Add the backprop loop that controls the iterations.

    This is added to the backprop loop. It is used to control the loop
    termination of the backprop loop. Called in the outer context of
    this grad context.

    The pseudocode is:
        `n = count; while (n >= 1) { n--; }`

    Note that a control dependency is added to `final_zero` to ensure the
    correct execution order of stack pop ops.

Args:
        count: The number of iterations for backprop.
        outer_grad_state: The outer grad state. None if not nested.

Returns:
        The loop index.


def __init__(self) -> None:
    self.errors: List[str] = []
    self.warnings: List[str] = []

A delegate for FHIRPath encoding errors.

    Errors are logged to the corresponding `logging` context (e.g. "warning" or
    "error") and any encountered messages are stored in the corresponding
    attribute (`warnings` and `errors`, respectively). These can then be retrieved
    by the caller within the context of the larger system.

Attributes:
        errors: A list of error messages encountered.
        warnings: A list of warning messages encountered.


def _segment_reduce(values, index, segment_reduce_fn, name):
    flat_index = flatten(index)
    vector_shape = values.size()[len(index.indices.size()):]
    flattened_shape = torch.cat([torch.as_tensor([-1], dtype=torch.long), torch.as_tensor(vector_shape, dtype=torch.long)], dim=0)
    flat_values = values.reshape(flattened_shape.tolist())
    out = torch.zeros(int(flat_index.num_segments), dtype=torch.float, device=flat_values.device)
    segment_means = out.scatter_reduce(dim=0, index=flat_index.indices.long(), src=flat_values.float(), reduce=segment_reduce_fn, include_self=False)
    device = index.num_segments.device
    new_shape = torch.cat([torch.as_tensor(index.batch_shape(), dtype=torch.long, device=device), torch.as_tensor([index.num_segments], dtype=torch.long, device=device), torch.as_tensor(vector_shape, dtype=torch.long, device=device)], dim=0)
    output_values = segment_means.clone().view(new_shape.tolist()).to(values.dtype)
    output_index = range_index_map(index.batch_shape(), index.num_segments)
    return (output_values, output_index)

Applies a segment reduction segment-wise.

Args:
        values (`torch.Tensor`):
            Tensor with segment values.
        index (`IndexMap`):
            IndexMap.
        segment_reduce_fn (`str`):
            Name for the reduce operation. One of "sum", "mean", "max" or "min".
        name (`str`):
            Name for the operation. Currently not used

Returns:
        (`IndexMap`): IndexMap of shape batch_shape with elements equal to range(num_segments).


def _Commands(component, depth=3):
    if inspect.isroutine(component) or inspect.isclass(component):
        for completion in Completions(component, verbose=False):
            yield (completion,)
    if inspect.isroutine(component):
        return
    if depth < 1:
        return
    for member_name, member in VisibleMembers(component, class_attrs={}, verbose=False):
        member_name = _FormatForCommand(member_name)
        yield (member_name,)
        for command in _Commands(member, depth - 1):
            yield ((member_name,) + command)

Yields tuples representing commands.

    To use the command from Python, insert '.' between each element of the tuple.
    To use the command from the command line, insert ' ' between each element of
    the tuple.

Args:
        component: The component considered to be the root of the yielded commands.
        depth: The maximum depth with which to traverse the member DAG for commands.

Yields:
        Tuples, each tuple representing one possible command for this CLI.
        Only traverses the member DAG up to a depth of depth.


def __init__(self, config):
    super().__init__()
    self.config = config
    self.layers = nn.ModuleList()
    for channels, factor in zip(config.neck_hidden_sizes, config.reassemble_factors):
        self.layers.append(PromptDepthAnythingReassembleLayer(config, channels=channels, factor=factor))

This class reassembles the hidden states of the backbone into image-like feature representations at various
    resolutions.

    This happens in 3 stages:
    1. Take the patch embeddings and reshape them to image-like feature representations.
    2. Project the channel dimension of the hidden states according to `config.neck_hidden_sizes`.
    3. Resizing the spatial dimensions (height, width).

Args:
        config (`[PromptDepthAnythingConfig]`):
            Model configuration class defining the model architecture.


def use_wrapped_call(layer, call_fn, default_training_value=None, return_method=False):
    expects_training_arg = layer_uses_training_bool(layer)
    if hasattr(call_fn, 'original_layer_call'):
        original_call = call_fn.original_layer_call
        call_fn = call_fn.__call__
    else:
        original_call = call_fn
    fn, arg_spec = maybe_add_training_arg(original_call, call_fn, expects_training_arg, default_training_value)

    def return_outputs_and_add_losses(*args, **kwargs):
        """Returns the outputs from the layer call function, and adds the losses."""
        if return_method:
            args = args[1:]
        outputs, losses = fn(*args, **kwargs)
        layer.add_loss(losses, inputs=True)
        if context.executing_eagerly():
            for i in layer._flatten_layers():
                if i is not layer:
                    i._eager_losses = [base_layer_utils.REVIVED_LOSS_PLACEHOLDER]
        return outputs
    decorated = tf_decorator.make_decorator(target=call_fn, decorator_func=return_outputs_and_add_losses, decorator_argspec=arg_spec)
    if return_method:
        return types.MethodType(decorated, layer)
    else:
        return decorated

Creates fn that adds the losses returned by call_fn & returns the outputs.

Args:
        layer: A Keras layer object
        call_fn: tf.function that takes layer inputs (and possibly a training arg),
        and returns a tuple of (outputs, list of losses).
        default_training_value: Default value of the training kwarg. If `None`, the
        default is `K.learning_phase()`.
        return_method: Whether to return a method bound to the layer.

Returns:
        function that calls call_fn and returns the outputs. Losses returned by
        call_fn are added to the layer losses.


def auto_cast_partition_dtype():
    return False

Whether incompatible row-partitioning dtypes should be auto-converted.

    If true, then operations that combine RaggedTensors but have different
    row-partitioning tensor dtypes will be automatically cast to a
    compatible dtype (`tf.int64`).  If false, then such operations will result
    in an error.

Returns:
        `bool`


def scatter(indices, values, shape):
    if any_symbolic_tensors((indices, values)):
        return Scatter(shape=shape).symbolic_call(indices, values)
    return backend.core.scatter(indices, values, shape)

Returns a tensor of shape `shape` where `indices` are set to `values`.

    At a high level, this operation does `zeros[indices] = updates` and
    returns the output. It is equivalent to:

    ```python
    zeros = keras.ops.zeros(shape)
    output = keras.ops.scatter_update(zeros, indices, values)
    ```

Args:
        indices: A tensor or list/tuple specifying
            indices for the values in `values`.
        values: A tensor, the values to be set at `indices`.
        shape: Shape of the output tensor.

Example:
    >>> indices = [[0, 1], [1, 1]]
    >>> values = np.array([1., 1.])
    >>> keras.ops.scatter(indices, values, shape=(2, 2))
    array([[0., 1.],
        [0., 1.]])


def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor]=None, pos_emb: Optional[torch.Tensor]=None, output_attentions: Optional[torch.Tensor]=False) -> Tuple[torch.Tensor, torch.Tensor]:
    bsz, q_len, _ = hidden_states.size()
    query_states = self.linear_q(hidden_states).view(bsz, -1, self.num_heads, self.head_dim)
    key_states = self.linear_k(hidden_states).view(bsz, -1, self.num_heads, self.head_dim)
    value_states = self.linear_v(hidden_states).view(bsz, -1, self.num_heads, self.head_dim)
    bsz_pos = pos_emb.size(0)
    pos_encoding = self.linear_pos(pos_emb).view(bsz_pos, -1, self.num_heads, self.head_dim)
    query_with_bias_u = (query_states + self.pos_bias_u).transpose(1, 2)
    query_with_bias_v = (query_states + self.pos_bias_v).transpose(1, 2)
    matrix_ac = torch.matmul(query_with_bias_u, key_states.permute(0, 2, 3, 1))
    matrix_bd = torch.matmul(query_with_bias_v, pos_encoding.permute(0, 2, 3, 1))
    matrix_bd = self.shift_relative_position_tensor(matrix_bd)
    scores = (matrix_ac + matrix_bd) / math.sqrt(self.dim_key)
    if attention_mask is not None:
        expected_size = (bsz, 1, q_len)
        if attention_mask.size() != expected_size:
            raise ValueError(f'Attention mask should be of size {expected_size}, but is {attention_mask.size()}')
        attention_mask = attention_mask.unsqueeze(1).eq(0)
        min_value = float(torch.finfo(scores.dtype).min)
        scores = scores.masked_fill(attention_mask, min_value)
        attn_weights = torch.softmax(scores, dim=-1).masked_fill(attention_mask, 0.0)
    else:
        attn_weights = torch.softmax(scores, dim=-1)
    attn_weights = self.dropout(attn_weights)
    attn_output = torch.matmul(attn_weights, value_states.transpose(1, 2))
    attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, q_len, -1)
    attn_output = self.linear_out(attn_output)
    if not output_attentions:
        attn_weights = None
    return (attn_output, attn_weights)

Compute 'Scaled Dot Product Attention' with rel. positional encoding.

Args:
        hidden_states (`torch.Tensor` of shape `(batch, time2, size)`): Values of the hidden states
        attention_mask (`torch.Tensor` of shape `(batch, time1, time2)`): Mask tensor.
        pos_emb (`torch.Tensor` of shape `(batch, 2*time1-1, size)`): Positional embedding tensor.
        output_attentions (`bool`, *optional*):
            Whether or not to return the attentions tensors of all attention layers. See `attentions` under
            returned tensors for more detail.

Returns:
        `torch.Tensor`: Output tensor of shape `(batch, time1, d_model)`.


def convert(self):
    if not _jit:
        raise ImportError('Cannot import jit from jax.')
    if not self._serving_funcs:
        raise ValueError('No serving func is specified.')
    if not self._inputs:
        raise ValueError('Input tensors are not specified.')
    if len(self._inputs) != len(self._serving_funcs):
        msg = 'Input tensor mapping len {} does not match serving func len {}.'.format(len(self._inputs), len(self._serving_funcs))
        raise ValueError(msg)
    if not isinstance(self._inputs, (tuple, list)):
        raise ValueError('Input tensors should be pass in a tuple list wrapped in an array.')
    if len(self._serving_funcs) > 1:
        raise ValueError('Currently only support single serving function.')
    if not isinstance(self._inputs[0], (tuple, list)):
        raise ValueError('The input placeholders are not a dictionary.')
    input_names = []
    ordered_inputs = []
    for input_name, tensor in self._inputs[0]:
        input_names.append(input_name)
        ordered_inputs.append(tensor)
    try:
        hlo_proto = _jit(self._serving_funcs[0]).trace(*ordered_inputs).lower(lowering_platforms=('cpu',)).compiler_ir('hlo').as_serialized_hlo_module_proto()
    except Exception:
        raise ValueError('Failed to convert the given Jax function to hlo.')
    converter_kwargs = {'input_content': hlo_proto, 'input_names': input_names, 'is_proto_format': True}
    converter_kwargs.update(self._get_base_converter_args())
    quant_mode = QuantizationMode(self.optimizations, self.target_spec, self.representative_dataset, None, experimental_qdq_annotation=self._experimental_strict_qdq)
    self._validate_inference_input_output_types(quant_mode)
    converter_kwargs.update(quant_mode.converter_flags())
    result = _convert_jax_hlo(**converter_kwargs)
    return self._optimize_tflite_model(result, quant_mode, _build_conversion_flags(**converter_kwargs).debug_options, quant_io=self.experimental_new_quantizer)

Converts a Jax serving func based on instance variables.

Returns:
        The converted data in serialized format.

Raises:
        ImportError:
        If cannot import the jit from jax.
        ValueError:
        No serving function is specified.
        Input tensors are not specified.
        The truth value of an array with more than one element is ambiguous.
        Failed to convert the given Jax function to hlo.


def __init__(self, config: JetMoeConfig):
    super(JetMoeMoE, self).__init__()
    self.input_size = config.hidden_size
    self.hidden_size = config.intermediate_size
    self.activation = ACT2FN[config.activation_function]
    self.bias = torch.nn.Parameter(torch.empty(self.input_size))
    self.input_linear = JetMoeParallelExperts(config.num_local_experts, self.input_size, self.hidden_size * 2)
    self.output_linear = JetMoeParallelExperts(config.num_local_experts, self.hidden_size, self.input_size)
    self.router = JetMoeTopKGating(input_size=self.input_size, num_experts=config.num_local_experts, top_k=config.num_experts_per_tok)

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

Args:
        config:
            Configuration object with model hyperparameters.


def _render_node_traceback(self, node_name):
    lines = [RL(''), RL(''), RL('Traceback of node construction:', 'bold')]
    try:
        node_stack = self._debug_dump.node_traceback(node_name)
        for depth, (file_path, line, function_name, text) in enumerate(node_stack):
            lines.append('%d: %s' % (depth, file_path))
            attribute = debugger_cli_common.MenuItem('', 'ps %s -b %d' % (file_path, line)) if text else None
            line_number_line = RL('  ')
            line_number_line += RL('Line:     %d' % line, attribute)
            lines.append(line_number_line)
            lines.append('  Function: %s' % function_name)
            lines.append('  Text:     ' + ('"%s"' % text if text else 'None'))
            lines.append('')
    except KeyError:
        lines.append('(Node unavailable in the loaded Python graph)')
    except LookupError:
        lines.append('(Unavailable because no Python graph has been loaded)')
    return debugger_cli_common.rich_text_lines_from_rich_line_list(lines)

Render traceback of a node's creation in Python, if available.

Args:
        node_name: (str) name of the node.

Returns:
        A RichTextLines object containing the stack trace of the node's
        construction.


def generate(arglist, git_tag_override=None):
    spec, head_symlink, _, dest_file = arglist
    data = json.load(open(spec))
    git_version = None
    if not data['git']:
        git_version = b'unknown'
    else:
        old_branch = data['branch']
        new_branch = parse_branch_ref(head_symlink)
        if new_branch != old_branch:
            raise RuntimeError("Run ./configure again, branch was '%s' but is now '%s'" % (old_branch, new_branch))
        git_version = get_git_version(data['path'], git_tag_override)
    write_version_info(dest_file, git_version)

Generate version_info.cc as given `destination_file`.

Args:
        arglist: should be a sequence that contains
            spec, head_symlink, ref_symlink, destination_file.

    `destination_file` is the filename where version_info.cc will be written

    `spec` is a filename where the file contains a JSON dictionary
        'git' bool that is true if the source is in a git repo
        'path' base path of the source code
        'branch' the name of the ref specification of the current branch/tag

    `head_symlink` is a filename to HEAD that is cross-referenced against
        what is contained in the json branch designation.

    `ref_symlink` is unused in this script but passed, because the build
        system uses that file to detect when commits happen.

        git_tag_override: Override the value for the git tag. This is useful for
        releases where we want to build the release before the git tag is
        created.

Raises:
        RuntimeError: If ./configure needs to be run, RuntimeError will be raised.


def symbolic_trace(model: 'PreTrainedModel', input_names: Optional[list[str]]=None, disable_check: bool=False, tracer_cls: type[HFTracer]=HFTracer) -> GraphModule:
    if input_names is None:
        input_names = model.dummy_inputs.keys()
    input_names = list(input_names)
    concrete_args = get_concrete_args(model, input_names)
    if not disable_check:
        check_if_model_is_supported(model)
    if 'past_key_values' in input_names and (not getattr(model.config, 'use_cache', False)):
        logger.warning('`past_key_values` were specified as input names, but model.config.use_cache = False, this might lead to unexpected behavior.')
    if 'past_key_values' not in input_names and getattr(model.config, 'use_cache', False):
        logger.warning('`past_key_values` were not specified as input names, but model.config.use_cache = True. Setting model.config.use_cache = False.')
        model.config.use_cache = False
    tracer = tracer_cls()
    traced_graph = tracer.trace(model, concrete_args=concrete_args)
    traced = torch.fx.GraphModule(model, traced_graph)
    traced.config = model.config
    traced.class_for_deserialization = model.__class__
    traced.device = model.device
    return traced

Performs symbolic tracing on the model.

Args:
        model ([`PretrainedModel`]):
            The model to trace.
        input_names (`List[str]`, *optional*):
            The names of the inputs of the traced model. If unset, model.dummy_inputs.keys() are used instead.
        disable_check (`bool`, *optional*, defaults to `False`):
            If `True`, no check is done before trying to trace the model, this is mostly usesul for debugging purposes.
        tracer_cls (`Type[HFTracer]`, *optional*, defaults to `HFTracer`):
            The tracer class to use for instantiating the tracer. If unset, `HFTracer` is used instead.

Returns:
        `torch.fx.GraphModule`: A GraphModule constructed by recording operations seen while tracing the model.

Example:
        ```python
        from transformers.utils.fx import symbolic_trace

        traced_model = symbolic_trace(model, input_names=["input_ids", "attention_mask", "token_type_ids"])
        ```


def _merge_call(self, fn, args, kwargs):
    t = threading.current_thread()
    assert isinstance(t, _MirroredReplicaThread)
    t.merge_fn = fn
    t.merge_args = args
    t.merge_kwargs = kwargs
    t.captured_name_scope = t.graph.get_name_scope()
    if t.captured_name_scope:
        t.captured_name_scope += '/'
    t.captured_var_scope = variable_scope.get_variable_scope()
    t.captured_control_deps = t.graph._current_control_dependencies()
    t.merge_call_entered_in_eager = context.context().executing_eagerly()
    if ops.get_default_graph() != t.graph:
        raise RuntimeError('`merge_call` called while defining a new graph or a tf.function. This can often happen if the function `fn` passed to `strategy.run()` contains a nested `@tf.function`, and the nested `@tf.function` contains a synchronization point, such as aggregating gradients (e.g, optimizer.apply_gradients), or if the function `fn` uses a control flow statement which contains a synchronization point in the body. Such behaviors are not yet supported. Instead, please avoid nested `tf.function`s or control flow statements that may potentially cross a synchronization boundary, for example, wrap the `fn` passed to `strategy.run` or the entire `strategy.run` inside a `tf.function` or move the control flow out of `fn`. If you are subclassing a `tf.keras.Model`, please avoid decorating overridden methods `test_step` and `train_step` in `tf.function`.')
    t.has_paused.set()
    t.should_run.wait()
    t.should_run.clear()
    if t.coord.should_stop():
        raise _RequestedStop()
    t.merge_call_entered_in_eager = None
    return t.merge_result

`merge_call()` implementation for synchronized replica.

    This pauses the current replica thread and passes `fn` and its arguments to
    the main thread. The main thread will wait until all replicas pause, then
    invoke `fn` with grouped arguments. The current replica thread will continue
    after `fn` completes.

    See `_call_for_each_replica` for the logic in the main thread.

Args:
        fn: a function that is called in cross replica context with grouped
        arguments from each replica. `fn` should returns grouped values.
        args: positional arguments to `fn`.
        kwargs: keyword arguments to `fn`.

Returns:
        Return value of `fn` for the current replica.

Raises:
        RuntimeError: when merge_call happens in a different graph, e.g. in a
        different tf.function, which is not supported now.
        _RequestedStop: when stop is requested.


def count_up_to(ref, limit, name=None):
    if ref.dtype._is_ref_dtype:
        return gen_state_ops.count_up_to(ref, limit=limit, name=name)
    return gen_state_ops.resource_count_up_to(ref.handle, limit, T=ref.dtype, name=name)

Increments 'ref' until it reaches 'limit'.

Args:
        ref: A Variable. Must be one of the following types: `int32`, `int64`.
        Should be from a scalar `Variable` node.
        limit: An `int`.
        If incrementing ref would bring it above limit, instead generates an
        'OutOfRange' error.
        name: A name for the operation (optional).

Returns:
        A `Tensor`. Has the same type as `ref`.
        A copy of the input before increment. If nothing else modifies the
        input, the values produced will all be distinct.


def byte_swap_tflite_buffer(tflite_model, from_endiness, to_endiness):
    if tflite_model is None:
        return None
    model = convert_bytearray_to_object(tflite_model)
    byte_swap_tflite_model_obj(model, from_endiness, to_endiness)
    return convert_object_to_bytearray(model)

Generates a new model byte array after byte swapping its buffers field.

Args:
        tflite_model: TFLite flatbuffer in a byte array.
        from_endiness: The original endianness format of the buffers in
        tflite_model.
        to_endiness: The destined endianness format of the buffers in tflite_model.

Returns:
        TFLite flatbuffer in a byte array, after being byte swapped to to_endiness
        format.


def from_dict(cls, feature_extractor_dict: dict[str, Any], **kwargs) -> PreTrainedFeatureExtractor:
    return_unused_kwargs = kwargs.pop('return_unused_kwargs', False)
    to_remove = []
    for key, value in kwargs.items():
        if key in feature_extractor_dict:
            feature_extractor_dict[key] = value
            to_remove.append(key)
    for key in to_remove:
        kwargs.pop(key, None)
    feature_extractor = cls(**feature_extractor_dict)
    logger.info(f'Feature extractor {feature_extractor}')
    if return_unused_kwargs:
        return (feature_extractor, kwargs)
    else:
        return feature_extractor

Instantiates a type of [`~feature_extraction_utils.FeatureExtractionMixin`] from a Python dictionary of
    parameters.

Args:
        feature_extractor_dict (`Dict[str, Any]`):
            Dictionary that will be used to instantiate the feature extractor object. Such a dictionary can be
            retrieved from a pretrained checkpoint by leveraging the
            [`~feature_extraction_utils.FeatureExtractionMixin.to_dict`] method.
        kwargs (`Dict[str, Any]`):
            Additional parameters from which to initialize the feature extractor object.

Returns:
        [`~feature_extraction_utils.FeatureExtractionMixin`]: The feature extractor object instantiated from those
        parameters.


def get_image_features(self, pixel_values: torch.FloatTensor, vision_feature_layer: Optional[Union[int, List[int]]]=None, vision_feature_select_strategy: Optional[str]=None, **kwargs):
    vision_feature_layer = vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer
    vision_feature_select_strategy = vision_feature_select_strategy if vision_feature_select_strategy is not None else self.config.vision_feature_select_strategy
    downsample_ratio = self.config.downsample_ratio
    if vision_feature_layer == -1:
        vision_features = self.vision_tower(pixel_values=pixel_values).last_hidden_state
    else:
        vision_features = self.vision_model(pixel_values=pixel_values).hidden_states[vision_feature_layer]
    if vision_feature_select_strategy == 'default':
        vision_features = vision_features[:, 1:, :]
    channels = vision_features.shape[1]
    feature_size = int(channels ** 0.5)
    batch_size = vision_features.shape[0]
    vision_features = vision_features.reshape(batch_size, feature_size, feature_size, -1)
    vision_features = self.pixel_shuffle(vision_features, scale_factor=downsample_ratio)
    vision_features = vision_features.reshape(batch_size, -1, vision_features.shape[-1])
    vision_features = self.multi_modal_projector(vision_features)
    return vision_features

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

Args:
        pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`)
        The tensors corresponding to the input images.
        vision_feature_layer (`int` or `List[int]`):
            Layer index or list of layer indices to extract features from.

Returns:
        vision_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`.


def load_supported_categories(categories_path: str):
    global _load_supported_categories
    if _load_supported_categories:
        return
    with open(categories_path, encoding='utf-8') as supported_categories:
        yaml_object = yaml.load(supported_categories.read(), Loader=yaml.SafeLoader)
    Tag.Config.supported_categories = yaml_object[TagFields.categories]
    _load_supported_categories = True

Load the list of supported categories from categories_path file
    into Tag model config

Args:
        categories_path: path to the file with categories.


def debug(msg: str, *args, **kwargs) -> None:
    _DEFAULT_LOGGER.debug(msg, *args, **kwargs)

Logs debug message.

Args:
        msg: Message with possible format string.
        *args: Values for variables in the format string.
        **kwargs: Keyword arguments for the logger.


def from_control_flow_context_def(context_def, import_scope=None):
    if context_def.HasField('cond_ctxt'):
        return CondContext.from_proto(context_def.cond_ctxt, import_scope=import_scope)
    if context_def.HasField('while_ctxt'):
        return WhileContext.from_proto(context_def.while_ctxt, import_scope=import_scope)
    raise NotImplementedError('Unknown ControlFlowContextDef field: %s' % context_def.WhichOneof('ctxt'))

Deserializes `context_def` into the appropriate ControlFlowContext.

Args:
        context_def: ControlFlowContextDef proto
        import_scope: Optional `string`. Name scope to add.

Returns:
        A ControlFlowContext subclass


def clean_up_tokenization(out_string: str) -> str:
    out_string = out_string.replace(' .', '.').replace(' ?', '?').replace(' !', '!').replace(' ,', ',').replace(" ' ", "'").replace(" n't", "n't").replace(" 'm", "'m").replace(" 's", "'s").replace(" 've", "'ve").replace(" 're", "'re")
    return out_string

Clean up a list of simple English tokenization artifacts like spaces before punctuations and abbreviated forms.

Args:
        out_string (`str`): The text to clean up.

Returns:
        `str`: The cleaned-up string.


def insert(self, keys, values, name=None):
    with ops.name_scope(name, '%s_lookup_table_insert' % self.name, [self.resource_handle, keys, values]):
        keys = ops.convert_to_tensor(keys, self._key_dtype, name='keys')
        values = ops.convert_to_tensor(values, self._value_dtype, name='values')
        with ops.colocate_with(self.resource_handle):
            op = gen_lookup_ops.lookup_table_insert_v2(self.resource_handle, keys, values)
    return op

Associates `keys` with `values`.

Args:
        keys: Keys to insert. Can be a tensor of any shape. Must match the table's
        key type.
        values: Values to be associated with keys. Must be a tensor of the same
        shape as `keys` and match the table's value type.
        name: A name for the operation (optional).

Returns:
        The created Operation.

Raises:
        TypeError: when `keys` or `values` doesn't match the table data
        types.


def _html_tree_view_render(self, *, view: 'HtmlTreeView', name: Optional[str]=None, parent: Any=None, root_path: Optional[KeyPath]=None, **kwargs) -> Html:
    return self._html_tree_view(view=view, name=name, parent=parent, root_path=root_path, **view.get_kwargs(kwargs, self._html_tree_view_config(), root_path or KeyPath())).add_style(*self._html_tree_view_css_styles())

The entrypoint of rendering the subtree represented by this extension.

Args:
        view: The view to render the object.
        name: The name of the object.
        parent: The parent of the object.
        root_path: The key path of the object relative to the root.
        **kwargs: kwargs to pass to `view.render()` on this extension.

Returns:
        The rendered HTML.


def with_hot_key_fanout(self, fanout):
    from apache_beam.transforms.combiners import curry_combine_fn
    if fanout is None:
        return self
    else:
        return _CombinePerKeyWithHotKeyFanout(curry_combine_fn(self.fn, self.args, self.kwargs), fanout)

A per-key combine operation like self but with two levels of aggregation.

    If a given key is produced by too many upstream bundles, the final
    reduction can become a bottleneck despite partial combining being lifted
    pre-GroupByKey.  In these cases it can be helpful to perform intermediate
    partial aggregations in parallel and then re-group to peform a final
    (per-key) combine.  This is also useful for high-volume keys in streaming
    where combiners are not generally lifted for latency reasons.

    Note that a fanout greater than 1 requires the data to be sent through
    two GroupByKeys, and a high fanout can also result in more shuffle data
    due to less per-bundle combining. Setting the fanout for a key at 1 or less
    places values on the "cold key" path that skip the intermediate level of
    aggregation.

Args:
        fanout: either None, for no fanout, an int, for a constant-degree fanout,
        or a callable mapping keys to a key-specific degree of fanout.

Returns:
        A per-key combining PTransform with the specified fanout.


def __init__(self, project, throttle_rampup=True, hint_num_workers=_Mutate._DEFAULT_HINT_NUM_WORKERS):
    mutate_fn = DeleteFromDatastore._DatastoreDeleteFn(project)
    super().__init__(mutate_fn, throttle_rampup, hint_num_workers)

Initialize the `DeleteFromDatastore` transform.

Args:
        project: (:class:`str`) The ID of the project from which the entities will
        be deleted.
        throttle_rampup: Whether to enforce a gradual ramp-up.
        hint_num_workers: A hint for the expected number of workers, used to
                        estimate appropriate limits during ramp-up throttling.


def _transform_features_v2(features, feature_columns, state_manager):
    feature_columns = _normalize_feature_columns(feature_columns)
    outputs = {}
    with ops.name_scope(None, default_name='transform_features', values=features.values()):
        transformation_cache = FeatureTransformationCache(features)
        for column in feature_columns:
            with ops.name_scope(None, default_name=_sanitize_column_name_for_variable_scope(column.name)):
                outputs[column] = transformation_cache.get(column, state_manager)
    return outputs

Returns transformed features based on features columns passed in.

    Please note that most probably you would not need to use this function. Please
    check `input_layer` and `linear_model` to see whether they will
    satisfy your use case or not.

Example:
    ```python
    # Define features and transformations
    crosses_a_x_b = crossed_column(
        columns=["sparse_feature_a", "sparse_feature_b"], hash_bucket_size=10000)
    price_buckets = bucketized_column(
        source_column=numeric_column("price"), boundaries=[...])

    columns = [crosses_a_x_b, price_buckets]
    features = tf.io.parse_example(..., features=make_parse_example_spec(columns))
    transformed = transform_features(features=features, feature_columns=columns)

    assertCountEqual(columns, transformed.keys())
    ```

Args:
        features: A mapping from key to tensors. `FeatureColumn`s look up via these
        keys. For example `numeric_column('price')` will look at 'price' key in
        this dict. Values can be a `SparseTensor` or a `Tensor` depends on
        corresponding `FeatureColumn`.
        feature_columns: An iterable containing all the `FeatureColumn`s.
        state_manager: A StateManager object that holds the FeatureColumn state.

Returns:
        A `dict` mapping `FeatureColumn` to `Tensor` and `SparseTensor` values.


def _maybe_truncate_traceback(traceback):
    if len(traceback) > MAX_TRACEBACK_LENGTH:
        return traceback[:MAX_TRACEBACK_LENGTH - 2] + [_ELLIPSIS, traceback[-1]]
    else:
        return traceback

Truncate the traceback if it is too long.

Args:
        traceback: A list representing an error's traceback. There should be one
        list item per entry in the traceback (in the right order); beyond that,
        this function does not care about the item types.

Returns:
        The traceback, possibly with some items removed and an _ELLIPSIS inserted.
        Guaranteed to be no longer than MAX_TRACEBACK_LENGTH.


def load_default_japanese_parser() -> Parser:
    with open(os.path.join(MODEL_DIR, 'ja.json'), encoding='utf-8') as f:
        model = json.load(f)
    return Parser(model)

Loads a parser equipped with the default Japanese model.

Returns:
        A parser (:obj:`budoux.Parser`).


def _spanner_io_read_test_preprocessor(test_spec: dict, expected: List[str], env: TestEnvironment):
    if (pipeline := test_spec.get('pipeline', None)):
        for transform in pipeline.get('transforms', []):
            if transform.get('type', '').startswith('ReadFromSpanner'):
                config = transform['config']
                instance, database = (config['instance_id'], config['database_id'])
                if (table := (config.get('table', None) is None)):
                    table = config.get('query', '').split('FROM')[-1].strip()
                transform['type'] = 'Create'
                transform['config'] = {k: v for k, v in config.items() if k.startswith('__')}
                elements = INPUT_TABLES[str(instance), str(database), str(table)]
                if config.get('query', None):
                    config['query'].replace('select ', 'SELECT ').replace(' from ', ' FROM ')
                    columns = set(''.join(config['query'].split('SELECT ')[1:]).split(' FROM', maxsplit=1)[0].split(', '))
                    if columns != {'*'}:
                        elements = [{column: element[column] for column in element if column in columns} for element in elements]
                transform['config']['elements'] = elements
    return test_spec

Preprocessor for tests that involve reading from Spanner.

    This preprocessor replaces any ReadFromSpanner transform with a Create
    transform that reads from a predefined in-memory dictionary. This allows
    the test to verify the pipeline's correctness without relying on external
    Spanner instances.

Args:
        test_spec: The dictionary representation of the YAML pipeline specification.
        expected: A list of strings representing the expected output of the
        pipeline.
        env: The TestEnvironment object providing utilities for creating temporary
        files.

Returns:
        The modified test_spec dictionary with ReadFromSpanner transforms replaced.


def GetShadowMap(self, since=None):
    return ShadowUpdateGetter().GetUpdates(self._GetClient(), self.conf['bucket'], self.conf['shadow_object'], since)

Return the shadow map from this source.

Args:
        since: Get data only changed since this timestamp (inclusive) or None
        for all data.

Returns:
        instance of shadow.ShadowMap


def _eval_indexed_slices(a):
    if isinstance(a, indexed_slices.IndexedSlices) and context.executing_eagerly():
        return indexed_slices.IndexedSlicesValue(indices=[x.numpy() for x in a.indices], values=[x.numpy() for x in a.values], dense_shape=a.dense_shape)
    return a

Converts IndexedSlices to IndexedSlicesValue with numpy indices/values.

    When eager execution is enabled, converts IndexedSlices
    to IndexedSlicesValue with numpy indices/values.

Args:
        a: any value.

Returns:
        If a is IndexedSlices and eager execution is enabled, calls numpy() on a's
        fields. Otherwise returns a unchanged.


def _overload_operator(cls, tensor_class, operator):
    tensor_oper = getattr(tensor_class, operator)
    tensor_oper = getattr(tensor_oper, '__func__', tensor_oper)
    setattr(cls, operator, tensor_oper)

Overload an operator with the same implementation as a base Tensor class.

    We pull the operator out of the class dynamically to avoid ordering issues.

Args:
        tensor_class: The (Composite)Tensor to get the method from.
        operator: string. The operator name.


def _strip_blank_lines(lines):
    start = 0
    num_lines = len(lines)
    while lines and start < num_lines and _is_blank(lines[start]):
        start += 1
    lines = lines[start:]
    while lines and _is_blank(lines[-1]):
        lines.pop()
    return lines

Removes lines containing only blank characters before and after the text.

Args:
        lines: A list of lines.

Returns:
        A list of lines without trailing or leading blank lines.


def __init__(self, output_bit, quant_mode=False, force_dequant='none'):
    super().__init__()
    self.output_bit = output_bit
    self.max_bit = 32
    self.quant_mode = quant_mode
    if force_dequant in ['nonlinear', 'softmax']:
        logger.info('Force dequantize softmax')
        self.quant_mode = False
    self.act = QuantAct(16, quant_mode=self.quant_mode)
    self.x0 = -0.6931
    self.const = 30
    self.coef = [0.35815147, 0.96963238, 1.0]
    self.coef[1] /= self.coef[0]
    self.coef[2] /= self.coef[0]

Quantized version of `torch.nn.Softmax`. Adds quantization-specific arguments on top of `torch.nn.Softmax`.

Args:
        output_bit (`int`):
            Bitwidth for the layer output activation.
        quant_mode (`bool`, *optional*, defaults to `False`):
            Whether or not the layer is quantized.
        force_dequant (`str`, *optional*, defaults to `"none"`):
            Force dequantize the layer if either "softmax" or "nonlinear" is given.


def recipe_drive_copy(config, auth_read, source, destination):
    drive(config, {'auth': auth_read, 'copy': {'source': source, 'destination': destination}})

Copy a drive document.

Args:
        auth_read (authentication) - Credentials used for reading data.
        source (string) - Name or URL of document to copy from.
        destination (string) - Name document to copy to.


def flat_values(self):
    rt_values = self.values
    while isinstance(rt_values, RaggedTensor):
        rt_values = rt_values.values
    return rt_values

The innermost `values` tensor for this ragged tensor.

    Concretely, if `rt.values` is a `Tensor`, then `rt.flat_values` is
    `rt.values`; otherwise, `rt.flat_values` is `rt.values.flat_values`.

    Conceptually, `flat_values` is the tensor formed by flattening the
    outermost dimension and all of the ragged dimensions into a single
    dimension.

    `rt.flat_values.shape = [nvals] + rt.shape[rt.ragged_rank + 1:]`
    (where `nvals` is the number of items in the flattened dimensions).

Returns:
        A `Tensor`.

    #### Example:

    >>> rt = tf.ragged.constant([[[3, 1, 4, 1], [], [5, 9, 2]], [], [[6], []]])
    >>> print(rt.flat_values)
    tf.Tensor([3 1 4 1 5 9 2 6], shape=(8,), dtype=int32)


def _single_column_cell_selection_loss(token_logits, column_logits, labels, cell_index, col_index, cell_mask):
    labels_per_column, _ = reduce_sum(tf.cast(labels, tf.float32), col_index)
    column_label = tf.argmax(labels_per_column, axis=-1, output_type=tf.int32)
    no_cell_selected = tf.equal(tf.reduce_max(labels_per_column, axis=-1), 0)
    column_label = tf.where(no_cell_selected, tf.zeros_like(column_label), column_label)
    column_dist = tfp.distributions.Categorical(logits=column_logits)
    column_loss_per_example = -column_dist.log_prob(column_label)
    logits_per_cell, _ = reduce_mean(token_logits, cell_index)
    labels_per_cell, labels_index = reduce_max(tf.cast(labels, tf.int32), cell_index)
    column_id_for_cells = cell_index.project_inner(labels_index).indices
    column_mask = tf.cast(tf.equal(column_id_for_cells, tf.expand_dims(column_label, axis=1)), tf.float32)
    cell_dist = tfp.distributions.Bernoulli(logits=logits_per_cell)
    cell_log_prob = cell_dist.log_prob(labels_per_cell)
    cell_loss = -tf.reduce_sum(cell_log_prob * column_mask * cell_mask, axis=1)
    cell_loss /= tf.reduce_sum(column_mask * cell_mask, axis=1) + EPSILON_ZERO_DIVISION
    selection_loss_per_example = column_loss_per_example
    selection_loss_per_example += tf.where(no_cell_selected, tf.zeros_like(selection_loss_per_example), cell_loss)
    selected_column_id = tf.argmax(column_logits, axis=-1, output_type=tf.int32)
    selected_column_mask = tf.cast(tf.equal(column_id_for_cells, tf.expand_dims(selected_column_id, axis=-1)), tf.float32)
    selected_column_mask = tf.where(tf.equal(column_id_for_cells, 0), tf.zeros_like(selected_column_mask), selected_column_mask)
    logits_per_cell += CLOSE_ENOUGH_TO_LOG_ZERO * (1.0 - cell_mask * selected_column_mask)
    logits = gather(logits_per_cell, cell_index)
    return (selection_loss_per_example, logits)

Computes the loss for cell selection constrained to a single column. The loss is a hierarchical log-likelihood. The
    model first predicts a column and then selects cells within that column (conditioned on the column). Cells outside
    the selected column are never selected.

Args:
        token_logits (`tf.Tensor` of shape `(batch_size, sequence_length)`):
            Tensor containing the logits per token.
        column_logits (`tf.Tensor` of shape `(batch_size, max_num_cols)`):
            Tensor containing the logits per column.
        labels (`tf.Tensor` of shape `(batch_size, sequence_length)`):
            Labels per token.
        cell_index (`ProductIndexMap`):
            Index that groups tokens into cells.
        col_index (`IndexMap`):
            Index that groups tokens into columns.
        cell_mask (`tf.Tensor` of shape `(batch_size, max_num_rows * max_num_cols)`):
            Mask for cells that exist in the table (i.e. that are not padding).

Returns:
        selection_loss_per_example (`tf.Tensor` of shape `(batch_size,)`): Loss for each example. logits (`tf.Tensor`
        of shape `(batch_size, sequence_length)`): New logits which are only allowed to select cells in a single
        column. Logits outside of the most likely column according to *column_logits* will be set to a very low value
        (such that the probabilities are 0).


def with_metadata_spec(self, column_name: str='metadata', python_type: Type=str, convert_fn: Optional[Callable[[Dict[str, Any]], Any]]=None, sql_typecast: Optional[str]='::jsonb') -> 'ColumnSpecsBuilder':

    def value_fn(chunk: Chunk) -> Any:
        if convert_fn:
            return convert_fn(chunk.metadata)
        return json.dumps(chunk.metadata) if python_type == str else chunk.metadata
    self._specs.append(ColumnSpec(column_name=column_name, python_type=python_type, value_fn=value_fn, sql_typecast=sql_typecast))
    return self

Add metadata :class:`.ColumnSpec` with optional type and conversion.

Args:
        column_name: Name for the metadata column (defaults to "metadata")
        python_type: Python type for the column (defaults to str)
        convert_fn: Optional function to convert the metadata dictionary
                    If None and python_type is str, converts to JSON string
        sql_typecast: Optional SQL type cast (defaults to "::jsonb")

Returns:
        Self for method chaining

Example:
        >>> builder.with_metadata_spec(
        ...     column_name="meta_tags",
        ...     python_type=list,
        ...     convert_fn=lambda meta: list(meta.keys()),
        ...     sql_typecast="::text[]"
        ... )


def __init__(self, resource_handle, create_op, name):
    stamp_token, serialized = gen_boosted_trees_ops.boosted_trees_serialize_ensemble(resource_handle)
    slice_spec = ''
    specs = [saver.BaseSaverBuilder.SaveSpec(stamp_token, slice_spec, name + '_stamp'), saver.BaseSaverBuilder.SaveSpec(serialized, slice_spec, name + '_serialized')]
    super(_TreeEnsembleSavable, self).__init__(resource_handle, specs, name)
    self.resource_handle = resource_handle
    self._create_op = create_op

Creates a _TreeEnsembleSavable object.

Args:
        resource_handle: handle to the decision tree ensemble variable.
        create_op: the op to initialize the variable.
        name: the name to save the tree ensemble variable under.


def basis(sample_paths, time_index):
    sample_paths = tf.convert_to_tensor(sample_paths, name='sample_paths')
    if sample_paths.shape.rank == 3:
        sample_paths = tf.expand_dims(sample_paths, axis=0)
    shape = tf.shape(sample_paths)
    num_samples = shape[1]
    batch_size = shape[0]
    dim = sample_paths.shape[-1]
    slice_samples = tf.slice(sample_paths, [0, 0, time_index, 0], [batch_size, num_samples, 1, dim])
    samples_centered = slice_samples - tf.math.reduce_mean(slice_samples, axis=1, keepdims=True)
    grid = tf.range(degree + 1, dtype=samples_centered.dtype)
    grid = tf.meshgrid(*dim * [grid])
    grid = tf.reshape(tf.stack(grid, -1), [-1, dim])
    basis_expansion = tf.reduce_prod(samples_centered ** grid, axis=-1)
    return tf.transpose(basis_expansion, [0, 2, 1])

Computes polynomial basis expansion at the given sample points.

Args:
        sample_paths: A `Tensor` of either `flaot32` or `float64` dtype and of
        either shape `[num_samples, num_times, dim]` or
        `[batch_size, num_samples, num_times, dim]`.
        time_index: An integer scalar `Tensor` that corresponds to the time
        coordinate at which the basis function is computed.

Returns:
        A `Tensor`s of shape `[batch_size, (degree + 1)**dim, num_samples]`.


def _find_max_under_constraint(self, constrained, dependent, predicate):
    feasible = array_ops.where_v2(predicate(constrained, self.value))
    feasible_exists = math_ops.greater(array_ops.size(feasible), 0)
    max_dependent = math_ops.reduce_max(array_ops.gather(dependent, feasible))
    return array_ops.where_v2(feasible_exists, max_dependent, 0.0)

Returns the maximum of dependent_statistic that satisfies the constraint.

Args:
        constrained: Over these values the constraint
        is specified. A rank-1 tensor.
        dependent: From these values the maximum that satiesfies the
        constraint is selected. Values in this tensor and in
        `constrained` are linked by having the same threshold at each
        position, hence this tensor must have the same shape.
        predicate: A binary boolean functor to be applied to arguments
        `constrained` and `self.value`, e.g. `tf.greater`.

    Returns maximal dependent value, if no value satiesfies the constraint 0.0.


def contains_saved_model(export_dir):
    if isinstance(export_dir, os.PathLike):
        export_dir = os.fspath(export_dir)
    return maybe_saved_model_directory(export_dir)

Checks whether the provided export directory could contain a SavedModel.

    Note that the method does not load any data by itself. If the method returns
    `false`, the export directory definitely does not contain a SavedModel. If the
    method returns `true`, the export directory may contain a SavedModel but
    provides no guarantee that it can be loaded.

Args:
        export_dir: Absolute path to possible export location. For example,
                '/my/foo/model'.

Returns:
        True if the export directory contains SavedModel files, False otherwise.


def check_annotation_type_mismatch(self, node, name, typ, value, stack, allow_none, details=None):
    if not typ or not value:
        return
    if value.data == [self.convert.ellipsis] or (allow_none and value.data == [self.convert.none]):
        return
    contained_type = abstract_utils.match_type_container(typ, ('typing.ClassVar', 'dataclasses.InitVar'))
    if contained_type:
        typ = contained_type
    bad = self.matcher(node).compute_one_match(value, typ).bad_matches
    for match in bad:
        self.errorlog.annotation_type_mismatch(stack, match.expected.typ, match.actual_binding, name, match.error_details, details)

Checks for a mismatch between a variable's annotation and value.

Args:
        node: node
        name: variable name
        typ: variable annotation
        value: variable value
        stack: a frame stack for error reporting
        allow_none: whether a value of None is allowed for any type
        details: any additional details to add to the error message


def update_shapes_dict_for_target_fn(target_fn, shapes_dict, call_spec, class_name):
    if utils.is_default(target_fn):
        return None
    sig = inspect.signature(target_fn)
    expected_names = []
    for name, param in sig.parameters.items():
        if param.kind in (param.POSITIONAL_OR_KEYWORD, param.POSITIONAL_ONLY, param.KEYWORD_ONLY):
            expected_names.append(name)
    if len(expected_names) == 1:
        key = expected_names[0]
        values = tuple(shapes_dict.values())
        if values:
            input_shape = values[0]
        else:
            input_shape = None
        return {key: input_shape}
    kwargs = {}
    for name in expected_names:
        method_name = target_fn.__name__
        error_preamble = f'For a `{method_name}()` method with more than one argument, all arguments should have a `_shape` suffix and match an argument from `call()`. E.g. `{method_name}(self, foo_shape, bar_shape)` '
        if not name.endswith('_shape'):
            raise ValueError(f"{error_preamble} For layer '{class_name}', Received `{method_name}()` argument `{name}`, which does not end in `_shape`.")
        expected_call_arg = utils.removesuffix(name, '_shape')
        if expected_call_arg not in call_spec.arguments_dict:
            raise ValueError(f"{error_preamble} For layer '{class_name}', received `{method_name}()` argument `{name}`, but `call()` does not have argument `{expected_call_arg}`.")
        if name in shapes_dict:
            kwargs[name] = shapes_dict[name]
    return kwargs

Updates a `shapes_dict` for `build()` or `compute_output_shape()`.

    This function will align a dictionary of the shapes of all tensor
    passed to `call`, with the signatures of `build()` or
    `compute_output_shape()`.

    The alignment is a follows:

    - If `build()` or `compute_output_shape()` accept only one argument,
        forward the shape of the first positional argument from call without
        checking any argument names.
    - If `build()` or `compute_output_shape()` accept multiple arguments,
        enforce that all argument names match a call argument name, e.g.
        `foo_shape` would match call argument `foo`.

Returns:
        An updated `shapes_dict` that can be used to invoke
        `target_fn(**shapes_dict)`.


def get_traverse_shallow_structure(traverse_fn, structure, expand_composites=False):
    is_nested_fn = _is_nested_or_composite if expand_composites else is_nested
    to_traverse = traverse_fn(structure)
    if not is_nested_fn(structure):
        if not isinstance(to_traverse, bool):
            raise TypeError('traverse_fn returned structure: %s for non-structure: %s' % (to_traverse, structure))
        return to_traverse
    level_traverse = []
    if isinstance(to_traverse, bool):
        if not to_traverse:
            return False
        else:
            for branch in nest_util.yield_value(nest_util.Modality.CORE, structure):
                level_traverse.append(get_traverse_shallow_structure(traverse_fn, branch, expand_composites=expand_composites))
    elif not is_nested_fn(to_traverse):
        raise TypeError('traverse_fn returned a non-bool scalar: %s for input: %s' % (to_traverse, structure))
    else:
        assert_shallow_structure(to_traverse, structure, expand_composites=expand_composites)
        for t, branch in zip(nest_util.yield_value(nest_util.Modality.CORE, to_traverse), nest_util.yield_value(nest_util.Modality.CORE, structure)):
            if not isinstance(t, bool):
                raise TypeError("traverse_fn didn't return a depth=1 structure of bools.  saw: %s  for structure: %s" % (to_traverse, structure))
            if t:
                level_traverse.append(get_traverse_shallow_structure(traverse_fn, branch))
            else:
                level_traverse.append(False)
    return nest_util.sequence_like(structure, level_traverse)

Generates a shallow structure from a `traverse_fn` and `structure`.

    `traverse_fn` must accept any possible subtree of `structure` and return
    a depth=1 structure containing `True` or `False` values, describing which
    of the top-level subtrees may be traversed.  It may also
    return scalar `True` or `False` "traversal is OK / not OK for all subtrees."

    Examples are available in the unit tests (nest_test.py).

Args:
        traverse_fn: Function taking a substructure and returning either a scalar
        `bool` (whether to traverse that substructure or not) or a depth=1
        shallow structure of the same type, describing which parts of the
        substructure to traverse.
        structure: The structure to traverse.
        expand_composites: If true, then composite tensors such as
        `tf.sparse.SparseTensor` and `tf.RaggedTensor` are expanded into their
        component tensors.

Returns:
        A shallow structure containing python bools, which can be passed to
        `map_structure_up_to` and `flatten_up_to`.

Raises:
        TypeError: if `traverse_fn` returns a nested structure for an atom input.
        or a structure with depth higher than 1 for a nested structure input,
        or if any leaf values in the returned structure or scalar are not type
        `bool`.


def get_bq_tableschema(schema):
    if isinstance(schema, (bigquery.TableSchema, value_provider.ValueProvider)) or callable(schema) or schema is None:
        return schema
    elif isinstance(schema, str):
        return get_table_schema_from_string(schema)
    elif isinstance(schema, dict):
        schema_string = json.dumps(schema)
        return parse_table_schema_from_json(schema_string)
    else:
        raise TypeError('Unexpected schema argument: %s.' % schema)

Convert the table schema to a TableSchema object.

Args:
        schema (str, dict, ~apache_beam.io.gcp.internal.clients.bigquery.bigquery_v2_messages.TableSchema):
        The schema to be used if the BigQuery table to write has to be created.
        This can either be a dict or string or in the TableSchema format.

Returns:
        ~apache_beam.io.gcp.internal.clients.bigquery.bigquery_v2_messages.TableSchema: The schema as a TableSchema object.


def _assert_sparse_compatible(sparse_tensors):
    checks = []
    first = sparse_tensors[0]
    for t in sparse_tensors[1:]:
        checks.append(check_ops.assert_equal(first.dense_shape, t.dense_shape, message='Mismatched shapes!'))
        checks.append(check_ops.assert_equal(first.indices, t.indices, message='Mismatched indices!'))
    return checks

Check that all of `sparse_tensors` have same `indices` and `dense_shape`.

Args:
        sparse_tensors: A list of sparse tensors.

Returns:
        An op to be used as a control dependency.


def nested_value_rowids(self, name=None):
    with ops.name_scope(name, 'RaggedNestedValueRowIds', [self]):
        rt_nested_ids = [self.value_rowids()]
        rt_values = self.values
        while isinstance(rt_values, RaggedTensor):
            rt_nested_ids.append(rt_values.value_rowids())
            rt_values = rt_values.values
        return tuple(rt_nested_ids)

Returns a tuple containing the value_rowids for all ragged dimensions.

    `rt.nested_value_rowids` is a tuple containing the `value_rowids` tensors
    for
    all ragged dimensions in `rt`, ordered from outermost to innermost.  In
    particular, `rt.nested_value_rowids = (rt.value_rowids(),) + value_ids`
    where:

    * `value_ids = ()` if `rt.values` is a `Tensor`.
    * `value_ids = rt.values.nested_value_rowids` otherwise.

Args:
        name: A name prefix for the returned tensors (optional).

Returns:
        A `tuple` of 1-D integer `Tensor`s.

    #### Example:

    >>> rt = tf.ragged.constant(
    ...     [[[[3, 1, 4, 1], [], [5, 9, 2]], [], [[6], []]]])
    >>> for i, ids in enumerate(rt.nested_value_rowids()):
    ...   print('row ids for dimension %d: %s' % (i+1, ids.numpy()))
    row ids for dimension 1: [0 0 0]
    row ids for dimension 2: [0 0 0 2 2]
    row ids for dimension 3: [0 0 0 0 2 2 2 3]


def rfft(x, fft_length=None):
    if any_symbolic_tensors((x,)):
        return RFFT(fft_length).symbolic_call(x)
    return backend.math.rfft(x, fft_length)

Real-valued Fast Fourier Transform along the last axis of the input.

    Computes the 1D Discrete Fourier Transform of a real-valued signal over the
    inner-most dimension of input.

    Since the Discrete Fourier Transform of a real-valued signal is
    Hermitian-symmetric, RFFT only returns the `fft_length / 2 + 1` unique
    components of the FFT: the zero-frequency term, followed by the
    `fft_length / 2` positive-frequency terms.

    Along the axis RFFT is computed on, if `fft_length` is smaller than the
    corresponding dimension of the input, the dimension is cropped. If it is
    larger, the dimension is padded with zeros.

Args:
        x: Input tensor.
        fft_length: An integer representing the number of the fft length. If not
            specified, it is inferred from the length of the last axis of `x`.
            Defaults to `None`.

Returns:
        A tuple containing two tensors - the real and imaginary parts of the
        output.

Example:
    >>> x = keras.ops.convert_to_tensor([0.0, 1.0, 2.0, 3.0, 4.0])
    >>> rfft(x)
    (array([10.0, -2.5, -2.5]), array([0.0, 3.4409548, 0.81229924]))

    >>> rfft(x, 3)
    (array([3.0, -1.5]), array([0.0, 0.8660254]))


def isnan(x):
    if any_symbolic_tensors((x,)):
        return Isnan().symbolic_call(x)
    return backend.numpy.isnan(x)

Test element-wise for NaN and return result as a boolean tensor.

Args:
        x: Input tensor.

Returns:
        Output boolean tensor.


def __init__(self, namespace: Optional[str], name: Optional[str], urn: Optional[str]=None, labels: Optional[Dict[str, str]]=None) -> None:
    if not urn:
        if not namespace:
            raise ValueError('Metric namespace must be non-empty')
        if not name:
            raise ValueError('Metric name must be non-empty')
    self.namespace = namespace
    self.name = name
    self.urn = urn
    self.labels = labels if labels else {}

Initializes ``MetricName``.

    Note: namespace and name should be set for user metrics,
    urn and labels should be set for an arbitrary metric to package into a
    MonitoringInfo.

Args:
        namespace: A string with the namespace of a metric.
        name: A string with the name of a metric.
        urn: URN to populate on a MonitoringInfo, when sending to RunnerHarness.
        labels: Labels to populate on a MonitoringInfo


def _sequence_like(instance, args):
    return nest_util.sequence_like(instance, args)

Converts the sequence `args` to the same type as `instance`.

Args:
        instance: an instance of `tuple`, `list`, `namedtuple`, `dict`,
        `collections.OrderedDict`, or `composite_tensor.Composite_Tensor`
        or `type_spec.TypeSpec`.
        args: items to be converted to the `instance` type.

Returns:
        `args` with the type of `instance`.


def solvevec(self, rhs, adjoint=False, name='solve'):
    with self._name_scope(name):
        rhs = tensor_conversion.convert_to_tensor_v2_with_dispatch(rhs, name='rhs')
        self._check_input_dtype(rhs)
        self_dim = -1 if adjoint else -2
        tensor_shape.dimension_at_index(self.shape, self_dim).assert_is_compatible_with(rhs.shape[-1])
        return self._solvevec(rhs, adjoint=adjoint)

Solve single equation with best effort: `A X = rhs`.

    The returned `Tensor` will be close to an exact solution if `A` is well
    conditioned. Otherwise closeness will vary. See class docstring for details.

Example:
    ```python
    # Make an operator acting like batch matrix A.  Assume A.shape = [..., M, N]
    operator = LinearOperator(...)
    operator.shape = [..., M, N]

    # Solve one linear system for every member of the batch.
    RHS = ... # shape [..., M]

    X = operator.solvevec(RHS)
    # X is the solution to the linear system
    # sum_j A[..., :, j] X[..., j] = RHS[..., :]

    operator.matvec(X)
    ==> RHS
    ```

Args:
        rhs: `Tensor` with same `dtype` as this operator.
        `rhs` is treated like a [batch] vector meaning for every set of leading
        dimensions, the last dimension defines a vector.  See class docstring
        for definition of compatibility regarding batch dimensions.
        adjoint: Python `bool`.  If `True`, solve the system involving the adjoint
        of this `LinearOperator`:  `A^H X = rhs`.
        name:  A name scope to use for ops added by this method.

Returns:
        `Tensor` with shape `[...,N]` and same `dtype` as `rhs`.

Raises:
        NotImplementedError:  If `self.is_non_singular` or `is_square` is False.


def function(inputs, outputs, updates=None, name=None, **kwargs):
    if ops.executing_eagerly_outside_functions():
        if kwargs:
            raise ValueError('Session keyword arguments are not supported during eager execution. You passed: %s' % (kwargs,))
        if updates:
            raise ValueError('`updates` argument is not supported during eager execution. You passed: %s' % (updates,))
        from tensorflow.python.keras import models
        from tensorflow.python.keras.utils import tf_utils
        model = models.Model(inputs=inputs, outputs=outputs)
        wrap_outputs = isinstance(outputs, list) and len(outputs) == 1

        def func(model_inputs):
            outs = model(model_inputs)
            if wrap_outputs:
                outs = [outs]
            return tf_utils.sync_to_numpy_or_python_type(outs)
        return func
    if kwargs:
        for key in kwargs:
            if key not in tf_inspect.getfullargspec(session_module.Session.run)[0] and key not in ['inputs', 'outputs', 'updates', 'name']:
                msg = 'Invalid argument "%s" passed to K.function with TensorFlow backend' % key
                raise ValueError(msg)
    return GraphExecutionFunction(inputs, outputs, updates=updates, name=name, **kwargs)

Instantiates a Keras function.

Args:
        inputs: List of placeholder tensors.
        outputs: List of output tensors.
        updates: List of update ops.
        name: String, name of function.
        **kwargs: Passed to `tf.Session.run`.

Returns:
        Output values as Numpy arrays.

Raises:
        ValueError: if invalid kwargs are passed in or if in eager execution.


def check_line_split(code_line):
    return re.search('\\\\\\s*\\n$', code_line)

Checks if a line was split with `\`.

Args:
        code_line: A line of Python code

Returns:
        If the line was split with `\`

    >>> skip_magic("!gcloud ml-engine models create ${MODEL} \\\n")
    True


def preprocess_input_examples_arg_string(input_examples_str):
    input_dict = preprocess_input_exprs_arg_string(input_examples_str)
    for input_key, example_list in input_dict.items():
        if not isinstance(example_list, list):
            raise ValueError('tf.Example input must be a list of dictionaries, but "%s" is %s' % (example_list, type(example_list)))
        input_dict[input_key] = [_create_example_string(example) for example in example_list]
    return input_dict

Parses input into dict that maps input keys to lists of tf.Example.

    Parses input string in the format of 'input_key1=[{feature_name:
    feature_list}];input_key2=[{feature_name:feature_list}];' into a dictionary
    that maps each input_key to its list of serialized tf.Example.

Args:
        input_examples_str: A string that specifies a list of dictionaries of
        feature_names and their feature_lists for each input.
        Each input is separated by semicolon. For each input key:
        'input=[{feature_name1: feature_list1, feature_name2:feature_list2}]'
        items in feature_list can be the type of float, int, long or str.

Returns:
        A dictionary that maps input keys to lists of serialized tf.Example.

Raises:
        ValueError: An error when the given tf.Example is not a list.


def matches_any(patterns: List[Pattern[str]], line: str) -> bool:
    stripped_line = line.strip()
    for pattern in patterns:
        if pattern.match(stripped_line):
            return True
    return False

Checks if the line matches any of the given patterns.

Args:
        patterns: A list of compiled regular expression patterns.
        line: The line to check for matches.

Returns:
        True if the line matches any of the patterns, False otherwise.


def _channel_flatten_input(x, data_format):
    graph = ops.get_default_graph()
    cache_key = (graph, x.ref(), data_format)
    if cache_key not in _channel_flatten_input_cache:
        x_shape = array_ops.shape(x)
        neg_ones = constant_op.constant([-1], dtype=x_shape.dtype)
        if data_format == b'NCHW':
            order = [1, 0, 2, 3, 4]
            shape = array_ops.concat([x_shape[1:2], neg_ones, x_shape[3:]], axis=0)
            reverse_order = order
        else:
            order = [1, 2, 3, 0, 4]
            shape = array_ops.concat([x_shape[1:4], neg_ones], axis=0)
            reverse_order = [3, 0, 1, 2, 4]
        x = array_ops.transpose(x, order)
        reverse_shape = array_ops.shape(x)
        x = array_ops.reshape(x, shape)
        outputs = (x, reverse_order, reverse_shape)
        _channel_flatten_input_cache[cache_key] = outputs
    else:
        outputs = _channel_flatten_input_cache[cache_key]
    return outputs

Merge the stack dimension with the channel dimension.

    If S is pfor's stacking dimension, then,
        - for SNCHW, we transpose to NSCHW. If N dimension has size 1, the transpose
        should be cheap.
        - for SNHWC, we transpose to NHWSC.
    We then merge the S and C dimension.

Args:
        x: tensor_lib.Tensor to transform.
        data_format: "NCHW" or "NHWC".

Returns:
        A 3-element tuple with the transformed value, along with the shape for
        reshape and order for transpose required to transform back.


def extrapolation_step():

    def step_fn(time, next_time, coord_grid, value_grid, boundary_conditions, second_order_coeff_fn, first_order_coeff_fn, zeroth_order_coeff_fn, inner_second_order_coeff_fn, inner_first_order_coeff_fn, num_steps_performed, dtype=None, name=None):
        """Performs the step."""
        del num_steps_performed
        name = name or 'extrapolation_step'
        return parabolic_equation_step(time, next_time, coord_grid, value_grid, boundary_conditions, second_order_coeff_fn, first_order_coeff_fn, zeroth_order_coeff_fn, inner_second_order_coeff_fn, inner_first_order_coeff_fn, time_marching_scheme=extrapolation_scheme, dtype=dtype, name=name)
    return step_fn

Creates a stepper function with Extrapolation time marching scheme.

    Extrapolation scheme combines two half-steps and the full time step to obtain
    desirable properties. See more details below in `extrapolation_scheme`.

    It is slower than Crank-Nicolson scheme, but deals better with value grids
    that have discontinuities. Consider also `oscillation_damped_crank_nicolson`,
    an efficient combination of Crank-Nicolson and Extrapolation schemes.

Returns:
        Callable to be used in finite-difference PDE solvers (see fd_solvers.py).


def multiply_no_nan(x, y, name=None):
    with ops.name_scope(name, 'multiply_no_nan', [x, y]) as name:
        x = ops.convert_to_tensor(x, name='x')
        y = ops.convert_to_tensor(y, name='y', dtype=x.dtype.base_dtype)
        x_dtype = x.dtype.base_dtype
        y_dtype = y.dtype.base_dtype
        if x_dtype != y_dtype:
            raise TypeError(f'`x` and `y` must have the same dtype, got {x_dtype!r} != {y_dtype!r}')
        return gen_math_ops.mul_no_nan(x, y, name=name)

Computes the product of x and y and returns 0 if the y is zero, even if x is NaN or infinite.

    Note this is noncommutative: if y is NaN or infinite and x is 0, the result
    will be NaN.

Args:
        x: A `Tensor`. Must be one of the following types: `float32`, `float64`.
        y: A `Tensor` whose dtype is compatible with `x`.
        name: A name for the operation (optional).

Returns:
        The element-wise value of the x times y.


def __init__(self, image_processor=None, tokenizer=None, **kwargs):
    feature_extractor = None
    if 'feature_extractor' in kwargs:
        warnings.warn('The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor` instead.', FutureWarning)
        feature_extractor = kwargs.pop('feature_extractor')
    image_processor = image_processor if image_processor is not None else feature_extractor
    if image_processor is None:
        raise ValueError('You need to specify an `image_processor`.')
    if tokenizer is None:
        raise ValueError('You need to specify a `tokenizer`.')
    super().__init__(image_processor, tokenizer)
    self.current_processor = self.image_processor

Constructs a FLAVA processor which wraps a FLAVA image processor and a FLAVA tokenizer into a single processor.

    [`FlavaProcessor`] offers all the functionalities of [`FlavaImageProcessor`] and [`BertTokenizerFast`]. See the
    [`~FlavaProcessor.__call__`] and [`~FlavaProcessor.decode`] for more information.

Args:
        image_processor ([`FlavaImageProcessor`], *optional*): The image processor is a required input.
        tokenizer ([`BertTokenizerFast`], *optional*): The tokenizer is a required input.


def to_struct(self, from_api: dict=None, from_json: dict=None, indent: int=2) -> str:
    if from_api:
        from_json = self.to_json(from_api=from_api)
    fields = []
    spaces = ' ' * indent
    for key, value in from_json.items():
        if not isinstance(value, dict):
            continue
        if value.get('type', 'record') == 'record':
            fields.append('%sSTRUCT(\n%s\n%s) AS %s' % (spaces, self.to_struct(from_json=value, indent=indent + 2), spaces, key))
        elif value['type'] == 'array':
            if 'enum' in value['items']:
                fields.append('%s[%s\n%s] AS %s' % (spaces, 'STRING', spaces, key))
            else:
                fields.append('%s[STRUCT(\n%s\n%s)] AS %s' % (spaces, self.to_struct(from_json=value['items'], indent=indent + 2), spaces, key))
        else:
            fields.append('%s%s AS %s' % (spaces, value['type'].upper(), key))
    return ',\n'.join(fields)

Translates a Discovery API Document schema to a BigQuery STRUCT.

    Recursively crawls the discovery document reference tree to build struct.
    Leverages recursion depth passed in constructor to stop if necessary.

Args:
        from_api: the api schema to extrapolate
        from_json: new object with references replaced, not passed by caller
        parents: used to track recursion depth for a specific schema branch

Returns:
        A BigQuery STRUCT object that can be pasted into a query.


def constant_value(pred):
    if isinstance(pred, tensor.Tensor):
        return tensor_util.constant_value(pred)
    if pred in {0, 1}:
        return bool(pred)
    if isinstance(pred, bool):
        return pred
    if isinstance(pred, variables.Variable):
        return None
    raise TypeError('`pred` must be a Tensor, or a Python bool, or 1 or 0. Found instead: %s' % type(pred))

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

Args:
        pred: A scalar, either a Python bool or a TensorFlow boolean variable
        or tensor, or the Python integer 1 or 0.

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

Raises:
        TypeError: If `pred` is not a Variable, Tensor or bool, or Python
        integer 1 or 0.


def WorkerMessages(self, request, global_params=None):
    config = self.GetMethodConfig('WorkerMessages')
    return self._RunMethod(config, request, global_params=global_params)

Send a worker_message to the service.

Args:
        request: (DataflowProjectsWorkerMessagesRequest) input message
        global_params: (StandardQueryParameters, default: None) global arguments

Returns:
        (SendWorkerMessagesResponse) The response message.


def coerce(lhs: FhirPathDataType, rhs: FhirPathDataType) -> FhirPathDataType:
    if not is_coercible(lhs, rhs):
        raise TypeError(f'Unsupported Standard SQL coercion between {lhs} and {rhs}.')
    if isinstance(rhs, _Any) or isinstance(lhs, _Any):
        return _Any
    if rhs in lhs.supported_coercion:
        return rhs
    else:
        return lhs

Performs implicit type coercion between two datatypes.

    See more at: https://hl7.org/fhirpath/#conversion.

Args:
        lhs: The left operand.
        rhs: The right operand.

Returns:
        The resulting coerced datatype, if successful.

Raises:
        TypeError: In the event that coercion is not supported.
        ValueError: In the event that a coercion cycle is detected.


def embedding_feature(self):
    return HardwareFeature._embedding_feature_proto_to_string(self.tpu_hardware_feature_proto.embedding_feature)

TPU embedding feature.

Returns:
        An EmbeddingFeature enum.


def __init__(self, initializer, num_oov_buckets, lookup_key_dtype=None, name=None, experimental_is_anonymous=False):
    if num_oov_buckets <= 0:
        raise ValueError('`num_oov_buckets` must be > 0; use StaticHashTable.')
    if name:
        name = name.rstrip('/')
    if initializer:
        if lookup_key_dtype is None:
            lookup_key_dtype = initializer.key_dtype
        supported_table_key_dtypes = (dtypes.int64, dtypes.string)
        if initializer.key_dtype not in supported_table_key_dtypes:
            raise TypeError('Invalid `key_dtype`, expected one of %s, but got %s.' % (supported_table_key_dtypes, initializer.key_dtype))
        if initializer.key_dtype.is_integer != lookup_key_dtype.is_integer:
            raise TypeError('Invalid `key_dtype`, expected %s but got %s.' % ('integer' if lookup_key_dtype.is_integer else 'non-integer', initializer.key_dtype))
        if initializer.value_dtype != dtypes.int64:
            raise TypeError('Invalid `value_dtype`, expected %s but got %s.' % (dtypes.int64, initializer.value_dtype))
        if isinstance(initializer, trackable_base.Trackable):
            self._initializer = self._track_trackable(initializer, '_initializer')
        self._table = HashTable(initializer, default_value=-1, experimental_is_anonymous=experimental_is_anonymous)
        name = name or self._table.name
    else:
        lookup_key_dtype = dtypes.string
        self._table = None
        name = name or 'hash_bucket'
    if not lookup_key_dtype.is_integer and dtypes.string != lookup_key_dtype:
        raise TypeError(f'Invalid `key_dtype`, expected integer or string, got {lookup_key_dtype}')
    self._num_oov_buckets = num_oov_buckets
    self._table_name = None
    if name is not None:
        self._table_name = name.split('/')[-1]
    super(StaticVocabularyTable, self).__init__(lookup_key_dtype, dtypes.int64)

Construct a `StaticVocabularyTable` object.

Args:
        initializer: A `TableInitializerBase` object that contains the data used
        to initialize the table. If None, then we only use out-of-vocab buckets.
        num_oov_buckets: Number of buckets to use for out-of-vocabulary keys. Must
        be greater than zero. If out-of-vocab buckets are not required, use
        `StaticHashTable` instead.
        lookup_key_dtype: Data type of keys passed to `lookup`. Defaults to
        `initializer.key_dtype` if `initializer` is specified, otherwise
        `tf.string`. Must be string or integer, and must be castable to
        `initializer.key_dtype`.
        name: A name for the operation (optional).
        experimental_is_anonymous: Whether to use anonymous mode for the
        table (default is False). In anonymous mode, the table
        resource can only be accessed via a resource handle. It can't
        be looked up by a name. When all resource handles pointing to
        that resource are gone, the resource will be deleted
        automatically.

Raises:
        ValueError: when `num_oov_buckets` is not positive.
        TypeError: when lookup_key_dtype or initializer.key_dtype are not
        integer or string. Also when initializer.value_dtype != int64.


def non_serializable():

    def _apply_fn(dataset):
        """Function from `Dataset` to `Dataset` that applies the transformation."""
        return _NonSerializableDataset(dataset)
    return _apply_fn

A non-serializable identity transformation.

Returns:
        A `Dataset` transformation function, which can be passed to
        `tf.data.Dataset.apply`.


def get_int(self, min_int=_MIN_INT, max_int=_MAX_INT):
    return self.fdp.ConsumeIntInRange(min_int, max_int)

Consume a signed integer with given constraints.

Args:
        min_int: Minimum allowed integer.
        max_int: Maximum allowed integer.

Returns:
        Consumed integer based on input bytes and constraints.


def run_schedule(inputs: Dict[EventSetNode, EventSet], schedule: Schedule, verbose: int, check_execution: bool, force_garbage_collector_interval: Optional[float]=10) -> Dict[EventSetNode, EventSet]:
    data = {**inputs}
    gc_begin_time = time.time()
    num_steps = len(schedule.steps)
    for step_idx, step in enumerate(schedule.steps):
        operator_def = step.op.definition
        implementation_cls = implementation_lib.get_implementation_class(operator_def.key)
        implementation = implementation_cls(step.op)
        if verbose == 1:
            print(f'    {step_idx + 1} / {num_steps}: {step.op.operator_key()}', file=sys.stderr, end='', flush=True)
        elif verbose >= 2:
            print('=============================', file=sys.stderr)
            print(f'{step_idx + 1} / {num_steps}: Run {step.op}', file=sys.stderr, flush=True)
        operator_inputs = {input_key: data[input_node] for input_key, input_node in step.op.inputs.items()}
        if verbose >= 2:
            print(f'Inputs:\n{operator_inputs}\n', file=sys.stderr, flush=True)
        begin_time = time.perf_counter()
        if check_execution:
            operator_outputs = implementation.call(**operator_inputs)
        else:
            operator_outputs = implementation(**operator_inputs)
        end_time = time.perf_counter()
        if verbose == 1:
            print(f' [{end_time - begin_time:.5f} s]', file=sys.stderr, flush=True)
        elif verbose >= 2:
            print(f'Outputs:\n{operator_outputs}\n', file=sys.stderr)
            print(f'Duration: {end_time - begin_time} s', file=sys.stderr, flush=True)
        for output_key, output_node in step.op.outputs.items():
            output_evset = operator_outputs[output_key]
            output_evset._internal_node = output_node
            data[output_node] = output_evset
        for node in step.released_nodes:
            assert node in data
            del data[node]
        if force_garbage_collector_interval is not None and time.time() - gc_begin_time >= force_garbage_collector_interval:
            begin_gc = time.time()
            if verbose >= 2:
                print('Garbage collection', file=sys.stderr, flush=True, end='')
            gc.collect()
            gc_begin_time = time.time()
            if verbose >= 2:
                print(f' [{gc_begin_time - begin_gc:.5f} s]', file=sys.stderr, flush=True)
    return data

Evaluates a schedule on a dictionary of input
    [`EventSets`][temporian.EventSet].

Args:
        inputs: Mapping of EventSetNodes to materialized EventSets.
        schedule: Sequence of operators to apply on the data.
        verbose: If >0, prints details about the execution on the standard error
            output. The larger the number, the more information is displayed.
        check_execution: If `True`, data of the intermediate results of the
            operators is checked against its expected structure and raises if
            it differs.
        force_garbage_collector_interval: If set, triggers the garbage
            collection every "force_garbage_collector_interval" seconds.


def __init__(self, name='hinge', dtype=None):
    super().__init__(fn=hinge, name=name, dtype=dtype)
    self._direction = 'down'

Computes the hinge metric between `y_true` and `y_pred`.

    `y_true` values are expected to be -1 or 1. If binary (0 or 1) labels are
    provided we will convert them to -1 or 1.

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

Example:
    >>> m = keras.metrics.Hinge()
    >>> m.update_state([[0, 1], [0, 0]], [[0.6, 0.4], [0.4, 0.6]])
    >>> m.result()
    1.3
    >>> m.reset_state()
    >>> m.update_state([[0, 1], [0, 0]], [[0.6, 0.4], [0.4, 0.6]],
    ...                sample_weight=[1, 0])
    >>> m.result()
    1.1


def minutes(value: Union[int, float]) -> Duration:
    return float(value * 60)

Converts input value from minutes to a `Duration` in seconds.

Example:
        ```python
        >>> timestamps = [tp.duration.minutes(i) for i in [5, 10, 30]]
        >>> timestamps
        [300.0, 600.0, 1800.0]

        >>> # Usage in a window operation
        >>> a = tp.event_set(timestamps=timestamps, features={"f1": [1, 5, -5]})
        >>> a.moving_sum(window_length=tp.duration.minutes(6))
        indexes: ...
                timestamps: [ 300. 600. 1800.]
                'f1': [ 1 6 -5]
        ...

        ```

Args:
        value: Number of minutes.

Returns:
        Equivalent number of seconds.


def _VerifyValues(self, image, ksizes, strides, rates, padding, patches):
    ksizes = [1] + ksizes + [1]
    strides = [1] + strides + [1]
    rates = [1] + rates + [1]
    for dtype in [np.float16, np.float32, np.float64, dtypes.bfloat16.as_numpy_dtype]:
        out_tensor = array_ops.extract_image_patches(constant_op.constant(image, dtype=dtype), ksizes=ksizes, strides=strides, rates=rates, padding=padding, name='im2col')
        self.assertAllClose(np.array(patches, dtype=dtype), self.evaluate(out_tensor))

Tests input-output pairs for the ExtractImagePatches op.

Args:
        image: Input tensor with shape: [batch, in_rows, in_cols, depth].
        ksizes: Patch size specified as: [ksize_rows, ksize_cols].
        strides: Output strides, specified as [stride_rows, stride_cols].
        rates: Atrous rates, specified as [rate_rows, rate_cols].
        padding: Padding type.
        patches: Expected output.


def initial_value(self):
    raise NotImplementedError

Returns the Tensor used as the initial value for the variable.

    Note that this is different from `initialized_value()` which runs
    the op that initializes the variable before returning its value.
    This method returns the tensor that is used by the op that initializes
    the variable.

Returns:
        A `Tensor`.


def from_seed(cls, seed, alg=None):
    if alg is None:
        alg = DEFAULT_ALGORITHM
    alg = random_ops_util.convert_alg_to_int(alg)
    state = create_rng_state(seed, alg)
    return cls(state=state, alg=alg)

Creates a generator from a seed.

    A seed is a 1024-bit unsigned integer represented either as a Python
    integer or a vector of integers. Seeds shorter than 1024-bit will be
    padded. The padding, the internal structure of a seed and the way a seed
    is converted to a state are all opaque (unspecified). The only semantics
    specification of seeds is that two different seeds are likely to produce
    two independent generators (but no guarantee).

Args:
        seed: the seed for the RNG.
        alg: (optional) the RNG algorithm. If None, it will be auto-selected. See
        `__init__` for its possible values.

Returns:
        The new generator.


def __init__(self, dataset, coordinator):
    if isinstance(dataset, input_lib.DistributedDataset):
        original_dataset = dataset._original_dataset
        serialized = serialize_dataset_to_graph(original_dataset)

        def dataset_fn():
            deserialized = deserialize_dataset_from_graph(serialized, original_dataset.element_spec)
            dataset.build(dataset_to_replace=deserialized)
            return dataset
    elif isinstance(dataset, input_lib.DistributedDatasetsFromFunction):

        def dataset_fn():
            dataset.build()
            return dataset
    elif isinstance(dataset, dataset_ops.Dataset):
        serialized = serialize_dataset_to_graph(dataset)

        def dataset_fn():
            return deserialize_dataset_from_graph(serialized, dataset.element_spec)
    else:
        raise ValueError('Unexpected dataset type!')
    super(PerWorkerDatasetFromDataset, self).__init__(dataset_fn, coordinator)

Makes an iterable from datasets created by the given dataset.

    It creates a dataset_fn which deserializes a dataset from a graph under the
    hood.

Args:
        dataset: A tf.data.Dataset, a DistributedDataset or a
        DistributedDatasetsFromFunction
        coordinator: a `ClusterCoordinator` object, used to create dataset
        resources.


def iterate_nested(nd: NestedDict, previous_keys: Optional[KeyList]=None) -> Iterator[Tuple[KeyList, LeafType]]:
    if previous_keys is None:
        previous_keys = []
    for k, v in nd.items():
        keys = previous_keys + [k]
        if not _is_nested(v):
            yield (keys, v)
        else:
            as_dict = dataclasses.asdict(v) if dataclasses.is_dataclass(v) else v
            for val in iterate_nested(as_dict, keys):
                yield val

Creates an iterator over every leaf value in depth first order.

    Iterates over a nested dictionary in depth first order. The order in which
    the peer keys are traversed is not guaranteed (same as for the keys of a
    dictionary).

    ```Example
    nested_dict = {'a': 1, 'b': [2, 3, 4], 'c': {'d': 8}}
    for k, v in iterate_nested(nested_dict):
        print('_'.join(k), v)
    # Prints out:
    # a: 1
    # b: [2, 3, 4]
    # c_d: 8
    ```

Args:
        nd: The dictionary to be traversed.
        previous_keys: If supplied, the computed key list will be a join of the
        previous_keys and the current keys.

Yields:
        A tuple of the key path and the value for each leaf node.


def get_c_function(self, name):
    self.ensure_initialized()
    return c_api_util.ScopedTFFunction(pywrap_tfe.TFE_ContextGetFunction(self._handle, name), name)

Get a C API TF_Function from the context.

Args:
        name: Name of the function to get.

Returns:
        A ScopedTFFunction wrapping the C API TF_Function.


def add_update(self, updates, inputs=None):
    if inputs is not None:
        tf_logging.warning('`add_update` `inputs` kwarg has been deprecated. You no longer need to pass a value to `inputs` as it is being automatically inferred.')
    call_context = base_layer_utils.call_context()
    if call_context.in_keras_graph:
        return
    if not call_context.frozen:
        for update in nest.flatten(updates):
            if callable(update):
                update()

Add update op(s), potentially dependent on layer inputs.

    Weight updates (for instance, the updates of the moving mean and variance
    in a BatchNormalization layer) may be dependent on the inputs passed
    when calling a layer. Hence, when reusing the same layer on
    different inputs `a` and `b`, some entries in `layer.updates` may be
    dependent on `a` and some on `b`. This method automatically keeps track
    of dependencies.

    This call is ignored when eager execution is enabled (in that case, variable
    updates are run on the fly and thus do not need to be tracked for later
    execution).

Args:
        updates: Update op, or list/tuple of update ops, or zero-arg callable
        that returns an update op. A zero-arg callable should be passed in
        order to disable running the updates by setting `trainable=False`
        on this Layer, when executing in Eager mode.
        inputs: Deprecated, will be automatically inferred.


def wait_at_barrier(self, barrier_id, timeout_in_ms):
    ensure_initialized()
    pywrap_tfe.TFE_WaitAtBarrier(self._context_handle, barrier_id, timeout_in_ms)

Blocks until all coordinated tasks are at the barrier.

    The barrier may fail if it times out or if one of the tasks is unhealthy.

Args:
        barrier_id: Unique string identifying the barrier.
        timeout_in_ms: Duration before the barrier times out and fails.


def param_static_shapes(cls, sample_shape):
    if isinstance(sample_shape, tensor_shape.TensorShape):
        if not sample_shape.is_fully_defined():
            raise ValueError('TensorShape sample_shape must be fully defined')
        sample_shape = sample_shape.as_list()
    params = cls.param_shapes(sample_shape)
    static_params = {}
    for name, shape in params.items():
        static_shape = tensor_util.constant_value(shape)
        if static_shape is None:
            raise ValueError('sample_shape must be a fully-defined TensorShape or list/tuple')
        static_params[name] = tensor_shape.TensorShape(static_shape)
    return static_params

param_shapes with static (i.e. `TensorShape`) shapes.

    This is a class method that describes what key/value arguments are required
    to instantiate the given `Distribution` so that a particular shape is
    returned for that instance's call to `sample()`. Assumes that the sample's
    shape is known statically.

    Subclasses should override class method `_param_shapes` to return
    constant-valued tensors when constant values are fed.

Args:
        sample_shape: `TensorShape` or python list/tuple. Desired shape of a call
        to `sample()`.

Returns:
        `dict` of parameter name to `TensorShape`.

Raises:
        ValueError: if `sample_shape` is a `TensorShape` and is not fully defined.


def stop(save=True):
    global _profiler
    with _profiler_lock:
        if _profiler is None:
            raise errors.UnavailableError(None, None, 'Cannot export profiling results. No profiler is running.')
        if save:
            try:
                _profiler.export_to_tb()
            except Exception:
                _profiler = None
                raise
        _profiler = None

Stops the current profiling session.

    The profiler session will be stopped and profile results can be saved.

Args:
        save: An optional variable to save the results to TensorBoard. Default True.

Raises:
        UnavailableError: If there is no active profiling session.


def _is_framework_filename(filename):
    for pattern in _EXTERNAL_FILENAME_PATTERNS:
        if pattern.search(filename):
            return False
    for pattern in _FRAMEWORK_FILENAME_PATTERNS:
        if pattern.search(filename):
            return True
    for prefix in _FRAMEWORK_PATH_PREFIXES:
        if filename.startswith(prefix):
            return True
    return False

Returns whether a filename should be considered a part of the framework.

    A file is part of the framework if it does not match a pattern in
    _EXTERNAL_FILENAME_PATTERNS and it either matches a pattern in
    _FRAMEWORK_FILENAME_PATTERNS or starts with a _FRAMEWORK_PATH_PREFIXES prefix.

Args:
        filename: A filename string.

Returns:
        Whether the filename should be considered to be internal to the
        TensorFlow framework for the purposes of reporting errors.


def recipe_bucket(config, auth_write, bucket_bucket, bucket_emails, bucket_groups):
    bucket(config, {'auth': auth_write, 'bucket': bucket_bucket, 'emails': bucket_emails, 'groups': bucket_groups})

Create and permission a bucket in Storage.

Args:
        auth_write (authentication) - Credentials used for writing data.
        bucket_bucket (string) - Name of Google Cloud Bucket to create.
        bucket_emails (string_list) - Comma separated emails.
        bucket_groups (string_list) - Comma separated groups.


def transform_feature(self, transformation_cache, state_manager):
    pass

Returns intermediate representation (usually a `Tensor`).

    Uses `transformation_cache` to create an intermediate representation
    (usually a `Tensor`) that other feature columns can use.

    Example usage of `transformation_cache`:
    Let's say a Feature column depends on raw feature ('raw') and another
    `FeatureColumn` (input_fc). To access corresponding `Tensor`s,
    transformation_cache will be used as follows:

    ```python
    raw_tensor = transformation_cache.get('raw', state_manager)
    fc_tensor = transformation_cache.get(input_fc, state_manager)
    ```

Args:
        transformation_cache: A `FeatureTransformationCache` object to access
        features.
        state_manager: A `StateManager` to create / access resources such as
        lookup tables.

Returns:
        Transformed feature `Tensor`.