def forward(self, trajectories: Optional[torch.LongTensor]=None, past_key_values: Optional[Tuple[Tuple[torch.Tensor]]]=None, targets: Optional[torch.FloatTensor]=None, attention_mask: Optional[torch.FloatTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[Tuple[torch.Tensor], TrajectoryTransformerOutput]:
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    if past_key_values is None:
        past_key_values = tuple([None] * len(self.blocks))
    batch_size, sequence_length = trajectories.size()
    if sequence_length > self.block_size:
        raise ValueError('Cannot forward, model block size is exhausted.')
    offset_trajectories = self.offset_tokens(trajectories)
    token_embeddings = self.tok_emb(offset_trajectories)
    position_embeddings = self.pos_emb[:, :sequence_length, :]
    hidden_states = self.drop(token_embeddings + position_embeddings)
    if self.gradient_checkpointing and self.training:
        if use_cache:
            logger.warning_once('`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...')
            use_cache = False
    presents = () if use_cache else None
    all_self_attentions = () if output_attentions else None
    all_hidden_states = () if output_hidden_states else None
    for i, (block, layer_past) in enumerate(zip(self.blocks, past_key_values)):
        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)
        if self.gradient_checkpointing and self.training:
            outputs = self._gradient_checkpointing_func(block.__call__, hidden_states, layer_past, use_cache, output_attentions)
        else:
            outputs = block(hidden_states, layer_past, use_cache, output_attentions)
        hidden_states = outputs[0]
        if use_cache is True:
            presents = presents + (outputs[1],)
        if output_attentions:
            all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)
    hidden_state = self.ln_f(hidden_states)
    if output_hidden_states:
        all_hidden_states = all_hidden_states + (hidden_states,)
    hidden_states_pad, n_pad = self.pad_to_full_observation(hidden_state)
    logits = self.head(hidden_states_pad)
    logits = logits.reshape(batch_size, sequence_length + n_pad, self.vocab_size + 1)
    logits = logits[:, :sequence_length]
    if targets is not None:
        loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.view(-1), reduction='none')
        if self.action_weight != 1 or self.reward_weight != 1 or self.value_weight != 1:
            n_states = int(np.ceil(sequence_length / self.transition_dim))
            weights = torch.cat([torch.ones(self.observation_dim, device=trajectories.device), torch.ones(self.action_dim, device=trajectories.device) * self.action_weight, torch.ones(1, device=trajectories.device) * self.reward_weight, torch.ones(1, device=trajectories.device) * self.value_weight])
            weights = weights.repeat(n_states)
            weights = weights[1:].repeat(batch_size, 1)
            loss = loss * weights.view(-1)
        loss = (loss * attention_mask.view(-1)).mean()
    else:
        loss = None
    if not return_dict:
        return tuple((v for v in [loss, logits, presents, all_hidden_states, all_self_attentions] if v is not None))
    return TrajectoryTransformerOutput(loss=loss, logits=logits, past_key_values=presents, hidden_states=all_hidden_states, attentions=all_self_attentions)

Returns:

Example:
    ```python
    >>> from transformers import TrajectoryTransformerModel
    >>> import torch

    >>> model = TrajectoryTransformerModel.from_pretrained(
    ...     "CarlCochet/trajectory-transformer-halfcheetah-medium-v2"
    ... )
    >>> model.to(device)
    >>> model.eval()

    >>> observations_dim, action_dim, batch_size = 17, 6, 256
    >>> seq_length = observations_dim + action_dim + 1

    >>> trajectories = torch.LongTensor([np.random.permutation(self.seq_length) for _ in range(batch_size)]).to(
    ...     device
    ... )
    >>> targets = torch.LongTensor([np.random.permutation(self.seq_length) for _ in range(batch_size)]).to(device)

    >>> outputs = model(
    ...     trajectories,
    ...     targets=targets,
    ...     use_cache=True,
    ...     output_attentions=True,
    ...     output_hidden_states=True,
    ...     return_dict=True,
    ... )
    ```


def __init__(self, thresholds=None, name=None, dtype=None):
    super(FalseNegatives, self).__init__(confusion_matrix_cond=metrics_utils.ConfusionMatrix.FALSE_NEGATIVES, thresholds=thresholds, name=name, dtype=dtype)

Calculates the number of false negatives.

    If `sample_weight` is given, calculates the sum of the weights of
    false negatives. This metric creates one local variable, `accumulator`
    that is used to keep track of the number of false negatives.

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

Args:
        thresholds: (Optional) Defaults to 0.5. A float value or a python
        list/tuple of float threshold values in [0, 1]. A threshold is compared
        with prediction values to determine the truth value of predictions
        (i.e., above the threshold is `true`, below is `false`). One metric
        value is generated for each threshold value.
        name: (Optional) string name of the metric instance.
        dtype: (Optional) data type of the metric result.

    Standalone usage:

    >>> m = tf.keras.metrics.FalseNegatives()
    >>> m.update_state([0, 1, 1, 1], [0, 1, 0, 0])
    >>> m.result().numpy()
    2.0

    >>> m.reset_state()
    >>> m.update_state([0, 1, 1, 1], [0, 1, 0, 0], sample_weight=[0, 0, 1, 0])
    >>> m.result().numpy()
    1.0

    Usage with `compile()` API:

    ```python
    model.compile(optimizer='sgd',
                loss='mse',
                metrics=[tf.keras.metrics.FalseNegatives()])
    ```


def _VerifyValues(self, image, kernel, strides, rates, padding, out, use_gpu, dtype):
    strides = [1] + strides + [1]
    rates = [1] + rates + [1]
    with self.cached_session(use_gpu=use_gpu):
        out_tensor = nn_ops.dilation2d(constant_op.constant(image, dtype=dtype), constant_op.constant(kernel, dtype=dtype), strides=strides, rates=rates, padding=padding, name='dilation2d')
        self.assertAllCloseAccordingToType(out, self.evaluate(out_tensor))

Verifies the output values of the dilation function.

Args:
        image: Input tensor with shape: [batch, in_height, in_width, channels].
        kernel: Filter tensor with shape: [filter_height, filter_width, channels].
        strides: Output strides, specified as [stride_height, stride_width].
        rates: Atrous rates, specified as [rate_height, rate_width].
        padding: Padding type.
        out: Expected output.
        use_gpu: Whether we are running on GPU.


def set_tuple_shapes(self, tuple_shapes):
    if len(tuple_shapes) != self.number_of_tuple_elements:
        raise ValueError(f'tuple_shapes is {str(tuple_shapes)}, but must be a list of length {self.number_of_tuple_elements}')
    try:
        tuple_shapes = [tensor_shape.as_shape(shape) for shape in tuple_shapes]
    except (ValueError, TypeError) as e:
        raise TypeError(f'tuple_shapes is {str(tuple_shapes)}, but must be a list of elements each convertible to TensorShape: got error {str(e)}') from e
    if self._frozen:
        for frozen, updated in zip(self._tuple_shapes, tuple_shapes):
            if frozen != updated:
                raise ValueError(f'Trying to update InfeedQueue with frozen configuration with an incompatible shape. Frozen shapes are {str(self._tuple_shapes)}, updated shapes are {str(tuple_shapes)}')
    else:
        self._tuple_shapes = tuple_shapes
    self._validate()

Sets the shape of each element of the queue.

    tuple_shapes must be a list of length
    self.number_of_tuple_elements, and each element must be
    convertible to a TensorShape.

Args:
        tuple_shapes: the shapes of each queue element.

Raises:
        ValueError: if tuple_shapes is not of length
        self.number_of_tuple_elements.
        TypeError: if an element of tuple_shapes cannot be converted to
        a TensorShape.


def NHWCToNCHW(input_tensor):
    if isinstance(input_tensor, tensor.Tensor):
        return array_ops.transpose(input_tensor, [0, 3, 1, 2])
    else:
        return [input_tensor[0], input_tensor[3], input_tensor[1], input_tensor[2]]

Convert the input from NHWC format to NCHW.

Args:
        input_tensor:  a 4-D tensor, or a 4-element array representing the same.

Returns:
        the converted tensor or a shape array


def set_options(cls, pipeline_options):
    cls._pipeline_options = pipeline_options

Set filesystem options.

Args:
        pipeline_options: Instance of ``PipelineOptions``.


def __init__(self, config: FastSpeech2ConformerConfig, module_config):
    super().__init__()
    input_channels = config.hidden_size
    hidden_channels = module_config['linear_units']
    kernel_size = config.positionwise_conv_kernel_size
    self.conv1 = nn.Conv1d(input_channels, hidden_channels, kernel_size, stride=1, padding=(kernel_size - 1) // 2)
    self.conv2 = nn.Conv1d(hidden_channels, input_channels, kernel_size, stride=1, padding=(kernel_size - 1) // 2)
    self.dropout = nn.Dropout(module_config['dropout_rate'])

Initialize FastSpeech2ConformerMultiLayeredConv1d module.

Args:
        input_channels (`int`): Number of input channels.
        hidden_channels (`int`): Number of hidden channels.
        kernel_size (`int`): Kernel size of conv1d.
        dropout_rate (`float`): Dropout rate.


def test_gradient(self, shape, rt_value, rt_grad, default_value, default_grad, output_value, output_grad, ragged_rank=None):
    rt_value = ragged_factory_ops.constant(rt_value, dtype=dtypes.float32, ragged_rank=ragged_rank)
    rt_grad = ragged_factory_ops.constant(rt_grad, dtype=dtypes.float32, ragged_rank=ragged_rank)
    default_value = constant_op.constant(default_value, dtype=dtypes.float32)
    default_grad = constant_op.constant(default_grad, dtype=dtypes.float32)
    output_value = constant_op.constant(output_value, dtype=dtypes.float32, shape=shape)
    output_grad = constant_op.constant(output_grad, dtype=dtypes.float32, shape=shape)
    shape = tensor_shape.as_shape(shape)
    for partition_type in ['row_splits', 'value_rowids']:
        rt_val = self.rt_with_partition_type(rt_value, partition_type)
        if context.executing_eagerly():
            self._test_gradient_helper(rt_val, default_value, shape, output_grad, output_value, rt_grad, default_grad)
        else:
            for shape_info in ['known', 'unknown_dims', 'unknown_rank']:
                rt_val = self.wrap_in_placeholder(rt_val, shape_info)
                default_val = self.wrap_in_placeholder(default_value, shape_info)
                shape_val = self.wrap_in_placeholder(shape, shape_info)
                self._test_gradient_helper(rt_val, default_val, shape_val, output_grad, output_value, rt_grad, default_grad)

Tests that ragged_to_dense generates the right gradient.

Args:
        shape: The `shape` arg for `ragged_to_dense`.
        rt_value: The `rt_input` arg for `ragged_to_dense`.
        rt_grad: The expected gradient for `rt_value`.  Corresponds 1:1 with
        `rt_value`.
        default_value: The `default_value` arg for `ragged_to_dense`.
        default_grad: The expected gradient for `default_value`.  Corresponds 1:1
        with `default_value`.
        output_value: The expected output of `ragged_to_dense`.
        output_grad: The gradient for the output (used to generate the gradients
        `rt_grad` and `default_grad`).  Corresponds 1:1 with `output_value`.
        ragged_rank: Ragged rank for `rt_value`.


def outer(x1, x2):
    if any_symbolic_tensors((x1, x2)):
        return Outer().symbolic_call(x1, x2)
    return backend.numpy.outer(x1, x2)

Compute the outer product of two vectors.

    Given two vectors `x1` and `x2`, the outer product is:

    ```
    out[i, j] = x1[i] * x2[j]
    ```

Args:
        x1: First input tensor.
        x2: Second input tensor.

Returns:
        Outer product of `x1` and `x2`.


def select(self, children: List['ColumnExpressionBuilder']) -> 'ColumnExpressionBuilder':
    if self._column_name:
        raise AttributeError(f'select() must not be called on a builder with name set already. Got select called on {str(self)}.')
    if self._builder.return_type.returns_collection() and (not self._needs_unnest):
        raise AttributeError(f'select() must not be called on a builder which returns collection. Got select called on {str(self)}.')
    if not isinstance(self._builder.return_type, _fhir_path_data_types.StructureDataType):
        raise AttributeError(f'select() can only be called on a FHIR path which returns a StructureDataType. Got type of {self._builder.return_type} from {str(self)}.')
    for child_builder in children:
        if not child_builder.column_name and (not child_builder.children):
            raise AttributeError(f'select() child builders must either have name names or children. Got {str(child_builder)}')
        if not child_builder.fhir_path.startswith(self._builder.fhir_path):
            raise AttributeError(f'select() child builders must be built by starting with their parent FHIR path. Got {str(child_builder)} in {self._builder.fhir_path}.')
        path_to_replace = self._builder.fhir_path
        reference_node = self._builder.node
        child_builder.node.replace_operand(path_to_replace, _evaluation.ReferenceNode(self._builder.node.context, reference_node, unnested=True))
    return ColumnExpressionBuilder(self._builder, self._column_name, children, self._needs_unnest, True)

The select() function.

    Selects the child fields from a FHIR path which returns a StructureDataType.
    Once this function is called, the FHIR path is sealed to be immutable.

Args:
        children: A list of selected FHIR Path.

Returns:
        A new ColumnExpressionBuilder with the given children.


def process_update(x):
    if callable(x):
        update = lambda: process_update(x())
        return update()
    elif isinstance(x, ops.Operation):
        update = x
    elif hasattr(x, 'op'):
        update = x.op
    else:
        update = tensor_conversion.convert_to_tensor_v2_with_dispatch(x)
    reachable = tf_utils.get_reachable_from_inputs(relevant_inputs, [update])
    update._unconditional_update = update not in reachable
    return update

Standardize update ops.

Args:
        x: Tensor, op, or callable.

Returns:
        An update op.


def create_profiler_ui(graph, run_metadata, ui_type='readline', on_ui_exit=None, config=None):
    del config
    analyzer = ProfileAnalyzer(graph, run_metadata)
    cli = ui_factory.get_ui(ui_type, on_ui_exit=on_ui_exit)
    cli.register_command_handler('list_profile', analyzer.list_profile, analyzer.get_help('list_profile'), prefix_aliases=['lp'])
    cli.register_command_handler('print_source', analyzer.print_source, analyzer.get_help('print_source'), prefix_aliases=['ps'])
    return cli

Create an instance of ReadlineUI based on a `tf.Graph` and `RunMetadata`.

Args:
        graph: Python `Graph` object.
        run_metadata: A `RunMetadata` protobuf object.
        ui_type: (str) requested UI type, e.g., "readline".
        on_ui_exit: (`Callable`) the callback to be called when the UI exits.
        config: An instance of `cli_config.CLIConfig`.

Returns:
        (base_ui.BaseUI) A BaseUI subtype object with a set of standard analyzer
        commands and tab-completions registered.


Updates the value of the loss scale.

    The loss scale will be potentially updated, based on the value of `grads`.
    The tensor returned by calling this class is only updated when this function
    is evaluated.

    In eager mode, this directly updates the loss scale, so that calling
    `__call__` will return the newly updated loss scale. In graph mode,
    this returns an op that, when evaluated, updates the loss scale.

    This function also returns a `should_apply_gradients` bool. If False,
    gradients should not be applied to the variables that step, as nonfinite
    gradients were found, and the loss scale has been be updated to reduce the
    chance of finding nonfinite gradients in the next step. Some loss scale
    classes will always return True, as they cannot adjust themselves in
    response to nonfinite gradients.

    When a DistributionStrategy is used, this function may only be called in a
    cross-replica context.

Args:
        grads: A nested structure of unscaled gradients, each which is the
        gradient of the loss with respect to a weight. The gradients should have
        already been divided by the loss scale being before passed to this
        function. 'None' gradients are accepted, and are ignored.

Returns:
        update_op: In eager mode, None. In graph mode, an op to update the loss
        scale.
        should_apply_gradients: Either a bool or a scalar boolean tensor. If
        False, the caller should skip applying `grads` to the variables this
        step.


def assert_equal(first, second, msg=None, extras=None):
    _call_unittest_assertion(_pyunit_proxy.assertEqual, first, second, msg=msg, extras=extras)

Asserts the equality of objects, otherwise fail the test.

    Error message is "first != second" by default. Additional explanation can
    be supplied in the message.

Args:
        first: The first object to compare.
        second: The second object to compare.
        msg: A string that adds additional info about the failure.
        extras: An optional field for extra information to be included in
        test result.


def sigmoid(x):
    output = ops.sigmoid(x)
    try:
        output._keras_logits = x
    except AttributeError:
        pass
    return output

Sigmoid activation function.

    It is defined as: `sigmoid(x) = 1 / (1 + exp(-x))`.

    For small values (<-5),
    `sigmoid` returns a value close to zero, and for large values (>5)
    the result of the function gets close to 1.

    Sigmoid is equivalent to a 2-element softmax, where the second element is
    assumed to be zero. The sigmoid function always returns a value between
    0 and 1.

Args:
        x: Input tensor.


def on_value_event(self, event):
    if self._dump_dir:
        self._write_value_event(event)
    else:
        value = event.summary.value[0]
        tensor_value = debug_data.load_tensor_from_event(event)
        self._event_listener_servicer.debug_tensor_values[value.node_name].append(tensor_value)
        items = event.summary.value[0].node_name.split(':')
        node_name = items[0]
        output_slot = int(items[1])
        debug_op = items[2]
        if (node_name, output_slot, debug_op) in self._event_listener_servicer.breakpoints:
            return debug_service_pb2.EventReply()

Implementation of the tensor value-carrying Event proto callback.

    Writes the Event proto to the file system for testing. The path written to
    follows the same pattern as the file:// debug URLs of tfdbg, i.e., the
    name scope of the op becomes the directory structure under the dump root
    directory.

Args:
        event: The Event proto carrying a tensor value.

Returns:
        If the debug node belongs to the set of currently activated breakpoints,
        a `EventReply` proto will be returned.


def _aggregate_gradients(self, grads_and_vars):
    return self.gradient_aggregator(grads_and_vars)

Called in `apply_gradients` to aggregate gradients across devices.

    Note that user subclasses may override this, so the interface should not be
    changed.

Args:
        grads_and_vars: List of (gradient, variable) pairs.

Returns:
        A list of (aggregated_gradient, variable) pairs. By default, this calls
        `self.gradient_aggregator`.


def add_test_class(self, config, test_class, tests=None, name_suffix=None):
    if self._log_dir != config.log_path:
        raise Error('TestRunner\'s log folder is "%s", but a test config with a different log folder ("%s") was added.' % (self._log_dir, config.log_path))
    if self._testbed_name != config.testbed_name:
        raise Error('TestRunner\'s test bed is "%s", but a test config with a different test bed ("%s") was added.' % (self._testbed_name, config.testbed_name))
    self._test_run_infos.append(TestRunner._TestRunInfo(config=config, test_class=test_class, tests=tests, test_class_name_suffix=name_suffix))

Adds tests to the execution plan of this TestRunner.

Args:
        config: config_parser.TestRunConfig, configuration to execute this
        test class with.
        test_class: class, test class to execute.
        tests: list of strings, optional list of test names within the
        class to execute.
        name_suffix: string, suffix to append to the class name for
        reporting. This is used for differentiating the same class
        executed with different parameters in a suite.

Raises:
        Error: if the provided config has a log_path or testbed_name which
        differs from the arguments provided to this TestRunner's
        constructor.


def find(self, predicate, first_n=0, device_name=None, exclude_node_names=None):
    if exclude_node_names:
        exclude_node_names = re.compile(exclude_node_names)
    matched_data = []
    for device in self._dump_tensor_data if device_name is None else (self._dump_tensor_data[device_name],):
        for datum in self._dump_tensor_data[device]:
            if exclude_node_names and exclude_node_names.match(datum.node_name):
                continue
            if predicate(datum, datum.get_tensor()):
                matched_data.append(datum)
                if first_n > 0 and len(matched_data) >= first_n:
                    return matched_data
    return matched_data

Find dumped tensor data by a certain predicate.

Args:
        predicate: A callable that takes two input arguments:

        ```python
        def predicate(debug_tensor_datum, tensor):
        # returns a bool
        ```

        where `debug_tensor_datum` is an instance of `DebugTensorDatum`, which
        carries the metadata, such as the `Tensor`'s node name, output slot
        timestamp, debug op name, etc.; and `tensor` is the dumped tensor value
        as a `numpy.ndarray`.
        first_n: (`int`) return only the first n `DebugTensotDatum` instances (in
        time order) for which the predicate returns True. To return all the
        `DebugTensotDatum` instances, let first_n be <= 0.
        device_name: optional device name.
        exclude_node_names: Optional regular expression to exclude nodes with
        names matching the regular expression.

Returns:
        A list of all `DebugTensorDatum` objects in this `DebugDumpDir` object
        for which predicate returns True, sorted in ascending order of the
        timestamp.


def index_of(self, value_str):
    if value_str is None:
        value_str = ''
    if value_str in self._string_to_index:
        return self._string_to_index[value_str]
    index = len(self._string_table)
    self._string_table.append(value_str)
    self._string_to_index[value_str] = index
    return index

Get index of value_str in the string table.

    If value_str is not in the string table, we will add it at the end
    and then return the new index.

Args:
        value_str: (string) Value to lookup/add in/to the string table.

Returns:
        Index of value_str in the string table.


def assertOutputStateMatches(self, **has_output):
    output_types = {'stdout', 'stderr', 'returncode'}
    assert len(output_types) == len(has_output)
    for output_type in output_types:
        output_value = getattr(self, output_type)
        if has_output[output_type]:
            self.assertTrue(output_value, output_type + ' unexpectedly empty')
        else:
            value = str(output_value)
            if len(value) > 50:
                value = value[:47] + '...'
            self.assertFalse(output_value, f'Unexpected output to {output_type}: {value!r}')

Check that the output state matches expectations.

    If, for example, you expect the program to print something to stdout and
    nothing to stderr before exiting with an error code, you would write
    assertOutputStateMatches(stdout=True, stderr=False, returncode=True).

Args:
        **has_output: Whether each output type should have output.


def __call__(self, shape, dtype=None):
    raise NotImplementedError('Initializer subclasses must implement the `__call__()` method.')

Returns a tensor object initialized as specified by the initializer.

Args:
        shape: Shape of the tensor.
        dtype: Optional dtype of the tensor.


def assert_shallow_structure(shallow_tree, input_tree, check_types=True, expand_composites=False):
    nest_util.assert_shallow_structure(nest_util.Modality.CORE, shallow_tree, input_tree, check_types, expand_composites)

Asserts that `shallow_tree` is a shallow structure of `input_tree`.

    That is, this function tests if the `input_tree` structure can be created from
    the `shallow_tree` structure by replacing its leaf nodes with deeper
    tree structures.

Example:
    The following code will raise an exception:
    ```python
        shallow_tree = {"a": "A", "b": "B"}
        input_tree = {"a": 1, "c": 2}
        assert_shallow_structure(shallow_tree, input_tree)
    ```

    The following code will raise an exception:
    ```python
        shallow_tree = ["a", "b"]
        input_tree = ["c", ["d", "e"], "f"]
        assert_shallow_structure(shallow_tree, input_tree)
    ```

Args:
        shallow_tree: an arbitrarily nested structure.
        input_tree: an arbitrarily nested structure.
        check_types: if `True` (default) the sequence types of `shallow_tree` and
        `input_tree` have to be the same. Note that even with check_types==True,
        this function will consider two different namedtuple classes with the same
        name and _fields attribute to be the same class.
        expand_composites: If true, then composite tensors such as
        `tf.sparse.SparseTensor` and `tf.RaggedTensor` are expanded into their
        component tensors.

Raises:
        TypeError: If `shallow_tree` is a sequence but `input_tree` is not.
        TypeError: If the sequence types of `shallow_tree` are different from
        `input_tree`. Only raised if `check_types` is `True`.
        ValueError: If the sequence lengths of `shallow_tree` are different from
        `input_tree`.


def __init__(self, opening_bracket, newline, opening_column):
    self.opening_bracket = opening_bracket
    self.has_split_before_first_param = newline
    self.opening_column = opening_column
    self.parameters = opening_bracket.parameters
    self.split_before_closing_bracket = False

Maintains the state of function parameter list formatting decisions.

Attributes:
        opening_bracket: The opening bracket of the parameter list.
        closing_bracket: The closing bracket of the parameter list.
        has_typed_return: True if the function definition has a typed return.
        ends_in_comma: True if the parameter list ends in a comma.
        last_token: Returns the last token of the function declaration.
        has_default_values: True if the parameters have default values.
        has_split_before_first_param: Whether there is a newline before the first
        parameter.
        opening_column: The position of the opening parameter before a newline.
        parameters: A list of parameter objects (Parameter).
        split_before_closing_bracket: Split before the closing bracket. Sometimes
        needed if the indentation would collide.


def _var_key(var):
    if hasattr(var, '_distributed_container'):
        var = var._distributed_container()
    if var._in_graph_mode:
        return var._shared_name
    return var._unique_id

Key for representing a primary variable, for looking up slots.

    In graph mode the name is derived from the var shared name.
    In eager mode the name is derived from the var unique id.
    If distribution strategy exists, get the primary variable first.

Args:
        var: the variable.

Returns:
        the unique name of the variable.


def preprocess_dataset(ingested_dataset_path: str, preprocessed_dataset_path: str, base_artifact_path: str, gcp_project_id: str, region: str, dataflow_staging_root: str, beam_runner: str):
    timestamp = time.time()
    target_path = f'{base_artifact_path}/preprocessing/preprocessed_dataset_{timestamp}'
    Path(preprocessed_dataset_path).parent.mkdir(parents=True, exist_ok=True)
    with open(preprocessed_dataset_path, 'w') as f:
        f.write(target_path)
    pipeline_options = PipelineOptions(runner=beam_runner, project=gcp_project_id, job_name=f'preprocessing-{int(time.time())}', temp_location=dataflow_staging_root, region=region, requirements_file='/requirements.txt', save_main_session=True)
    with beam.Pipeline(options=pipeline_options) as pipeline:
        pipeline | 'Read input jsonlines file' >> beam.io.ReadFromText(ingested_dataset_path) | 'Load json' >> beam.Map(json.loads) | 'Filter licenses' >> beam.Filter(valid_license) | 'Download image from URL' >> beam.FlatMap(download_image_from_url) | 'Resize image' >> beam.Map(resize_image, size=IMAGE_SIZE) | 'Clean Text' >> beam.Map(clean_text) | 'Serialize Example' >> beam.Map(serialize_example) | 'Write to Avro files' >> beam.io.WriteToAvro(file_path_prefix=target_path, schema={'namespace': 'preprocessing.example', 'type': 'record', 'name': 'Sample', 'fields': [{'name': 'id', 'type': 'int'}, {'name': 'caption', 'type': 'string'}, {'name': 'image', 'type': 'bytes'}]}, file_name_suffix='.avro')

Preprocess the ingested raw dataset and write the result to avro format.

Args:
        ingested_dataset_path (str): Path to the ingested dataset
        preprocessed_dataset_path (str): Path to where the preprocessed dataset will be saved
        base_artifact_path (str): path to the base directory of where artifacts can be stored for
        this component.
        gcp_project_id (str): ID for the google cloud project to deploy the pipeline to.
        region (str): Region in which to deploy the pipeline.
        dataflow_staging_root (str): Path to staging directory for the dataflow runner.
        beam_runner (str): Beam runner: DataflowRunner or DirectRunner.


def absolute(x):
    if any_symbolic_tensors((x,)):
        return Absolute().symbolic_call(x)
    return backend.numpy.absolute(x)

Compute the absolute value element-wise.

    `keras.ops.abs` is a shorthand for this function.

Args:
        x: Input tensor.

Returns:
        An array containing the absolute value of each element in `x`.

Example:
    >>> x = keras.ops.convert_to_tensor([-1.2, 1.2])
    >>> keras.ops.absolute(x)
    array([1.2, 1.2], dtype=float32)


def multiplier_with_docstring(num, rate=2):
    return num * rate

Multiplies num by rate.

Args:
        num (int): the num you want to multiply
        rate (int): the rate for multiplication

Returns:
        Multiplication of num by rate


def __init__(self, source_shape, target_shape, layer_broadcasters, dtype=None):
    if not isinstance(source_shape, DynamicRaggedShape):
        raise TypeError('source_shape is not a DynamicRaggedShape')
    if not isinstance(target_shape, DynamicRaggedShape):
        raise TypeError('target_shape is not a DynamicRaggedShape')
    if not isinstance(layer_broadcasters, list):
        raise TypeError('layer_broadcasters not a list: ' + str(layer_broadcasters))
    for bc in layer_broadcasters:
        if not isinstance(bc, _LayerBroadcaster):
            raise TypeError('Not a LayerBroadcaster: ' + str(bc))
    dtype = _find_dtype(source_shape, dtype)
    dtype = _find_dtype(target_shape, dtype)
    dtype = _find_dtype_iterable(layer_broadcasters, dtype)
    dtype = _find_dtype(dtypes.int64, dtype)
    self._source_shape = source_shape.with_dtype(dtype)
    self._target_shape = target_shape.with_dtype(dtype)
    self._layer_broadcasters = [x.with_dtype(dtype) for x in layer_broadcasters]

Create a broadcaster.

    Do not call directly.
    The source_shape, target_shape, and layer_broadcasters are converted
    to have the same dtype.

    Note: source_shape.rank and target_shape.rank must be known.

Args:
        source_shape: the source DynamicRaggedShape
        target_shape: the target DynamicRaggedShape
        layer_broadcasters: List[_LayerBroadcaster] of length source_shape.rank.
        dtype: the preferred dtype of the broadcaster.

Raises:
        TypeError: if the input types don't match.


def __init__(self, initial_learning_rate: float, decay_schedule_fn: Callable, warmup_steps: int, power: float=1.0, name: Optional[str]=None):
    super().__init__()
    self.initial_learning_rate = initial_learning_rate
    self.warmup_steps = warmup_steps
    self.power = power
    self.decay_schedule_fn = decay_schedule_fn
    self.name = name

Applies a warmup schedule on a given learning rate decay schedule.

Args:
        initial_learning_rate (`float`):
            The initial learning rate for the schedule after the warmup (so this will be the learning rate at the end
            of the warmup).
        decay_schedule_fn (`Callable`):
            The schedule function to apply after the warmup for the rest of training.
        warmup_steps (`int`):
            The number of steps for the warmup part of training.
        power (`float`, *optional*, defaults to 1.0):
            The power to use for the polynomial warmup (defaults is a linear warmup).
        name (`str`, *optional*):
            Optional name prefix for the returned tensors during the schedule.


def wait_for_session(self, master: str, config=None, max_wait_secs=float('Inf')) -> Optional[session.Session]:
    self._target = master
    if max_wait_secs is None:
        max_wait_secs = float('Inf')
    timer = _CountDownTimer(max_wait_secs)
    while True:
        sess = session.Session(self._target, graph=self._graph, config=config)
        not_ready_msg = None
        not_ready_local_msg = None
        local_init_success, not_ready_local_msg = self._try_run_local_init_op(sess)
        if local_init_success:
            is_ready, not_ready_msg = self._model_ready(sess)
            if is_ready:
                return sess
        self._safe_close(sess)
        remaining_ms_after_wait = timer.secs_remaining() - self._recovery_wait_secs
        if remaining_ms_after_wait < 0:
            raise errors.DeadlineExceededError(None, None, 'Session was not ready after waiting %d secs.' % (max_wait_secs,))
        logging.info('Waiting for model to be ready.  Ready_for_local_init_op:  %s, ready: %s', not_ready_local_msg, not_ready_msg)
        time.sleep(self._recovery_wait_secs)

Creates a new `Session` and waits for model to be ready.

    Creates a new `Session` on 'master'.  Waits for the model to be
    initialized or recovered from a checkpoint.  It's expected that
    another thread or process will make the model ready, and that this
    is intended to be used by threads/processes that participate in a
    distributed training configuration where a different thread/process
    is responsible for initializing or recovering the model being trained.

    NB: The amount of time this method waits for the session is bounded
    by max_wait_secs. By default, this function will wait indefinitely.

Args:
        master: `String` representation of the TensorFlow master to use.
        config: Optional ConfigProto proto used to configure the session.
        max_wait_secs: Maximum time to wait for the session to become available.

Returns:
        A `Session`. May be None if the operation exceeds the timeout
        specified by config.operation_timeout_in_ms.

Raises:
        tf.DeadlineExceededError: if the session is not available after
        max_wait_secs.


def query(self, src: Any, use_inferred: bool=False) -> Any:
    return self._query(0, src, use_inferred)

Query the value from the source object based on current path.

    Example::

        @pg.members([
        ('x', pg.typing.Int()),
        ('y', pg.typing.Str())
        ])
        class A(pg.Object):
        pass

        @pg.members([
        ('z', pg.typing.Object(A))
        ])
        class B(pg.Object):
        pass

        b = B(z=A(x=1, y='foo'))
        assert pg.KeyPath.parse('z.x').query(b) == 1

Args:
        src: Source value to query.
        use_inferred: If True, infer `pg.Inferential` values. Otherwise returns
        their symbolic form. Applicable only for symbolic values.

Returns:
        Value from src if path exists.

Raises:
        KeyError: Path doesn't exist in src.
        RuntimeError: Called on a KeyPath that is considered as removed.


Sends the JSON RPC request to the server and gets a response.

    Note that the request and response are both in string format. So if the
    connection with server provides interfaces in bytes format, please
    transform them to string in the implementation of this function.

Args:
        request: str, a string of the RPC request.

Returns:
        A string of the RPC response.

Raises:
        errors.ProtocolError: something went wrong when exchanging data with the
        server.


def _StatelessRandomGammaV3Grad(op: ops.Operation, grad):
    shape = op.inputs[0]
    alpha = op.inputs[4]
    sample = op.outputs[0]
    with ops.control_dependencies([grad]):
        return (None, None, None, None, _StatelessGammaGradAlpha(shape, alpha, sample, grad))

Returns the gradient of a Gamma sample w.r.t. alpha.

    The gradient is computed using implicit differentiation
    (Figurnov et al., 2018).

Args:
        op: A `StatelessRandomGamma` operation. We assume that the inputs to the
        operation are `shape`, `key`, `counter`, `alg`, and `alpha` tensors, and
        the output is the `sample` tensor.
        grad: The incoming gradient `dloss / dsample` of the same shape as
        `op.outputs[0]`.

Returns:
        A `Tensor` with derivatives `dloss / dalpha`.

    References:
        Implicit Reparameterization Gradients:
        [Figurnov et al., 2018]
        (http://papers.nips.cc/paper/7326-implicit-reparameterization-gradients)
        ([pdf]
        (http://papers.nips.cc/paper/7326-implicit-reparameterization-gradients.pdf))


def _create_hash_str(self, input_arg, output_arg, node_def):
    hasher = hashlib.sha1()

    def update_num(n):
        hasher.update(compat.as_bytes('%x' % n))

    def update_str(s):
        update_num(len(s))
        hasher.update(compat.as_bytes(s))

    def update_strs(slist):
        update_num(len(slist))
        for s in slist:
            update_str(s)
    for adef in input_arg:
        update_str(adef.SerializeToString())
    for adef in output_arg:
        update_str(adef.SerializeToString())
    for n in sorted(node_def, key=lambda n: n.name):
        update_str(n.name)
        update_str(n.op)
        update_strs(n.input)
        update_num(len(n.attr))
        for k in sorted(n.attr):
            update_str(k)
            update_str(n.attr[k].SerializeToString())
    return hasher.hexdigest()[:8]

Creates an 8-character string unique to this input.

Args:
        input_arg: the input_arg field of an OpDef
                (e.g. self._definition.signature.input_arg)
        output_arg: the output_arg field of an OpDef
                (e.g. self._definition.signature.output_arg)
        node_def: the node_def field of a FunctionDef
                (e.g. self._definition.node_def)

Returns:
        The unique string for this input


def parallel_walk(node, other):
    if isinstance(node, (list, tuple)):
        node_stack = list(node)
    else:
        node_stack = [node]
    if isinstance(other, (list, tuple)):
        other_stack = list(other)
    else:
        other_stack = [other]
    while node_stack and other_stack:
        assert len(node_stack) == len(other_stack)
        n = node_stack.pop()
        o = other_stack.pop()
        if not isinstance(n, (ast.AST, gast.AST, str)) and n is not None or (not isinstance(o, (ast.AST, gast.AST, str)) and n is not None) or n.__class__.__name__ != o.__class__.__name__:
            raise ValueError('inconsistent nodes: {} ({}) and {} ({})'.format(n, n.__class__.__name__, o, o.__class__.__name__))
        yield (n, o)
        if isinstance(n, str):
            assert isinstance(o, str), 'The check above should have ensured this'
            continue
        if n is None:
            assert o is None, 'The check above should have ensured this'
            continue
        for f in n._fields:
            n_child = getattr(n, f, None)
            o_child = getattr(o, f, None)
            if f.startswith('__') or n_child is None or o_child is None:
                continue
            if isinstance(n_child, (list, tuple)):
                if not isinstance(o_child, (list, tuple)) or len(n_child) != len(o_child):
                    raise ValueError('inconsistent values for field {}: {} and {}'.format(f, n_child, o_child))
                node_stack.extend(n_child)
                other_stack.extend(o_child)
            elif isinstance(n_child, (gast.AST, ast.AST)):
                node_stack.append(n_child)
                other_stack.append(o_child)
            elif n_child != o_child:
                raise ValueError('inconsistent values for field {}: {} and {}'.format(f, n_child, o_child))

Walks two ASTs in parallel.

    The two trees must have identical structure.

Args:
        node: Union[ast.AST, Iterable[ast.AST]]
        other: Union[ast.AST, Iterable[ast.AST]]

Yields:
        Tuple[ast.AST, ast.AST]

Raises:
        ValueError: if the two trees don't have identical structure.


def _skip_remaining_tests(self, exception):
    for test_name in self.results.requested:
        if not self.results.is_test_executed(test_name):
            test_record = records.TestResultRecord(test_name, self.TAG)
            test_record.test_skip(exception)
            self.results.add_record(test_record)
            self.summary_writer.dump(test_record.to_dict(), records.TestSummaryEntryType.RECORD)

Marks any requested test that has not been executed in a class as
    skipped.

    This is useful for handling abort class signal.

Args:
        exception: The exception object that was thrown to trigger the
        skip.


def _read_protocol_line(self):
    while True:
        line = self._proc.stdout.readline().decode('utf-8')
        if not line:
            raise jsonrpc_client_base.AppStartError(self._ad, 'Unexpected EOF waiting for app to start')
        line = line.strip()
        if line.startswith('INSTRUMENTATION_RESULT:') or line.startswith('SNIPPET '):
            self.log.debug('Accepted line from instrumentation output: "%s"', line)
            return line
        self.log.debug('Discarded line from instrumentation output: "%s"', line)

Reads the next line of instrumentation output relevant to snippets.

    This method will skip over lines that don't start with 'SNIPPET' or
    'INSTRUMENTATION_RESULT'.

Returns:
        (str) Next line of snippet-related instrumentation output, stripped.

Raises:
        jsonrpc_client_base.AppStartError: If EOF is reached without any
        protocol lines being read.


def _get_unpad_data(attention_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, int]:
    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
    max_seqlen_in_batch = seqlens_in_batch.max().item()
    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
    return (indices, cu_seqlens, max_seqlen_in_batch)

Retrieves indexing data required to repad unpadded (ragged) tensors.

Args:
        attention_mask (`torch.Tensor`):
            Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid.

Returns:
        indices (`torch.Tensor`):
            The indices of non-masked tokens from the flattened input sequence.
        cu_seqlens (`torch.Tensor`):
            The cumulative sequence lengths, used to index into ragged (unpadded) tensors. `cu_seqlens` shape is (batch_size + 1,).
        max_seqlen_in_batch (`int`):
            Maximum sequence length in batch.


def load_model(filepath, custom_objects=None, compile=True, safe_mode=True):
    is_keras_zip = str(filepath).endswith('.keras') and zipfile.is_zipfile(filepath)
    is_keras_dir = file_utils.isdir(filepath) and file_utils.exists(file_utils.join(filepath, 'config.json'))
    is_hf = str(filepath).startswith('hf://')
    if file_utils.is_remote_path(filepath) and (not file_utils.isdir(filepath)) and (not is_keras_zip) and (not is_hf):
        local_path = file_utils.join(saving_lib.get_temp_dir(), os.path.basename(filepath))
        file_utils.copy(filepath, local_path)
        if zipfile.is_zipfile(local_path):
            filepath = local_path
            is_keras_zip = True
    if is_keras_zip or is_keras_dir or is_hf:
        return saving_lib.load_model(filepath, custom_objects=custom_objects, compile=compile, safe_mode=safe_mode)
    if str(filepath).endswith(('.h5', '.hdf5')):
        return legacy_h5_format.load_model_from_hdf5(filepath, custom_objects=custom_objects, compile=compile)
    elif str(filepath).endswith('.keras'):
        raise ValueError(f'File not found: filepath={filepath}. Please ensure the file is an accessible `.keras` zip file.')
    else:
        raise ValueError(f"File format not supported: filepath={filepath}. Keras 3 only supports V3 `.keras` files and legacy H5 format files (`.h5` extension). Note that the legacy SavedModel format is not supported by `load_model()` in Keras 3. In order to reload a TensorFlow SavedModel as an inference-only layer in Keras 3, use `keras.layers.TFSMLayer({filepath}, call_endpoint='serving_default')` (note that your `call_endpoint` might have a different name).")

Loads a model saved via `model.save()`.

Args:
        filepath: `str` or `pathlib.Path` object, path to the saved model file.
        custom_objects: Optional dictionary mapping names
            (strings) to custom classes or functions to be
            considered during deserialization.
        compile: Boolean, whether to compile the model after loading.
        safe_mode: Boolean, whether to disallow unsafe `lambda` deserialization.
            When `safe_mode=False`, loading an object has the potential to
            trigger arbitrary code execution. This argument is only
            applicable to the Keras v3 model format. Defaults to `True`.

Returns:
        A Keras model instance. If the original model was compiled,
        and the argument `compile=True` is set, then the returned model
        will be compiled. Otherwise, the model will be left uncompiled.

Example:
    ```python
    model = keras.Sequential([
        keras.layers.Dense(5, input_shape=(3,)),
        keras.layers.Softmax()])
    model.save("model.keras")
    loaded_model = keras.saving.load_model("model.keras")
    x = np.random.random((10, 3))
    assert np.allclose(model.predict(x), loaded_model.predict(x))
    ```

    Note that the model variables may have different name values
    (`var.name` property, e.g. `"dense_1/kernel:0"`) after being reloaded.
    It is recommended that you use layer attributes to
    access specific variables, e.g. `model.get_layer("dense_1").kernel`.


def get_platform():
    global PLATFORM
    cmd = 'uname'
    out, err = run_shell_cmd(cmd)
    platform_detected = out.strip().lower()
    if platform_detected != 'linux':
        if err and FLAGS.debug:
            print('Error in detecting platform:\n %s' % str(err))
        print('Error: Detected unsupported operating system.\nStopping...')
        sys.exit(1)
    else:
        PLATFORM = platform_detected
    return PLATFORM

Retrieves platform information.

    Currently the script only support linux. If other platoforms such as Windows
    or MacOS is detected, it throws an error and terminates.

Returns:
        String that is platform type.
        e.g. 'linux'


def solve(self):
    hierarchy = type_match.get_all_subclasses([self.ast, self.builtins])
    factory_protocols = type_match.TypeMatch(hierarchy)
    factory_partial = type_match.TypeMatch(hierarchy)
    solver_protocols = factory_protocols.solver
    solver_partial = factory_partial.solver
    unknown_classes = set()
    partial_classes = set()
    complete_classes = set()
    for cls in self.ast.classes:
        if is_unknown(cls):
            solver_protocols.register_variable(cls.name)
            solver_partial.register_variable(cls.name)
            unknown_classes.add(cls)
        elif is_partial(cls):
            partial_classes.add(cls)
        else:
            complete_classes.add(cls)
    protocol_classes_and_aliases = set(self.protocols.classes)
    for alias in self.protocols.aliases:
        if not isinstance(alias.type, pytd.AnythingType) and alias.name != 'protocols.Protocol':
            protocol_classes_and_aliases.add(alias.type.cls)
    for protocol in protocol_classes_and_aliases:
        for unknown in unknown_classes:
            self.match_unknown_against_protocol(factory_protocols, solver_protocols, unknown, protocol)
    for complete in complete_classes.union(self.builtins.classes):
        for partial in partial_classes:
            if escape.unpack_partial(partial.name) == complete.name:
                self.match_partial_against_complete(factory_partial, solver_partial, partial, complete)
    partial_functions = set()
    complete_functions = set()
    for f in self.ast.functions:
        if is_partial(f):
            partial_functions.add(f)
        else:
            complete_functions.add(f)
    for partial in partial_functions:
        for complete in complete_functions.union(self.builtins.functions):
            if escape.unpack_partial(partial.name) == complete.name:
                self.match_call_record(factory_partial, solver_partial, partial, complete)
    log.info('=========== Equations to solve =============\n%s', solver_protocols)
    log.info('=========== Equations to solve (end) =======')
    solved_protocols = solver_protocols.solve()
    log.info('=========== Call trace equations to solve =============\n%s', solver_partial)
    log.info('=========== Call trace equations to solve (end) =======')
    solved_partial = solver_partial.solve()
    merged_solution = {}
    for unknown in itertools.chain(solved_protocols, solved_partial):
        if unknown in solved_protocols and unknown in solved_partial:
            merged_solution[unknown] = solved_protocols[unknown].union(solved_partial[unknown])
            merged_solution[unknown].discard('?')
        elif unknown in solved_protocols:
            merged_solution[unknown] = solved_protocols[unknown]
        else:
            merged_solution[unknown] = solved_partial[unknown]
    return merged_solution

Solve the equations generated from the pytd.

Returns:
        A dictionary (str->str), mapping unknown class names to known class names.

Raises:
        AssertionError: If we detect an internal error.


def _mark_func_graph_as_unsaveable(graph, learning_phase):
    if graph.building_function and is_placeholder(learning_phase):
        graph.mark_as_unsaveable('The keras learning phase placeholder was used inside a function. Exporting placeholders is not supported when saving out a SavedModel. Please call `tf.keras.backend.set_learning_phase(0)` in the function to set the learning phase to a constant value.')

Mark func graph as unsaveable due to use of symbolic keras learning phase.

    Functions that capture the symbolic learning phase cannot be exported to
    SavedModel. Mark the funcgraph as unsaveable, so that an error will be raised
    if it is exported.

Args:
        graph: Graph or FuncGraph object.
        learning_phase: Learning phase placeholder or int defined in the graph.


def format_speech_generation_kwargs(kwargs):
    kwargs_text = {}
    kwargs_speech = {}
    for key, value in kwargs.items():
        if key.startswith('text_'):
            key = key[len('text_'):]
            kwargs_text[key] = value
        elif key.startswith('speech_'):
            key = key[len('speech_'):]
            kwargs_speech[key] = value
        elif key == 'generation_config':
            kwargs_text[key] = value
        else:
            if key not in kwargs_text:
                kwargs_text[key] = value
            if key not in kwargs_speech:
                kwargs_speech[key] = value
    return (kwargs_text, kwargs_speech)

Format kwargs for SeamlessM4T models that generate speech, attribute kwargs to either the text generation or the
    speech generation models.

Args:
        kwargs (`dict`)`:
            Keyword arguments are of two types:

                - Without a prefix, they will be entered as `**kwargs` for the `generate` method of each sub-model,
                except for `decoder_input_ids` which will only be passed through the text components.
                - With a *text_* or *speech_* prefix, they will be input for the `generate` method of the
                text model and speech model respectively. It has the priority over the keywords without a prefix.

                This means you can, for example, specify a generation strategy for one generation but not for the
                other.


def assign_add(self, delta, use_locking=None, name=None, read_value=True):
    with _handle_graph(self.handle), self._assign_dependencies():
        assign_add_op = gen_resource_variable_ops.assign_add_variable_op(self.handle, ops.convert_to_tensor(delta, dtype=self.dtype), name=name)
    if read_value:
        return self._lazy_read(assign_add_op)
    return assign_add_op

Adds a value to this variable.

Args:
        delta: A `Tensor`. The value to add to this variable.
        use_locking: If `True`, use locking during the operation.
        name: The name to use for the operation.
        read_value: A `bool`. Whether to read and return the new value of the
        variable or not.

Returns:
        If `read_value` is `True`, this method will return the new value of the
        variable after the assignment has completed. Otherwise, when in graph mode
        it will return the `Operation` that does the assignment, and when in eager
        mode it will return `None`.


def register_proto_function(collection_name, proto_type=None, to_proto=None, from_proto=None) -> None:
    if to_proto and (not callable(to_proto)):
        raise TypeError('to_proto must be callable.')
    if from_proto and (not callable(from_proto)):
        raise TypeError('from_proto must be callable.')
    _proto_function_registry.register((proto_type, to_proto, from_proto), collection_name)

Registers `to_proto` and `from_proto` functions for collection_name.

    `to_proto` function converts a Python object to the corresponding protocol
    buffer, and returns the protocol buffer.

    `from_proto` function converts protocol buffer into a Python object, and
    returns the object..

Args:
        collection_name: Name of the collection.
        proto_type: Protobuf type, such as `saver_pb2.SaverDef`,
        `variable_pb2.VariableDef`, `queue_runner_pb2.QueueRunnerDef`..
        to_proto: Function that implements Python object to protobuf conversion.
        from_proto: Function that implements protobuf to Python object conversion.


def counter(namespace: Union[Type, str], name: str) -> 'Metrics.DelegatingCounter':
    namespace = Metrics.get_namespace(namespace)
    return Metrics.DelegatingCounter(MetricName(namespace, name))

Obtains or creates a Counter metric.

Args:
        namespace: A class or string that gives the namespace to a metric
        name: A string that gives a unique name to a metric

Returns:
        A Counter object.


def run(code: str, *, global_vars: Optional[Dict[str, Any]]=None, permission: Optional[permissions.CodePermission]=None, returns_stdout: bool=False, outputs_intermediate: bool=False, sandbox: Optional[bool]=None, timeout: Optional[float]=None) -> Union[Any, Dict[str, Any]]:
    return maybe_sandbox_call(evaluate, code=code, global_vars=global_vars, permission=permission, returns_stdout=returns_stdout, outputs_intermediate=outputs_intermediate, sandbox=sandbox, timeout=timeout)

Executes Python code.

    Features:
        * Fine-grained execution policy for limiting what APIs could be executed.
        This eliminates the need for sandboxing.
        * It exposes both the final results and intermediate results (variables).

Args:
        code: Python code to run.
        global_vars: An optional dict of
        permission: Permission for the Python code to run.
        returns_stdout: If True, the stdout (a str) will be returned.
        outputs_intermediate: Applicable when returns_stdout is False. If True,
        intermediate output will be outputted as a dict, with the last line's
        value accessible by key '__result__' and the std output accessible by
        key '__stdout__'. Otherwise the value of the last line will be returned.
        sandbox: If True, run code in sandbox; If False, run code in current
        process. If None, run in sandbox first, if the output could not be
        serialized and pass to current process, run the code again in current
        process.
        timeout: Execution timeout in seconds. If None, wait the code the complete.

Returns:
        The value of the last line of the code block. Or a dict of variable
        names of all locals to their evaluated values as the output of the code to
        run. The value for the last line can be accessed by key '__result__'. Or the
        stdout as a str.

Raises:
        TimeoutError: If the execution time exceeds the timeout.
        Exception: Exception  that are raised from the code.


def scramble_generating_matrices(generating_matrices: types.IntTensor, scrambling_matrices: types.IntTensor, num_digits: types.IntTensor, validate_args: bool=False, dtype: tf.DType=None, name: str=None):
    with tf.name_scope(name or 'scramble_generating_matrices'):
        if dtype is None:
            generating_matrices = tf.convert_to_tensor(generating_matrices)
        dtype = dtype or generating_matrices.dtype
        num_digits = tf.convert_to_tensor(num_digits, dtype=dtype, name='num_digits')
        generating_matrices = tf.cast(generating_matrices, dtype=dtype, name='generating_matrices')
        scrambling_matrices = tf.cast(scrambling_matrices, dtype=dtype, name='scrambling_matrices')
        control_deps = []
        if validate_args:
            control_deps.append(tf.debugging.assert_equal(tf.rank(generating_matrices), tf.rank(scrambling_matrices), message='input matrices must have the same rank'))
            control_deps.append(tf.debugging.assert_positive(num_digits, message='num_digits must be positive'))
        with tf.control_dependencies(control_deps):

            def loop_predicate_fn(matrix, shift):
                del matrix
                return shift < num_digits

            def loop_body_fn(matrix, shift):
                shifted_scrambling_matrices = tf.bitwise.right_shift(tf.gather(scrambling_matrices, [shift], axis=1), shift)
                updated_matrix = tf.bitwise.bitwise_xor(matrix, utils.filter_tensor(shifted_scrambling_matrices, generating_matrices, num_digits - 1 - shift))
                return (updated_matrix, shift + 1)
            matrix, _ = tf.while_loop(loop_predicate_fn, loop_body_fn, loop_vars=(tf.zeros_like(generating_matrices), tf.constant(0, dtype=dtype)), maximum_iterations=tf.cast(num_digits, tf.int32))
            return matrix

Scrambles a generating matrix.

    #### Examples

    ```python
    import tf_quant_finance as tff

    # Example: Scrambling the 2D Sobol generating matrices.

    dim = 2
    num_results = 1000
    num_digits = 10
    seed = (2, 3)

    tff.math.qmc.scramble_generating_matrices(
        tff.math.qmc.sobol_generating_matrices(dim, num_results, num_digits),
        tff.math.qmc.random_scrambling_matrices(dim, num_digits, seed=seed),
        num_digits)
    # ==> tf.Tensor([
    #             [586, 505, 224, 102,  34,  31,  13,   6,   2,   1],
    #             [872, 695, 945, 531, 852, 663, 898, 568, 875, 693],
    #         ], shape=(2, 10), dtype=int32)
    ```

Args:
        generating_matrices: Positive scalar `Tensor` of integers.
        scrambling_matrices: Positive Scalar `Tensor` of integers with the same
        `shape` as `generating_matrices`.
        num_digits: Positive scalar `Tensor` of integers with rank 0. The base-2
        precision of the points which can be sampled from `generating_matrices`.
        validate_args: Python `bool` indicating whether to validate arguments.
        Default value: `False`.
        dtype: Optional `dtype`. The `dtype` of the output `Tensor` (either `int32`
        or `int64`).
        Default value: `None` which maps to `generating_matrices.dtype`.
        name: Python `str` name prefixed to ops created by this function.
        Default value: `None` which maps to `scramble_generating_matrices`.

Returns:
        A `Tensor` with the same `shape` and `dtype` as `generating_matrices`.


def resize(self, image: torch.Tensor, size: SizeDict, patch_size: SizeDict, interpolation: 'F.InterpolationMode'=None, **kwargs) -> torch.Tensor:
    interpolation = interpolation if interpolation is not None else F.InterpolationMode.BILINEAR
    if size.longest_edge:
        size = (size.longest_edge, size.longest_edge)
    elif size.height and size.width:
        size = (size.height, size.width)
    else:
        raise ValueError("size must contain either 'longest_edge' or 'height' and 'width'.")
    if patch_size.height and patch_size.width:
        patch_size = (patch_size.height, patch_size.width)
    else:
        raise ValueError("patch_size must contain either 'shortest_edge' or 'height' and 'width'.")
    output_size = get_resize_output_image_size(image, size=size, patch_size=patch_size)
    return F.resize(image, size=output_size, interpolation=interpolation, **kwargs)

Resize an image. The shortest edge of the image is resized to size["shortest_edge"], with the longest edge
    resized to keep the input aspect ratio.

Args:
        image (`torch.Tensor`):
            Image to resize.
        size (`SizeDict`):
            Dict containing the longest possible edge of the image.
        patch_size (`SizeDict`):
            Patch size used to calculate the size of the output image.
        interpolation (`InterpolationMode`, *optional*, defaults to `InterpolationMode.BILINEAR`):
            Resampling filter to use when resiizing the image.


def forward(self, pixel_values: torch.FloatTensor, labels: Optional[torch.FloatTensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[Tuple[torch.Tensor], DepthEstimatorOutput]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    outputs = self.glpn(pixel_values, output_attentions=output_attentions, output_hidden_states=True, return_dict=return_dict)
    hidden_states = outputs.hidden_states if return_dict else outputs[1]
    out = self.decoder(hidden_states)
    predicted_depth = self.head(out)
    loss = None
    if labels is not None:
        loss_fct = SiLogLoss()
        loss = loss_fct(predicted_depth, labels)
    if not return_dict:
        if output_hidden_states:
            output = (predicted_depth,) + outputs[1:]
        else:
            output = (predicted_depth,) + outputs[2:]
        return (loss,) + output if loss is not None else output
    return DepthEstimatorOutput(loss=loss, predicted_depth=predicted_depth, hidden_states=outputs.hidden_states if output_hidden_states else None, attentions=outputs.attentions)

labels (`torch.FloatTensor` of shape `(batch_size, height, width)`, *optional*):
    Ground truth depth estimation maps for computing the loss.

Example:
    ```python
    >>> from transformers import AutoImageProcessor, GLPNForDepthEstimation
    >>> import torch
    >>> import numpy as np
    >>> from PIL import Image
    >>> import requests

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

    >>> image_processor = AutoImageProcessor.from_pretrained("vinvino02/glpn-kitti")
    >>> model = GLPNForDepthEstimation.from_pretrained("vinvino02/glpn-kitti")

    >>> # prepare image for the model
    >>> inputs = image_processor(images=image, return_tensors="pt")

    >>> with torch.no_grad():
    ...     outputs = model(**inputs)

    >>> # interpolate to original size
    >>> post_processed_output = image_processor.post_process_depth_estimation(
    ...     outputs,
    ...     target_sizes=[(image.height, image.width)],
    ... )

    >>> # visualize the prediction
    >>> predicted_depth = post_processed_output[0]["predicted_depth"]
    >>> depth = predicted_depth * 255 / predicted_depth.max()
    >>> depth = depth.detach().cpu().numpy()
    >>> depth = Image.fromarray(depth.astype("uint8"))
    ```


def _update_sample_weight_modes(self, sample_weights=None):
    if not self._is_compiled:
        return
    if sample_weights and any((s is not None for s in sample_weights)):
        for endpoint in self._training_endpoints:
            endpoint.sample_weight_mode = endpoint.sample_weight_mode or 'samplewise'
    else:
        for endpoint in self._training_endpoints:
            endpoint.sample_weight_mode = None

Updates sample weight modes based on training/eval inputs.

    Sample weight placeholders will be created for all or no outputs
    based on whether sample_weight is provided for any output.

    If model contains `_sample_weight_modes` we check if the input
    `sample_weights` corresponds to the sample weight modes.
        1. Set sample weight mode to be 'temporal' for output i, if `compile`
        sample_weight_mode was set to `temporal` and sample weight inputs
        are given for one or more outputs.
        2. Set sample weight mode to be 'samplewise' for output i, if `compile`
        sample_weight_mode was not set and sample weight inputs are given for
        one or more outputs.
        3. Reset sample weight mode to None for output i if sample weight mode
        was set but there is no sample weight input.

Args:
        sample_weights: List of sample weights of the same length as model outputs
        or None.


def materialize_value_sets(self, value_set_protos: Iterable[value_set_pb2.ValueSet], batch_size: int=500) -> None:
    self._create_valueset_codes_table_if_not_exists()
    columns = [sqlalchemy.column('valueseturi'), sqlalchemy.column('valuesetversion'), sqlalchemy.column('system'), sqlalchemy.column('code')]
    sa_table = sqlalchemy.table(self._value_set_codes_table, *columns)
    queries = value_set_tables.valueset_codes_insert_statement_for(value_set_protos, sa_table, batch_size=batch_size)
    for query in queries:
        with self._client.connect() as curs:
            curs.execute(query)

Materialize the given value sets into the value_set_codes_table.

    Then writes these expanded codes into the database
    named after the `value_set_codes_table` provided at class initialization.
    Builds a valueset_codes table as described by
    https://github.com/FHIR/sql-on-fhir/blob/master/sql-on-fhir.md#valueset-support

    The table will be created if it does not already exist.

    The function will avoid inserting duplicate rows if some of the codes are
    already present in the given table. It will not attempt to perform an
    'upsert' or modify any existing rows.

    Note that value sets provided to this function should already be expanded,
    in that they contain the code values to write. Users should also see
    `materialize_value_set_expansion` below to retrieve an expanded set from
    a terminology server.

Args:
        value_set_protos: An iterable of FHIR ValueSet protos.
        batch_size: The maximum number of rows to insert in a single query.


def loop(self, timer_interval_secs, target, args=None, kwargs=None):
    looper = coordinator.LooperThread(self._coord, timer_interval_secs, target=target, args=args, kwargs=kwargs)
    looper.start()
    return looper

Start a LooperThread that calls a function periodically.

    If `timer_interval_secs` is None the thread calls `target(*args, **kwargs)`
    repeatedly.  Otherwise it calls it every `timer_interval_secs`
    seconds.  The thread terminates when a stop is requested.

    The started thread is added to the list of threads managed by the supervisor
    so it does not need to be passed to the `stop()` method.

Args:
        timer_interval_secs: Number. Time boundaries at which to call `target`.
        target: A callable object.
        args: Optional arguments to pass to `target` when calling it.
        kwargs: Optional keyword arguments to pass to `target` when calling it.

Returns:
        The started thread.


def __init__(self, dim, means=0.0, volatilities=1.0, corr_matrix=None, dtype=None, name=None):
    self._name = name or 'multivariate_geometric_brownian_motion'
    with tf.name_scope(self._name):
        self._means = tf.convert_to_tensor(means, dtype=dtype, name='means')
        self._dtype = self._means.dtype
        self._vols = tf.convert_to_tensor(volatilities, dtype=self._dtype, name='volatilities')
        self._dim = dim
        if corr_matrix is None:
            self._corr_matrix = None
        else:
            self._corr_matrix = tf.convert_to_tensor(corr_matrix, dtype=self._dtype, name='corr_matrix')
            if self._corr_matrix.shape.as_list() != [dim, dim]:
                raise ValueError('`corr_matrix` must be of shape [{0}, {0}] but is of shape {1}'.format(dim, self._corr_matrix.shape.as_list()))

Initializes the Multivariate Geometric Brownian Motion.

Args:
        dim: A Python scalar. The dimensionality of the process
        means:  A real `Tensor` of shape broadcastable to `[dim]`.
        Corresponds to the vector of means of the GBM components `X_i`.
        Default value: 0.0.
        volatilities: A `Tensor` of the same `dtype` as `means` and of shape
        broadcastable to `[dim]`. Corresponds to the volatilities of the GBM
        components `X_i`.
        Default value: 1.0.
        corr_matrix: An optional `Tensor` of the same `dtype` as `means` and of
        shape `[dim, dim]`. Correlation of the GBM components `W_i`.
        Default value: `None` which maps to a process with
        independent GBM components `X_i`.
        dtype: The default dtype to use when converting values to `Tensor`s.
        Default value: `None` which means that default dtypes inferred by
        TensorFlow are used.
        name: Python string. The name to give to the ops created by this class.
        Default value: `None` which maps to the default name
        'multivariate_geometric_brownian_motion'.

Raises:
        ValueError: If `corr_matrix` is supplied and is not of shape `[dim, dim]`


def forward(self, patch_input: torch.Tensor):
    if self.mask_type == 'random':
        masked_input, mask = random_masking(inputs=patch_input, mask_ratio=self.random_mask_ratio, unmasked_channel_indices=self.unmasked_channel_indices, channel_consistent_masking=self.channel_consistent_masking, mask_value=self.mask_value)
    elif self.mask_type == 'forecast':
        masked_input, mask = forecast_masking(inputs=patch_input, num_forecast_mask_patches=self.num_forecast_mask_patches, unmasked_channel_indices=self.unmasked_channel_indices, mask_value=self.mask_value)
    else:
        raise ValueError(f'Invalid mask type {self.mask_type}.')
    mask = mask.bool()
    return (masked_input, mask)

Parameters:
    patch_input (`torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`, *required*):
        Patch input

Returns:
    masked_input (`torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`)
        Masked patched input
    mask (`torch.Tensor` of shape `(batch_size, num_channels, num_patches)`)
        Bool tensor indicating True on masked points


def to_variable(self, node: 'cfg.CFGNode') -> 'cfg.Variable':
    return self.ctx.program.NewVariable([self], source_set=[], where=node)

Build a variable out of this abstract value.

Args:
        node: The current CFG node.

Returns:
        A cfg.Variable.


def flip_channel_order(self, image):
    self._ensure_format_supported(image)
    if isinstance(image, PIL.Image.Image):
        image = self.to_numpy_array(image)
    return image[::-1, :, :]

Flips the channel order of `image` from RGB to BGR, or vice versa. Note that this will trigger a conversion of
    `image` to a NumPy array if it's a PIL Image.

Args:
        image (`PIL.Image.Image` or `np.ndarray` or `torch.Tensor`):
            The image whose color channels to flip. If `np.ndarray` or `torch.Tensor`, the channel dimension should
            be first.


def _format_value(cls, value, type_):
    res = value
    if type_ == 'CLASS':
        res = '{}.{}'.format(value.__module__, value.__name__)
    elif type_ == 'DURATION':
        res = value.total_seconds() * 1000
    elif type_ == 'TIMESTAMP':
        res = calendar.timegm(value.timetuple()) * 1000 + value.microsecond // 1000
    return res

Returns the API representation of a value given its type.

Args:
        value: The value of the item that needs to be shortened.
        type_(string): The type of the value.

Returns:
        A formatted value in the form of a float, int, or string.


def find_labels(model_class):
    model_name = model_class.__name__
    framework = infer_framework(model_class)
    if framework == 'tf':
        signature = inspect.signature(model_class.call)
    elif framework == 'pt':
        signature = inspect.signature(model_class.forward)
    else:
        signature = inspect.signature(model_class.__call__)
    if 'QuestionAnswering' in model_name:
        return [p for p in signature.parameters if 'label' in p or p in ('start_positions', 'end_positions')]
    else:
        return [p for p in signature.parameters if 'label' in p]

Find the labels used by a given model.

Args:
        model_class (`type`): The class of the model.


def get_merged_audio_embeddings(self, input_ids: torch.Tensor, audio_features: torch.Tensor, input_features_mask: Optional[torch.Tensor]=None) -> torch.Tensor:
    is_audio_index = input_ids == self.config.audio_token_id
    llm_input_ids = torch.where(is_audio_index, 0, input_ids)
    inputs_embeds = self.language_model.get_input_embeddings()(llm_input_ids)
    special_audio_mask = is_audio_index.unsqueeze(-1)
    audio_features = audio_features.to(inputs_embeds.device, inputs_embeds.dtype)
    if input_features_mask is not None:
        if torch.all(is_audio_index.int().sum(dim=1) != input_features_mask.int().sum(dim=1)).item():
            raise ValueError('Number of audio tokens does not match number of audio features')
        audio_features = audio_features[input_features_mask]
    inputs_embeds = inputs_embeds.masked_scatter(special_audio_mask, audio_features)
    return inputs_embeds

Adds the audio token to the model's LLM vocabulary so that we can pass it
    through the tokenizer; it's assumed that the embeddings corresponding to the
    <|audio|> token will be clobbered with speech features.

Args:
        input_ids (`torch.Tensor`):
            Input IDs containing one or more audio tokens.
        audio_features (`torch.Tensor`):
            Audio features to be masked into the language embeddings to form multimodal embeddings.
        input_features_mask (`torch.Tensor`, *optional*, defaults to `None`)
            Mask to be applied to audio features prior to scattering into the language embeddings.


def predict_signature_def(inputs, outputs):
    if inputs is None or not inputs:
        raise ValueError('Prediction `inputs` cannot be None or empty.')
    if outputs is None or not outputs:
        raise ValueError('Prediction `outputs` cannot be None or empty.')
    signature_inputs = {key: utils.build_tensor_info(tensor) for key, tensor in inputs.items()}
    signature_outputs = {key: utils.build_tensor_info(tensor) for key, tensor in outputs.items()}
    signature_def = build_signature_def(signature_inputs, signature_outputs, signature_constants.PREDICT_METHOD_NAME)
    return signature_def

Creates prediction signature from given inputs and outputs.

    This function produces signatures intended for use with the TensorFlow Serving
    Predict API (tensorflow_serving/apis/prediction_service.proto). This API
    imposes no constraints on the input and output types.

Args:
        inputs: dict of string to `Tensor`.
        outputs: dict of string to `Tensor`.

Returns:
        A prediction-flavored signature_def.

Raises:
        ValueError: If inputs or outputs is `None`.


def is_copy_node(node_name):
    return node_name.startswith('__copy_')

Determine whether a node name is that of a debug Copy node.

    Such nodes are inserted by TensorFlow core upon request in
    RunOptions.debug_options.debug_tensor_watch_opts.

Args:
        node_name: Name of the node.

Returns:
        A bool indicating whether the input argument is the name of a debug Copy
        node.


def send_log_message(self, message: LogMessage) -> None:
    self.send_log_messages([message])

Sends the log message to BigQuery.

Args:
        * message: LogMessage dictionary

Returns:
        * None

Raises:
        * RuntimeError: if BigQuery insert fails


def view_options(**kwargs) -> Iterator[Dict[str, Any]]:
    parent_options = utils.thread_local_peek(_TLS_KEY_VIEW_OPTIONS, {})
    options = utils.merge([parent_options, kwargs])
    utils.thread_local_push(_TLS_KEY_VIEW_OPTIONS, options)
    try:
        yield options
    finally:
        utils.thread_local_pop(_TLS_KEY_VIEW_OPTIONS)

Context manager to inject rendering args to view.

Example:
        with pg.view_options(enable_summary_tooltip=False):
        MyObject().to_html()

Args:
        **kwargs: Keyword arguments for View.render method.

Yields:
        The merged keyword arguments.


def wait_until_finish(self, duration=None):
    if not PipelineState.is_terminal(self._state):
        raise NotImplementedError()

Waits until the pipeline finishes and returns the final status.

Args:
        duration (int): The time to wait (in milliseconds) for job to finish.
        If it is set to :data:`None`, it will wait indefinitely until the job
        is finished.

Raises:
        IOError: If there is a persistent problem getting job
        information.
        NotImplementedError: If the runner does not support this
        operation.

Returns:
        The final state of the pipeline, or :data:`None` on timeout.


def aggregate_and_return_name_for_output(self, fused_op_name, output_index, out_graphdef):
    del fused_op_name, output_index, out_graphdef
    raise RuntimeError('Unimplemented abstract method.')

Add node(s) to graph representing output operands and returns type.

Args:
        fused_op_name: name of the fused op stub name.
        output_index: Output index that we are currently processing from stub.
        out_graphdef: The destination graphdef we are currently building up.

Returns:
        The datatype of this identity.

Raises:
        RuntimeError: if the method is not implemented.


def interpolate(self, x: types.RealTensor, y: types.RealTensor, name: str=None):
    name = name or self._name + '_interpolate'
    with tf.name_scope(name):
        x = tf.convert_to_tensor(x, dtype=self._dtype, name='x')
        y = tf.convert_to_tensor(y, dtype=self._dtype, name='y')
        y = tf.expand_dims(y, axis=-2)
        xy = cubic.interpolate(y, self._spline_yz, name='interpolation_in_y_direction')
        xy_rank = xy.shape.rank
        perm = [xy_rank - 1] + list(range(xy_rank - 1))
        yx = tf.transpose(xy, perm=perm)
        perm_original = list(range(1, xy_rank)) + [0]
        x = tf.expand_dims(tf.transpose(x, [xy_rank - 2] + list(range(xy_rank - 2))), axis=-1)
        z_values = linear.interpolate(x, self._xdata, yx)
        return tf.squeeze(tf.transpose(z_values, perm=perm_original), axis=-2)

Performs 2-D interpolation on a specified set of points.

Args:
        x: Real-valued `Tensor` of shape `batch_shape + [num_points]`.
        Defines the x-coordinates at which the interpolation should be
        performed. Note that `batch_shape` should be the same as in the
        underlying data.
        y: A `Tensor` of the same shape and `dtype` as `x`.
        Defines the y-coordinates at which the interpolation should be
        performed.
        name: Python `str` name prefixed to ops created by this function.
        Default value: `None` which is mapped to the default name
        `interpolate`.

Returns:
        A `Tensor` of the same shape and `dtype` as `x`. Represents the
        interpolated values of the function on for the coordinates
        `(x, y)`.


def elementwise_division(func):

    @functools.wraps(func)
    def sparse_wrapper(x1, x2):
        if isinstance(x1, tf.SparseTensor):
            if isinstance(x2, tf.SparseTensor):
                x1 = sparse_to_dense(x1)
                x2 = sparse_to_dense(x2)
            elif not hasattr(x2, 'shape') or len(x2.shape) == 0:
                return sparse_with_values(x1, func(x1.values, x2))
            else:
                x2_zeros_and_nans = tf.equal(x2, 0)
                if not tf.as_dtype(x2.dtype).is_integer:
                    x2_zeros_and_nans = tf.math.logical_or(x2_zeros_and_nans, tf.math.is_nan(x2))

                def func_for_x1_indices():
                    return sparse_with_values(x1, func(x1.values, tf.gather_nd(x2, x1.indices)))

                def func_for_union_indices():
                    x2_zeros_and_nan_indices = tf.where(x2_zeros_and_nans)
                    union_indices, x1_values_for_union, _ = sparse_union_indices_and_values(x1, x2_zeros_and_nan_indices)
                    output = tf.SparseTensor(union_indices, func(x1_values_for_union, tf.gather_nd(x2, union_indices)), x1.dense_shape)
                    output.set_shape(x1.shape)
                    return output
                return tf.cond(tf.reduce_any(x2_zeros_and_nans), func_for_union_indices, func_for_x1_indices)
        elif isinstance(x2, tf.SparseTensor):
            x2 = sparse_to_dense(x2)
        elif isinstance(x1, tf.IndexedSlices):
            if isinstance(x2, tf.IndexedSlices):
                x1 = tf.convert_to_tensor(x1)
                x2 = tf.convert_to_tensor(x2)
            elif not hasattr(x2, 'shape') or len(x2.shape) == 0:
                return tf.IndexedSlices(func(x1.values, x2), x1.indices, x1.dense_shape)
            else:
                x2_zeros_and_nans = tf.equal(x2, 0)
                if not tf.as_dtype(x2.dtype).is_integer:
                    x2_zeros_and_nans = tf.math.logical_or(x2_zeros_and_nans, tf.math.is_nan(x2))
                x2_zeros_and_nans = tf.reduce_any(x2_zeros_and_nans, axis=tuple(range(1, x2.shape.rank)))

                def func_for_x1_indices():
                    return tf.IndexedSlices(func(x1.values, tf.gather(x2, x1.indices)), x1.indices, x1.dense_shape)

                def func_for_union_indices():
                    x2_zeros_and_nan_indices = tf.squeeze(tf.where(x2_zeros_and_nans), axis=-1)
                    union_indices, x1_values_for_union, _ = indexed_slices_union_indices_and_values(x1, x2_zeros_and_nan_indices)
                    return tf.IndexedSlices(func(x1_values_for_union, tf.gather(x2, union_indices)), union_indices, x1.dense_shape)
                return tf.cond(tf.reduce_any(x2_zeros_and_nans), func_for_union_indices, func_for_x1_indices)
        elif isinstance(x2, tf.IndexedSlices):
            x2 = tf.convert_to_tensor(x2)
        return func(x1, x2)
    return sparse_wrapper

Decorator to add support for `tf.SparseTensor` and `tf.IndexedSlices` to
    element-wise binary division and related operators.

    This decorator is designed for operations related to the division of two
    operands (e.g. `divide`). It accepts `tf.SparseTensor` and
    `tf.IndexedSlices` for both the dividend and the divisor, but handles them
    differently based on whether they are the dividend or the divisor.

    - If the divisor is a `tf.SparseTensor` or `tf.IndexedSlices`, it is
        densified and the result is dense because the result contains Inf or Nan
        outside of the indices of the dividend.
    - If the dividend is a `tf.SparseTensor` or `tf.IndexedSlices` and the
        divisor is dense, it finds occurrences of zeros and NaNs in the divisor.
        The result may therefore have more indices than there were in the dividend
        to return correct values where the divisor was zero or NaN.
    - If the dividend is a `tf.SparseTensor` or `tf.IndexedSlices` and the
        divisor is a scalar, it does the division element-wise. Note that the
        result is incorrectly sparse if the scalar divisor is zero.

Args:
        func: The function to wrap.

Returns:
        Wrapped function that supports `tf.SparseTensor` and `tf.IndexedSlices`.


def _valid_dtypes(self):
    return _DEFAULT_VALID_DTYPES

Valid types for loss, variables and gradients.

    Subclasses should override to allow other float types.

Returns:
        Valid types for loss, variables and gradients.


def _scalar(tf_fn, x, promote_to_float=False):
    x = np_array_ops.asarray(x)
    if promote_to_float and (not np.issubdtype(x.dtype.as_numpy_dtype, np.inexact)):
        x = x.astype(np_utils.result_type(float))
    return tf_fn(x)

Computes the tf_fn(x) for each element in `x`.

Args:
        tf_fn: function that takes a single Tensor argument.
        x: array_like. Could be an ndarray, a Tensor or any object that can be
        converted to a Tensor using `ops.convert_to_tensor`.
        promote_to_float: whether to cast the argument to a float dtype if it is not
        already.

Returns:
        An ndarray with the same shape as `x`. The default output dtype is
        determined by `np_utils.result_type(float)`, unless x is an ndarray with a
        floating point type, in which case the output type is same as x.dtype.


def export(self, input_ids: Optional[torch.Tensor]=None, cache_position: Optional[torch.Tensor]=None, dynamic_shapes: Optional[dict]=None, strict: Optional[bool]=None) -> torch.export.ExportedProgram:
    ALL_MASK_ATTENTION_FUNCTIONS.register('sdpa_without_vmap', sdpa_mask_without_vmap)
    ALL_ATTENTION_FUNCTIONS.register('sdpa_without_vmap', ALL_ATTENTION_FUNCTIONS['sdpa'])
    self.model.model.config._attn_implementation = 'sdpa_without_vmap'
    example_input_ids = input_ids if input_ids is not None else torch.tensor([[1]], dtype=torch.long)
    example_cache_position = cache_position if cache_position is not None else torch.tensor([0], dtype=torch.long)
    exported_program = torch.export.export(self.model, args=(example_input_ids, example_cache_position), kwargs={}, dynamic_shapes=dynamic_shapes, strict=strict if strict is not None else True)
    return exported_program

Export the wrapped module using `torch.export`.

Args:
        input_ids (`Optional[torch.Tensor]`):
            Tensor representing current input token id to the module. If not provided, a default tensor will be used.
        cache_position (`Optional[torch.Tensor]`):
            Tensor representing current input position in the cache. If not provided, a default tensor will be used.
        dynamic_shapes (`Optional[dict]`):
            Dynamic shapes to use for export if specified.
        strict(`Optional[bool]`):
            Flag to instruct `torch.export` to use `torchdynamo`.


def __init__(self, aggregation=None, metric_name=None, **kwargs):
    super(AddMetric, self).__init__(**kwargs)
    self.aggregation = aggregation
    self.metric_name = metric_name

Adds its inputs as a metric.

Attributes:
        aggregation: 'mean' or None. How the inputs should be aggregated.
        metric_name: The name to use for this metric.


def to_proto(self, export_scope=None):
    if context.executing_eagerly():
        raise RuntimeError('This operation is not supported when eager execution is enabled.')
    if export_scope is None or self.handle.name.startswith(export_scope):
        var_def = variable_pb2.VariableDef()
        var_def.variable_name = ops.strip_name_scope(self.handle.name, export_scope)
        if self._initial_value is not None:
            var_def.initial_value_name = ops.strip_name_scope(self._initial_value.name, export_scope)
        var_def.initializer_name = ops.strip_name_scope(self.initializer.name, export_scope)
        if self._cached_value is not None:
            var_def.snapshot_name = ops.strip_name_scope(self._cached_value.name, export_scope)
        else:
            var_def.snapshot_name = ops.strip_name_scope(self._graph_element.name, export_scope)
        var_def.is_resource = True
        var_def.trainable = self.trainable
        var_def.synchronization = self.synchronization.value
        var_def.aggregation = self.aggregation.value
        if self._save_slice_info:
            var_def.save_slice_info_def.MergeFrom(self._save_slice_info.to_proto(export_scope=export_scope))
        return var_def
    else:
        return None

Converts a `ResourceVariable` to a `VariableDef` protocol buffer.

Args:
        export_scope: Optional `string`. Name scope to remove.

Raises:
        RuntimeError: If run in EAGER mode.

Returns:
        A `VariableDef` protocol buffer, or `None` if the `Variable` is not
        in the specified name scope.


def _dump_file_name_to_datum(self, dir_name, file_name):
    debug_dump_rel_path = os.path.join(os.path.relpath(dir_name, self._dump_root), file_name)
    return DebugTensorDatum(self._dump_root, debug_dump_rel_path)

Obtain a DebugTensorDatum from the directory and file name.

Args:
        dir_name: (`str`) Name of the directory in which the dump file resides.
        file_name: (`str`) Base name of the dump file.

Returns:
        (`DebugTensorDatum`) The `DebugTensorDatum` loaded from the dump file.


def assert_compatible_matrix_dimensions(operator, x):
    assert_same_dd = check_ops.assert_equal(array_ops.shape(x)[-2], operator.domain_dimension_tensor(), message='Dimensions are not compatible.  shape[-2] of argument to be the same as this operator')
    return assert_same_dd

Assert that an argument to solve/matmul has proper domain dimension.

    If `operator.shape[-2:] = [M, N]`, and `x.shape[-2:] = [Q, R]`, then
    `operator.matmul(x)` is defined only if `N = Q`.  This `Op` returns an
    `Assert` that "fires" if this is not the case.  Static checks are already
    done by the base class `LinearOperator`.

Args:
        operator:  `LinearOperator`.
        x:  `Tensor`.

Returns:
        `Assert` `Op`.


def __getattr__(self, name: str):
    if name.startswith('__'):
        raise AttributeError(name)
    attr = getattr(self._builder, name)
    if isinstance(attr, expressions.Builder) and self._sealed:
        raise self._fhir_path_sealed_error(name)
    return ColumnExpressionBuilder._wrap_any(self, attr)

Redirects to the expressions.Builder when the attribute is not here.

    Note that in Python, '__getattribute__' always gets called first (the
    highest priority). Thus for attributes which has already been defined in
    this class, they won't be redirected to the expressions.Builder.

Args:
        name: The attribute name as a string.

Returns:
        The attribute get from expressions.Builder wrapped with _wrap_any.

Raises:
        AttributeError: if the FHIR path in this class is already sealed, or if
        getting the attribute from self._builder fails.


def __init__(self, agg_func: Callable[[Iterable[float]], float], agg_model_id: Optional[str]=None, include_source_predictions: bool=False):
    self._agg = agg_func
    _AggModelIdMixin.__init__(self, agg_model_id)
    _SourcePredictionMixin.__init__(self, include_source_predictions)

Aggregates anomaly predictions based on their scores.

    This is an abstract base class for `AggregationFn`s that combine multiple
    `AnomalyPrediction` objects into a single `AnomalyPrediction` based on
    the scores of the input predictions.

Args:
        agg_func (Callable[[Iterable[float]], float]): A function that aggregates
        a collection of anomaly scores (floats) into a single score.
        agg_model_id (Optional[str]): The model id used in aggregated predictions.
        Defaults to None.
        include_source_predictions (bool): If True, include the input predictions in
        the `source_predictions` of the output. Defaults to False.


def empty(shape, dtype=None):
    return backend.numpy.empty(shape, dtype=dtype)

Return a tensor of given shape and type filled with uninitialized data.

Args:
        shape: Shape of the empty tensor.
        dtype: Desired data type of the empty tensor.

Returns:
        The empty tensor.


def make_vjp(f, params=None, persistent=True):

    def decorated(*args, **kwds):
        """Computes the value and gradient of the decorated function."""
        parameter_positions = _get_arg_spec(f, params, args)
        assert not kwds, "The gradient function can't take keyword arguments."
        this_tape = tape.push_new_tape(persistent=persistent)
        try:
            sources = []
            args = [ops.convert_to_tensor(arg) if i in parameter_positions else arg for i, arg in enumerate(args)]
            args = _ensure_unique_tensor_objects(parameter_positions, args)
            for i in parameter_positions:
                if getattr(args[i], 'is_packed', False):
                    raise ValueError('GradientTape.gradient is not supported on packed EagerTensorsyet.')
                sources.append(args[i])
                tape.watch(this_tape, args[i])
            result = f(*args)
            if result is None:
                raise ValueError('Cannot differentiate a function that returns None; did you forget to return a value from {}?'.format(f.__name__))
            flat_result = nest.flatten(result)
            flat_result = [gen_array_ops.identity(x) for x in flat_result]
            result = nest.pack_sequence_as(result, flat_result)
        finally:
            tape.pop_tape(this_tape)

        def vjp(dy=None):
            if dy is not None:
                dy = [ops.convert_to_tensor(x) for x in nest.flatten(dy)]
            return imperative_grad.imperative_grad(this_tape, nest.flatten(result), sources, output_gradients=dy)
        return (result, vjp)
    return decorated

Returns a function that computes f and its vjp w.r.t.

    params.

    The term "vjp" here is an abbreviation for vector-jacobian product.

Args:
        f: the function to be differentiated.
        params: the parameters (numbers or names) to differentiate with respect to.
        A value of None will differentiate with respect to all parameters.
        persistent: Boolean controlling whether the VJP function can be re-used.
        Must be True or False.

Returns:
        A function, which when called, returns a tuple (value, vjp), where:
        - value is the result of calling f.
        - vjp is a function, which takes a vector as an argument and
        returns the product of that vector with the Jacobian of f.
        Providing no argument to vjp is equivalent to providing a
        vector of ones.

        For example,
        ```python
        def f(x):
        return x * x

        wrapped_fn = tfe.make_vjp(f)
        result, vjp = wrapped_fn(tf.constant(3.0))
        # result is 9.0
        vjp()  # the vjp function returns 6.0

Raises:
        ValueError: if `f` returns None.


def dense_block(x, blocks, name):
    for i in range(blocks):
        x = conv_block(x, 32, name=name + '_block' + str(i + 1))
    return x

A dense block.

Args:
        x: input tensor.
        blocks: integer, the number of building blocks.
        name: string, block label.

Returns:
        Output tensor for the block.


def restore(self, file_prefix: str, options: Optional[checkpoint_options.CheckpointOptions]=None) -> Dict[str, ops.Operation]:
    if options is not None and options.experimental_io_device is not None:
        raise ValueError('Specified experimental_io_device in DTensor checkpoint is not supported.')
    del options
    restore_specs = []
    tensor_structure = []
    for saveable in self._saveable_objects:
        saveable_tensor_structure = []
        tensor_structure.append(saveable_tensor_structure)
        for spec in saveable.specs:
            saveable_tensor_structure.append(spec.name)
            if isinstance(spec, d_variable.DSaveSpec):
                restore_specs.append((spec.name, spec.slice_spec, spec.dtype, spec.layout, spec.global_shape))
            elif isinstance(spec, saveable_object.SaveSpec):
                restore_specs.append((spec.name, spec.slice_spec, spec.dtype, layout.Layout.replicated(self._mesh.host_mesh(), spec.tensor.shape.rank).to_string(), spec.tensor.shape.as_list()))
    tensor_names, tensor_slices, tensor_dtypes, layouts, global_shapes = zip(*restore_specs)
    with ops.device(api.device_name()):
        restored_tensors = gen_dtensor_ops.d_tensor_restore_v2(prefix=file_prefix, tensor_names=tensor_names, shape_and_slices=tensor_slices, input_shapes=global_shapes, input_layouts=layouts, dtypes=tensor_dtypes)
    structured_restored_tensors = nest.pack_sequence_as(tensor_structure, restored_tensors)
    restore_ops = {}
    for saveable, restored_tensors in zip(self._saveable_objects, structured_restored_tensors):
        restore_ops[saveable.name] = saveable.restore(restored_tensors, restored_shapes=None)
    return restore_ops

Restore the saveable objects from a checkpoint with `file_prefix`.

Args:
        file_prefix: A string or scalar string Tensor containing the prefix for
        files to read from.
        options: Optional `CheckpointOptions` object. This is unused in DTensor.

Returns:
        A dictionary mapping from SaveableObject names to restore operations.


def insert_or_update(table, columns, values):
    rows = len(values)
    cells = len(columns) * len(values)
    return _Mutator(mutation=Mutation(insert_or_update=batch._make_write_pb(table, columns, values)), operation=WriteMutation._OPERATION_INSERT_OR_UPDATE, rows=rows, cells=cells, kwargs={'table': table, 'columns': columns, 'values': values})

Insert/update one or more table rows.

Args:
        table: Name of the table to be modified.
        columns: Name of the table columns to be modified.
        values: Values to be modified.