def _joined_array(num_dims, reduce_dim):
    formatter = '{:0%db}' % (num_dims - 1)
    result = np.zeros(shape=[2] * (num_dims - 1), dtype='S%d' % (2 * num_dims))
    flat = result.ravel()
    for i in range(2 ** (num_dims - 1)):
        dims = formatter.format(i)
        flat[i] = ''.join([(dims[:reduce_dim] + '%d' + dims[reduce_dim:]) % j for j in range(2)])
    return result

Creates an ndarray with the result from reduce_join on input_array.

Args:
        num_dims: The number of dimensions of the original input array.
        reduce_dim: The dimension to reduce.

Returns:
        An ndarray of shape [2] * (num_dims - 1).


def export(self, name=None):
    with tf.name_scope(name or '%s_lookup_table_export' % self._name):
        keys, values = gen_simple_hash_table_op.examples_simple_hash_table_export(self.resource_handle, key_dtype=self._key_dtype, value_dtype=self._value_dtype)
        return (keys, values)

Export all `key` and `value` pairs.

Args:
        name: A name for the operation (optional).

Returns:
        A tuple of two tensors, the first with the `keys` and the second with
        the `values`.


def assert_splits_match(nested_splits_lists):
    error_msg = 'Inputs must have identical ragged splits'
    for splits_list in nested_splits_lists:
        if len(splits_list) != len(nested_splits_lists[0]):
            raise ValueError(error_msg)
    return [check_ops.assert_equal(s1, s2, message=error_msg) for splits_list in nested_splits_lists[1:] for s1, s2 in zip(nested_splits_lists[0], splits_list)]

Checks that the given splits lists are identical.

    Performs static tests to ensure that the given splits lists are identical,
    and returns a list of control dependency op tensors that check that they are
    fully identical.

Args:
        nested_splits_lists: A list of nested_splits_lists, where each split_list is
        a list of `splits` tensors from a `RaggedTensor`, ordered from outermost
        ragged dimension to innermost ragged dimension.

Returns:
        A list of control dependency op tensors.

Raises:
        ValueError: If the splits are not identical.


def call(self, hidden_states: tf.Tensor, attention_mask: Optional[tf.Tensor]=None, position_ids: Optional[tf.Tensor]=None, past_key_value: Optional[Tuple[tf.Tensor]]=None, output_attentions: Optional[bool]=False, use_cache: Optional[bool]=False, training=False) -> Tuple[tf.Tensor, Optional[Tuple[tf.Tensor, tf.Tensor]]]:
    residual = hidden_states
    hidden_states = self.input_layernorm(hidden_states)
    hidden_states, self_attn_weights, present_key_value = self.self_attn(hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, output_attentions=output_attentions, use_cache=use_cache)
    hidden_states = tf.nn.dropout(hidden_states, rate=self.dropout)
    hidden_states = residual + hidden_states
    residual = hidden_states
    hidden_states = self.post_attention_layernorm(hidden_states)
    hidden_states = self.mlp(hidden_states)
    hidden_states = tf.nn.dropout(hidden_states, rate=self.dropout)
    hidden_states = residual + hidden_states
    outputs = (hidden_states,)
    if output_attentions:
        outputs += (self_attn_weights,)
    if use_cache:
        outputs += (present_key_value,)
    return outputs

Args:
    hidden_states (`tf.Tensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
    attention_mask (`tf.Tensor`, *optional*): attention mask of size
        `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
    output_attentions (`bool`, *optional*):
        Whether or not to return the attentions tensors of all attention layers. See `attentions` under
        returned tensors for more detail.
    use_cache (`bool`, *optional*):
        If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
        (see `past_key_values`).
    past_key_value (`Tuple(tf.Tensor)`, *optional*): cached past key and value projection states


def prediction_step(self, model: nn.Module, inputs: dict[str, Union[torch.Tensor, Any]], prediction_loss_only: bool, ignore_keys: Optional[list[str]]=None, **gen_kwargs) -> tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
    if not self.args.predict_with_generate or prediction_loss_only:
        return super().prediction_step(model, inputs, prediction_loss_only=prediction_loss_only, ignore_keys=ignore_keys)
    has_labels = 'labels' in inputs
    inputs = self._prepare_inputs(inputs)
    if len(gen_kwargs) == 0 and hasattr(self, '_gen_kwargs'):
        gen_kwargs = self._gen_kwargs.copy()
    if 'num_beams' in gen_kwargs and gen_kwargs['num_beams'] is None:
        gen_kwargs.pop('num_beams')
    if 'max_length' in gen_kwargs and gen_kwargs['max_length'] is None:
        gen_kwargs.pop('max_length')
    default_synced_gpus = is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self.model)
    gen_kwargs['synced_gpus'] = gen_kwargs.get('synced_gpus', default_synced_gpus)
    generation_inputs = inputs.copy()
    if 'labels' in generation_inputs and 'decoder_input_ids' in generation_inputs and (generation_inputs['labels'].shape == generation_inputs['decoder_input_ids'].shape):
        generation_inputs = {k: v for k, v in inputs.items() if k not in ('decoder_input_ids', 'decoder_attention_mask')}
    summon_full_params_context = FullyShardedDataParallel.summon_full_params(self.model) if isinstance(self.model, FullyShardedDataParallel) else contextlib.nullcontext()
    with summon_full_params_context:
        generated_tokens = self.model.generate(**generation_inputs, **gen_kwargs)
    if self.model.generation_config._from_model_config:
        self.model.generation_config._from_model_config = False
    gen_config = self.model.generation_config
    if generated_tokens.shape[-1] < gen_config.max_length:
        generated_tokens = self._pad_tensors_to_max_len(generated_tokens, gen_config.max_length)
    elif gen_config.max_new_tokens is not None and generated_tokens.shape[-1] < gen_config.max_new_tokens + 1:
        generated_tokens = self._pad_tensors_to_max_len(generated_tokens, gen_config.max_new_tokens + 1)
    with torch.no_grad():
        if has_labels:
            with self.compute_loss_context_manager():
                outputs = model(**inputs)
            if self.label_smoother is not None:
                loss = self.label_smoother(outputs, inputs['labels']).detach().mean()
            else:
                loss = (outputs['loss'] if isinstance(outputs, dict) else outputs[0]).detach().mean()
        else:
            loss = None
    if self.args.prediction_loss_only:
        return (loss, None, None)
    if has_labels:
        labels = inputs['labels']
        if labels.shape[-1] < gen_config.max_length:
            labels = self._pad_tensors_to_max_len(labels, gen_config.max_length)
        elif gen_config.max_new_tokens is not None and labels.shape[-1] < gen_config.max_new_tokens + 1:
            labels = self._pad_tensors_to_max_len(labels, gen_config.max_new_tokens + 1)
    else:
        labels = None
    return (loss, generated_tokens, labels)

Perform an evaluation step on `model` using `inputs`.

    Subclass and override to inject custom behavior.

Args:
        model (`nn.Module`):
            The model to evaluate.
        inputs (`Dict[str, Union[torch.Tensor, Any]]`):
            The inputs and targets of the model.

            The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
            argument `labels`. Check your model's documentation for all accepted arguments.
        prediction_loss_only (`bool`):
            Whether or not to return the loss only.
        gen_kwargs:
            Additional `generate` specific kwargs.

Returns:
        Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]: A tuple with the loss, logits and
        labels (each being optional).


def _delete_minibatch(self, bucket, keys):
    request = messages.DeleteBatchRequest(bucket, keys)
    results = {}
    try:
        response = self.client.delete_batch(request)
        for key in response.deleted:
            results[bucket, key] = None
        for key, error in zip(response.failed, response.errors):
            results[bucket, key] = error
    except messages.S3ClientError as e:
        for key in keys:
            results[bucket, key] = e
    return results

A helper method. Boto3 allows batch deletions
    for files within the same bucket.

Args:
        bucket: String bucket name
        keys: List of keys to be deleted in the bucket

    Returns: dict of the form {(bucket, key): error}, where error is None if the
    operation succeeded


def elementwise_binary_union(linear, use_sparsify):

    def wrap_elementwise_binary_union(func):
        sparse_func = jax_sparse.sparsify(func) if use_sparsify else None

        @functools.wraps(func)
        def sparse_wrapper(x1, x2):
            if isinstance(x1, jax_sparse.JAXSparse):
                if isinstance(x2, jax_sparse.JAXSparse):
                    if x1.indices is x2.indices and isinstance(x1, jax_sparse.BCOO) and isinstance(x2, jax_sparse.BCOO):
                        if not linear and (not x1.unique_indices):
                            x1 = jax_sparse.bcoo_sum_duplicates(x1)
                            x2 = jax_sparse.bcoo_sum_duplicates(x2)
                        return jax_sparse.BCOO((func(x1.data, x2.data), x1.indices), shape=x1.shape, indices_sorted=x1.indices_sorted, unique_indices=x1.unique_indices)
                    elif use_sparsify:
                        return sparse_func(x1, x2)
                    elif isinstance(x1, jax_sparse.BCOO) and isinstance(x2, jax_sparse.BCOO):
                        x1 = bcoo_add_indices(x1, x2, sum_duplicates=not linear)
                        x2 = bcoo_add_indices(x2, x1, sum_duplicates=not linear)
                        return jax_sparse.BCOO((func(x1.data, x2.data), x1.indices), shape=x1.shape, indices_sorted=True, unique_indices=True)
                    else:
                        ValueError(f'Unsupported sparse format: {x1.__class__} and {x2.__class__}')
                else:
                    x1 = x1.todense()
            elif isinstance(x2, jax_sparse.JAXSparse):
                x2 = x2.todense()
            return func(x1, x2)
        return sparse_wrapper
    return wrap_elementwise_binary_union

Decorator to add support for `JAXSparse` tensors (including `BCOO`) to an
    element-wise binary operator such that the indices present in the result are
    are the union of the indices in the two operand.

    The primary use case for this is the `add` and `subtract` operators.

    There are requirements on the operator for this decorator to work correctly:

    - The operator must be element-wise.
    - The operator must be binary (two input tensors and one output tensor).
    - Both inputs must be of the same shape or one input must be a scalar.
    - The output must be of the same shape as the (non scalar) inputs.
    - The indices of the output must be the union of the indices of the inputs.
        This implies that func(0, 0) must be 0. As a result, if one operand is
        dense or a scalar, then the result will be dense.

    Additional arguments to the function (besides the input tensors) are not
    supported.

    Note that if the result of the operation is zero at some indices, including
    because the operands were zero at these indices, the zeros and indices are
    preserved.

    The `BCOO` format is the only supported one in all cases. Other formats are
    not supported when `use_sparsify` is `False`.

Args:
        use_sparsify: indicates that the JAX `sparsify` transform supports this
            operation.
        linear: if `True`, mean that the operation is such that
            `op(a + b, c) == op(a, c) + op(b, c)` and
            `op(a, c + d) == op(a, c) + op(a, d)`.

Returns:
        Wrapped function that supports `JAXSparse`.


def test_binomial_vs_alternative_formulation(y_true, y_pred, global_dtype):

    def alt_loss(y, raw_pred):
        z = 2 * y - 1
        return np.mean(np.log(1 + np.exp(-z * raw_pred)))

    def alt_gradient(y, raw_pred):
        z = 2 * y - 1
        return -z / (1 + np.exp(z * raw_pred))
    bin_loss = HalfBinomialLoss()
    y_true = y_true.astype(global_dtype)
    y_pred = y_pred.astype(global_dtype)
    datum = (y_true, y_pred)
    assert bin_loss(*datum) == approx(alt_loss(*datum))
    assert_allclose(bin_loss.gradient(*datum), alt_gradient(*datum))

Test that both formulations of the binomial deviance agree.

    Often, the binomial deviance or log loss is written in terms of a variable
    z in {-1, +1}, but we use y in {0, 1}, hence z = 2 * y - 1.
    ESL II Eq. (10.18):

        -loglike(z, f) = log(1 + exp(-2 * z * f))

Note:
        - ESL 2*f = raw_prediction, hence the factor 2 of ESL disappears.
        - Deviance = -2*loglike + .., but HalfBinomialLoss is half of the
        deviance, hence the factor of 2 cancels in the comparison.


def input(self):
    return self._nested_inputs

Retrieves the input tensor(s) of a layer.

    Only applicable if the layer has exactly one input,
    i.e. if it is connected to one incoming layer.

Returns:
        Input tensor or list of input tensors.

Raises:
        RuntimeError: If called in Eager mode.
        AttributeError: If no inbound nodes are found.


def normalize(x, axis=-1, order=2):
    l2 = np.atleast_1d(np.linalg.norm(x, order, axis))
    l2[l2 == 0] = 1
    return x / np.expand_dims(l2, axis)

Normalizes a Numpy array.

Args:
        x: Numpy array to normalize.
        axis: axis along which to normalize.
        order: Normalization order (e.g. `order=2` for L2 norm).

Returns:
        A normalized copy of the array.


def _get_or_create_arg_by_name(state, name, is_kwarg=False):
    for arg in state.args + state.kwargs:
        if arg.name == name:
            return arg
    arg = Namespace()
    arg.name = name
    arg.type.lines = []
    arg.description.lines = []
    if is_kwarg:
        state.kwargs.append(arg)
    else:
        state.args.append(arg)
    return arg

Gets or creates a new Arg.

    These Arg objects (Namespaces) are turned into the ArgInfo namedtuples
    returned by parse. Each Arg object is used to collect the name, type, and
    description of a single argument to the docstring's function.

Args:
        state: The state of the parser.
        name: The name of the arg to create.
        is_kwarg: A boolean representing whether the argument is a keyword arg.

Returns:
        The new Arg.


def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub: bool=False, **kwargs):
    use_auth_token = kwargs.pop('use_auth_token', None)
    if use_auth_token is not None:
        warnings.warn('The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.', FutureWarning)
        if kwargs.get('token', None) is not None:
            raise ValueError('`token` and `use_auth_token` are both specified. Please set only the argument `token`.')
        kwargs['token'] = use_auth_token
    if os.path.isfile(save_directory):
        raise AssertionError(f'Provided path ({save_directory}) should be a directory, not a file')
    os.makedirs(save_directory, exist_ok=True)
    if push_to_hub:
        commit_message = kwargs.pop('commit_message', None)
        repo_id = kwargs.pop('repo_id', save_directory.split(os.path.sep)[-1])
        repo_id = self._create_repo(repo_id, **kwargs)
        files_timestamps = self._get_files_timestamps(save_directory)
    if self._auto_class is not None:
        custom_object_save(self, save_directory, config=self)
    output_image_processor_file = os.path.join(save_directory, IMAGE_PROCESSOR_NAME)
    self.to_json_file(output_image_processor_file)
    logger.info(f'Image processor saved in {output_image_processor_file}')
    if push_to_hub:
        self._upload_modified_files(save_directory, repo_id, files_timestamps, commit_message=commit_message, token=kwargs.get('token'))
    return [output_image_processor_file]

Save an image processor object to the directory `save_directory`, so that it can be re-loaded using the
    [`~image_processing_utils.ImageProcessingMixin.from_pretrained`] class method.

Args:
        save_directory (`str` or `os.PathLike`):
            Directory where the image processor JSON file will be saved (will be created if it does not exist).
        push_to_hub (`bool`, *optional*, defaults to `False`):
            Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the
            repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
            namespace).
        kwargs (`Dict[str, Any]`, *optional*):
            Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.


def GetOpeningBracket(node):
    return getattr(node, _NODE_ANNOTATION_PREFIX + 'container_bracket', None)

Get opening bracket value from a node.

Args:
        node: the node.

Returns:
        The opening bracket node or None if it couldn't find one.


def from_features(cls, features, types):
    params = cls()
    if features:
        for key in sorted(features.keys()):
            feature = features[key]
            if not isinstance(feature, tuple(types)):
                raise ValueError(f"Unsupported {type(feature).__name__} {feature} for key '{key}'")
            params._add_feature(key, feature)
    params._validate()
    return params

Builds _ParseOpParams for a given set of features and allowed types.

Args:
        features: A `dict` mapping feature keys to objects of a type in `types`.
        types: Type of features to allow, among `FixedLenFeature`,
        `VarLenFeature`, `SparseFeature`, and `FixedLenSequenceFeature`.

Returns:
        A `_ParseOpParams` containing the raw parameters for `gen_parsing_ops`.

Raises:
        ValueError: if `features` contains an item not in `types`, or an invalid
        feature.
        ValueError: if sparse and dense key sets intersect.
        ValueError: if input lengths do not match up.


def random_shuffle(value, seed=None, name=None):
    with ops.name_scope(name, 'shuffle', [value, seed]):
        if value.rank == 0:
            raise ValueError('Cannot shuffle a scalar StructuredTensor')
        first_dimension = value.nrows()
        index = random_ops.random_shuffle(math_ops.range(first_dimension), seed=seed)
        return gather(value, index, axis=0)

Shuffle a structured tensor on the zeroth axis.

Args:
        value: a structured tensor of rank at least one.
        seed: the seed for shuffling.
        name: the name for shuffle.

Returns:
        The shuffled structured tensor.


def forward(self, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.FloatTensor]=None, token_type_ids: Optional[torch.LongTensor]=None, position_ids: Optional[torch.LongTensor]=None, head_mask: Optional[torch.FloatTensor]=None, inputs_embeds: Optional[torch.FloatTensor]=None, encoder_hidden_states: Optional[torch.FloatTensor]=None, encoder_attention_mask: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, **kwargs) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    if labels is not None:
        use_cache = False
    outputs = self.data2vec_text(input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, past_key_values=past_key_values, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    sequence_output = outputs[0]
    prediction_scores = self.lm_head(sequence_output)
    lm_loss = None
    if labels is not None:
        lm_loss = self.loss_function(prediction_scores, labels, vocab_size=self.config.vocab_size, **kwargs)
    if not return_dict:
        output = (prediction_scores,) + outputs[2:]
        return (lm_loss,) + output if lm_loss is not None else output
    return CausalLMOutputWithCrossAttentions(loss=lm_loss, logits=prediction_scores, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, cross_attentions=outputs.cross_attentions)

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

Example:
    ```python
    >>> from transformers import AutoTokenizer, Data2VecTextForCausalLM, Data2VecTextConfig
    >>> import torch

    >>> tokenizer = AutoTokenizer.from_pretrained("facebook/data2vec-text-base")
    >>> config = Data2VecTextConfig.from_pretrained("facebook/data2vec-text-base")
    >>> config.is_decoder = True
    >>> model = Data2VecTextForCausalLM.from_pretrained("facebook/data2vec-text-base", config=config)

    >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
    >>> outputs = model(**inputs)

    >>> prediction_logits = outputs.logits
    ```


def class_from_typename(cls, type_name: str) -> Optional[Type['JSONConvertible']]:
    return cls._TYPE_REGISTRY.class_from_typename(type_name)

Gets the class for a registered type name.

Args:
        type_name: A string as the global unique type identifier for requested
        class.

Returns:
        A type object if registered, otherwise None.


def random_sample(value: Any, num_examples: Optional[int]=None, where: Optional[Callable[[base.HyperPrimitive], bool]]=None, seed: Optional[int]=None):
    return iterate(value, num_examples, geno.Random(seed), where=where)

Returns an iterator of random sampled examples.

    Example::

        hyper_dict = pg.Dict(x=pg.oneof(range(3)), y=pg.floatv(0.0, 1.0))

        # Generate one random example from the hyper_dict.
        d = next(pg.random_sample(hyper_dict))

        # Generate 5 random examples with random seed.
        ds = list(pg.random_sample(hyper_dict, 5, seed=1))

        # Generate 3 random examples of `x` with `y` intact.
        ds = list(pg.random_sample(hyper_dict, 3,
        where=lambda x: isinstance(x, pg.hyper.OneOf)))

Args:
        value: A (maybe) hyper value.
        num_examples: An optional integer as number of examples to propose. If None,
        propose will return an iterator that iterates forever.
        where: Function to filter hyper primitives. If None, all hyper primitives in
        `value` will be included in the encoding/decoding process. Otherwise only
        the hyper primitives on which 'where' returns True will be included.
        `where` can be useful to partition a search space into separate
        optimization processes. Please see 'Template' docstr for details.
        seed: An optional integer as random seed.

Returns:
        Iterator of random examples.


def GetTermSize(self):
    return self._term_size

Returns the terminal (x, y) dimensions in characters.

Returns:
        (x, y): A tuple of the terminal x and y dimensions.


def block_embedding_to(self, device):
    self.block_emb = self.block_emb.to(device)

Send `self.block_emb` to a specific device.

Args:
        device (`str` or `torch.device`):
            The device to which `self.block_emb` will be sent.


def random_hermitian_matrix(num_qubits):
    dim = 2 ** num_qubits
    val_range = 2
    random_real = tf.cast(tf.random.uniform([dim, dim], -val_range, val_range), tf.complex128)
    random_imag = 1j * tf.cast(tf.random.uniform([dim, dim], -val_range, val_range), tf.complex128)
    random_matrix = random_real + random_imag
    return random_matrix + tf.linalg.adjoint(random_matrix)

Returns a random Hermitian matrix.

    Uses the property that A + A* is Hermitian for any matrix A.

Args:
        num_qubits: Number of qubits on which the matrix acts.


def _reset_build_compile_trackers(model):
    model.built = False
    model.inputs = None
    model.outputs = None
    model._is_compiled = False
    if not ops.executing_eagerly_outside_functions():
        model._v1_compile_was_called = False
    model.optimizer = None

Reset state trackers for model.

    Note that we do not actually zero out attributes such as optimizer,
    but instead rely on the expectation that all of the attrs will be
    over-written on calling build/compile/etc. This is somewhat fragile,
    insofar as we check elsewhere for the presence of these attributes as
    evidence of having been built/compiled/etc. Pending a better way to do this,
    we reset key attributes here to allow building and compiling.

Args:
        model: the model that is being reset


def global_step(sess, global_step_tensor):
    if context.executing_eagerly():
        return int(global_step_tensor.numpy())
    return int(sess.run(global_step_tensor))

Small helper to get the global step.

    ```python
    # Create a variable to hold the global_step.
    global_step_tensor = tf.Variable(10, trainable=False, name='global_step')
    # Create a session.
    sess = tf.compat.v1.Session()
    # Initialize the variable
    sess.run(global_step_tensor.initializer)
    # Get the variable value.
    print('global_step: %s' % tf.compat.v1.train.global_step(sess,
    global_step_tensor))

    global_step: 10
    ```

Args:
        sess: A TensorFlow `Session` object.
        global_step_tensor:  `Tensor` or the `name` of the operation that contains
        the global step.

Returns:
        The global step value.


def log(self: EventSetOrNode) -> EventSetOrNode:
    from temporian.core.operators.unary import log
    return log(self)

Calculates the natural logarithm of an [`EventSet`][temporian.EventSet]'s
    features.

    Can only be used on floating point features.

Example:
        ```python
        >>> a = tp.event_set(
        ...     timestamps=[1, 2, 3, 4, 5],
        ...     features={"M": [np.e, 1., 2., 10., -1.]},
        ... )
        >>> a.log()
        indexes: ...
                timestamps: [1. 2. 3. 4. 5.]
                'M': [1. 0. 0.6931 2.3026 nan]
        ...

        ```

Returns:
        EventSetOr with logarithm of input features.


def get_num_bytes(self, batch: Sequence[numpy.ndarray]) -> int:
    return sum((np_array.itemsize for np_array in batch))

Returns:
    The number of bytes of data for a batch.


def remove_comp_items(self, context_word, comp_items):
    if context_word not in self._comp_dict:
        raise KeyError('Context word "%s" has not been registered' % context_word)
    for item in comp_items:
        self._comp_dict[context_word].remove(item)

Remove a list of completion items from a completion context.

Args:
        context_word: A single completion word as a string. The removal will
        also apply to all other context words of the same context.
        comp_items: Completion items to remove.

Raises:
        KeyError: if the context word has not been registered.


def __call__(self, *args: Union[str, 'Image.Image', List['Image.Image'], List[str]], **kwargs: Any) -> List[Any]:
    return super().__call__(*args, **kwargs)

Extract the features of the input(s).

Args:
        images (`str`, `List[str]`, `PIL.Image` or `List[PIL.Image]`):
            The pipeline handles three types of images:

            - A string containing a http link pointing to an image
            - A string containing a local path to an image
            - An image loaded in PIL directly

            The pipeline accepts either a single image or a batch of images, which must then be passed as a string.
            Images in a batch must all be in the same format: all as http links, all as local paths, or all as PIL
            images.
        timeout (`float`, *optional*, defaults to None):
            The maximum time in seconds to wait for fetching images from the web. If None, no timeout is used and
            the call may block forever.

Returns:
        A nested list of `float`: The features computed by the model.


def name_based_restore(mesh: layout_lib.Mesh, checkpoint_prefix: str, name_tensor_dict: Dict[str, Union[tensor_lib.Tensor, tf_variables.Variable]]):
    if not context.executing_eagerly():
        raise ValueError('name based restore must run eagerly.')
    ordered_name_tensor_dict = name_tensor_dict
    if not isinstance(name_tensor_dict, collections.OrderedDict):
        ordered_name_tensor_dict = collections.OrderedDict(name_tensor_dict)
    for name, tensor in ordered_name_tensor_dict.items():
        try:
            if api.fetch_layout(tensor).mesh.device_type().upper() != 'CPU':
                raise ValueError('Restoring a non CPU Tensor is not supported currently. Offending tensor name : {tensor_name}'.format(tensor_name=name))
        except errors_impl.OpError as op_error:
            raise ValueError('Saving/Restoring tensor must be a DTensor') from op_error
    checkpoint_prefix = api.pack([checkpoint_prefix] * mesh.num_local_devices(), layout_lib.Layout.replicated(mesh.host_mesh(), rank=0))
    tensor_names = api.pack([list(ordered_name_tensor_dict.keys())] * mesh.num_local_devices(), layout_lib.Layout.replicated(mesh.host_mesh(), rank=1))
    shape_and_slices = api.pack([[''] * len(ordered_name_tensor_dict)] * mesh.num_local_devices(), layout_lib.Layout.replicated(mesh.host_mesh(), rank=1))
    input_shapes = [tensor.shape for tensor in ordered_name_tensor_dict.values()]
    input_layouts = [api.fetch_layout(tensor).to_string() for tensor in ordered_name_tensor_dict.values()]
    with ops.device(api.device_name()):
        restored_cpu_tensors = gen_dtensor_ops.d_tensor_restore_v2(prefix=checkpoint_prefix, tensor_names=tensor_names, shape_and_slices=shape_and_slices, input_shapes=input_shapes, input_layouts=input_layouts, dtypes=[tensor.dtype for tensor in ordered_name_tensor_dict.values()])
    return collections.OrderedDict(zip(ordered_name_tensor_dict.keys(), restored_cpu_tensors))

Restores from checkpoint_prefix to name based DTensors.

    It is required to have already-initialized DTensor variables that have same
    shape/dtype for the tensors being restored.

    Also, we currently only support a named based restore on a single mesh.

Args:
        mesh: The single mesh that all Tensors would be restored to.
        checkpoint_prefix : The prefix of checkpoint to be restored.
        name_tensor_dict: A ordered dictionary of tensor_names to a DTensor. The
        DTensor shape/dtype must match the tensors being saved/restored for now.

Returns:
        A dictionary of name to its restored DTensor value.


def filter_tests(output_file: str, filters: List[str]):
    if not os.path.isfile(output_file):
        print('No test file found.')
        return
    with open(output_file, 'r', encoding='utf-8') as f:
        test_files = f.read().split(' ')
    if len(test_files) == 0 or test_files == ['']:
        print('No tests to filter.')
        return
    if test_files == ['tests']:
        test_files = [os.path.join('tests', f) for f in os.listdir('tests') if f not in ['__init__.py'] + filters]
    else:
        test_files = [f for f in test_files if f.split(os.path.sep)[1] not in filters]
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(' '.join(test_files))

Reads the content of the output file and filters out all the tests in a list of given folders.

Args:
        output_file (`str` or `os.PathLike`): The path to the output file of the tests fetcher.
        filters (`List[str]`): A list of folders to filter.


def call(self, input_ids: Optional[tf.Tensor]=None, position_ids: Optional[tf.Tensor]=None, inputs_embeds: Optional[tf.Tensor]=None) -> tf.Tensor:
    if input_ids is None and inputs_embeds is None:
        raise ValueError('You have to specify either input_ids or inputs_embeds')
    if inputs_embeds is None:
        check_embeddings_within_bounds(input_ids, self.config.vocab_size)
        inputs_embeds = tf.gather(params=self.weight, indices=input_ids)
    input_shape = shape_list(inputs_embeds)[:-1]
    if position_ids is None:
        position_ids = tf.expand_dims(tf.range(start=0, limit=input_shape[-1]), axis=0)
    position_embeds = tf.gather(params=self.position_embedding, indices=position_ids)
    position_embeds = tf.tile(input=position_embeds, multiples=(input_shape[0], 1, 1))
    final_embeddings = inputs_embeds + position_embeds
    return final_embeddings

Applies embedding based on inputs tensor.

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


def node_device(self, node_name):
    if not self._debug_graphs:
        raise LookupError('Node devices are not loaded from partition graphs yet.')
    if node_name not in self._node_devices:
        raise ValueError("Node '%s' does not exist in partition graphs." % node_name)
    output = list(self._node_devices[node_name])
    return output[0] if len(output) == 1 else output

Get the names of the devices that has nodes of the specified name.

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

Returns:
        (`str` or `list` of `str`) name of the device(s) on which the node of the
        given name is found. Returns a `str` if there is only one such device,
        otherwise return a `list` of `str`.

Raises:
        LookupError: If node inputs and control inputs have not been loaded
        from partition graphs yet.
        ValueError: If the node does not exist in partition graphs.


def _expected_exercise_fn(design, calibration_indices, continuation_value, exercise_value):
    mask = exercise_value > 0
    design_t = tf.transpose(design, [0, 2, 1])
    masked = tf.where(tf.expand_dims(tf.transpose(mask), axis=-1), design_t, tf.zeros_like(design_t))
    if calibration_indices is None:
        submask = masked
        mask_cont_value = continuation_value
    else:
        submask = tf.gather(masked, calibration_indices, axis=1)
        mask_cont_value = tf.gather(continuation_value, calibration_indices)
    lhs = tf.matmul(submask, submask, transpose_a=True)
    lhs_pinv = tf.linalg.pinv(lhs)
    rhs = tf.matmul(submask, tf.expand_dims(tf.transpose(mask_cont_value), axis=-1), transpose_a=True)
    beta = tf.matmul(lhs_pinv, rhs)
    continuation = tf.matmul(design_t, beta)
    return tf.nn.relu(tf.transpose(tf.squeeze(continuation, axis=-1)))

Returns the expected continuation value for each path.

Args:
        design: A real `Tensor` of shape `[batch_size, basis_size, num_samples]`.
        calibration_indices: A rank 1 integer `Tensor` denoting indices of samples
        used for regression.
        continuation_value: A `Tensor` of shape `[num_samples, batch_size]` and of
        the same dtype as `design`. The optimal value of the option conditional on
        not exercising now or earlier, taking future information into account.
        exercise_value: A `Tensor` of the same shape and dtype as
        `continuation_value`. Value of the option if exercised immideately at
        the current time

Returns:
        A `Tensor` of the same shape and dtype as `continuation_value` whose
        `(n, v)`-th entry represents the expected continuation value of sample path
        `n` under the `v`-th payoff scheme.


def __init__(self, config: Kosmos2VisionConfig):
    super().__init__()
    self.config = config
    self.layers = nn.ModuleList([Kosmos2VisionEncoderLayer(config) for _ in range(config.num_hidden_layers)])
    self.gradient_checkpointing = False

Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
    [`Kosmos2VisionEncoderLayer`].

Args:
        config: Kosmos2VisionConfig


def __init__(self, config):
    super().__init__()
    self.use_batch_norm = config.use_batch_norm_in_fusion_residual
    use_bias_in_fusion_residual = config.use_bias_in_fusion_residual if config.use_bias_in_fusion_residual is not None else not self.use_batch_norm
    self.activation1 = nn.ReLU()
    self.convolution1 = nn.Conv2d(config.fusion_hidden_size, config.fusion_hidden_size, kernel_size=3, stride=1, padding=1, bias=use_bias_in_fusion_residual)
    self.activation2 = nn.ReLU()
    self.convolution2 = nn.Conv2d(config.fusion_hidden_size, config.fusion_hidden_size, kernel_size=3, stride=1, padding=1, bias=use_bias_in_fusion_residual)
    if self.use_batch_norm:
        self.batch_norm1 = nn.BatchNorm2d(config.fusion_hidden_size)
        self.batch_norm2 = nn.BatchNorm2d(config.fusion_hidden_size)

ResidualConvUnit, pre-activate residual unit.

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


def is_subtype_of(self, other: trace.TraceType) -> bool:
    if type(self) is not type(other):
        return False
    is_subtype = True

    def check_attribute(attribute_self, attribute_other):
        nonlocal is_subtype
        if not is_subtype:
            return
        if isinstance(attribute_self, trace.TraceType):
            if not attribute_self.is_subtype_of(attribute_other):
                is_subtype = False
                return
        elif attribute_self != attribute_other:
            is_subtype = False
    try:
        nest.map_structure(check_attribute, self._serialize(), other._serialize())
    except (ValueError, TypeError):
        return False
    return is_subtype

Returns True if `self` is a subtype of `other`.

    Implements the tf.types.experimental.func.TraceType interface.

    If not overridden by a subclass, the default behavior is to assume the
    TypeSpec is covariant upon attributes that implement TraceType and
    invariant upon rest of the attributes as well as the structure and type
    of the TypeSpec.

Args:
        other: A TraceType object.


def _concat(self):
    if len(self._variable_list) == 1:
        with ops.name_scope(None):
            return array_ops.identity(self._variable_list[0], name=self._name)
    partition_axes = self._partition_axes()
    if len(partition_axes) > 1:
        raise NotImplementedError('Cannot concatenate along more than one dimension: %s.  Multi-axis partition concat is not supported' % str(partition_axes))
    partition_ix = partition_axes[0]
    with ops.name_scope(self._name + '/ConcatPartitions/'):
        concatenated = array_ops.concat(self._variable_list, partition_ix)
    with ops.name_scope(None):
        return array_ops.identity(concatenated, name=self._name)

Returns the overall concatenated value as a `Tensor`.

    This is different from using the partitioned variable directly as a tensor
    (through tensor conversion and `as_tensor`) in that it creates a new set of
    operations that keeps the control dependencies from its scope.

Returns:
        `Tensor` containing the concatenated value.


def create_test_methods(spec):
    for suffix, providers in provider_sets(spec):

        def test(self, providers=providers):
            vars = {}
            with contextlib.ExitStack() as stack:
                stack.enter_context(mock.patch('apache_beam.yaml.yaml_provider.standard_providers', lambda: providers))
                for fixture in spec.get('fixtures', []):
                    vars[fixture['name']] = stack.enter_context(python_callable.PythonCallableWithSource.load_from_fully_qualified_name(fixture['type'])(**yaml_transform.SafeLineLoader.strip_metadata(fixture.get('config', {}))))
                for pipeline_spec in spec['pipelines']:
                    with beam.Pipeline(options=PipelineOptions(pickle_library='cloudpickle', **yaml_transform.SafeLineLoader.strip_metadata(pipeline_spec.get('options', {})))) as p:
                        yaml_transform.expand_pipeline(p, replace_recursive(pipeline_spec, vars))
        yield (f'test_{suffix}', test)

Dynamically creates test methods based on a YAML specification.

    This function takes a YAML specification (`spec`) which defines pipelines,
    fixtures, and potentially options. It iterates through different
    combinations of "providers" (which determine how YAML transforms are
    implemented, e.g., using Python or SQL).

    For each combination of providers:
        1. It constructs a unique test method name (e.g., `test_only`).
        2. It defines a test method that:
        a. Sets up any specified fixtures, making their values available as
            variables.
        b. Mocks the standard YAML providers to use the current combination
            of providers for this test run.
        c. For each pipeline defined in the `spec`:
            i. Creates a `beam.Pipeline` instance with specified options.
            ii. Expands the YAML pipeline definition using
                `yaml_transform.expand_pipeline`, substituting any fixture
                variables.
            iii. Runs the Beam pipeline.

    The function yields tuples of (test_method_name, test_method_function),
    which can then be used to populate a `unittest.TestCase` class.

Args:
        spec (dict): A dictionary parsed from a YAML test specification file.
        It's expected to have keys like 'fixtures' (optional) and 'pipelines'.

Yields:
        tuple: A tuple containing:
        - str: The generated name for the test method (e.g., "test_only").
        - function: The dynamically generated test method.


def forward(self, inputs_embeds=None, attention_mask=None, object_queries=None, output_attentions=None, output_hidden_states=None, return_dict=None):
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    hidden_states = inputs_embeds
    hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
    if attention_mask is not None:
        attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype)
    encoder_states = () if output_hidden_states else None
    all_attentions = () if output_attentions else None
    for i, encoder_layer in enumerate(self.layers):
        if output_hidden_states:
            encoder_states = encoder_states + (hidden_states,)
        to_drop = False
        if self.training:
            dropout_probability = torch.rand([])
            if dropout_probability < self.layerdrop:
                to_drop = True
        if to_drop:
            layer_outputs = (None, None)
        else:
            layer_outputs = encoder_layer(hidden_states, attention_mask, object_queries=object_queries, output_attentions=output_attentions)
            hidden_states = layer_outputs[0]
        if output_attentions:
            all_attentions = all_attentions + (layer_outputs[1],)
    if output_hidden_states:
        encoder_states = encoder_states + (hidden_states,)
    if not return_dict:
        return tuple((v for v in [hidden_states, encoder_states, all_attentions] if v is not None))
    return BaseModelOutput(last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions)

Args:
    inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
        Flattened feature map (output of the backbone + projection layer) that is passed to the encoder.

    attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
        Mask to avoid performing attention on padding pixel features. Mask values selected in `[0, 1]`:

        - 1 for pixel features that are real (i.e. **not masked**),
        - 0 for pixel features that are padding (i.e. **masked**).

        [What are attention masks?](../glossary#attention-mask)

    object_queries (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
        Object queries that are added to the queries in each self-attention layer.

    output_attentions (`bool`, *optional*):
        Whether or not to return the attentions tensors of all attention layers. See `attentions` under
        returned tensors for more detail.
    output_hidden_states (`bool`, *optional*):
        Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
        for more detail.
    return_dict (`bool`, *optional*):
        Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.


def _prefix_to_checkpoint_path(prefix, format_version):
    if format_version == saver_pb2.SaverDef.V2:
        return prefix + '.index'
    return prefix

Returns the pathname of a checkpoint file, given the checkpoint prefix.

    For V1 checkpoint, simply returns the prefix itself (the data file).  For V2,
    returns the pathname to the index file.

Args:
        prefix: a string, the prefix of a checkpoint.
        format_version: the checkpoint format version that corresponds to the
        prefix.

Returns:
        The pathname of a checkpoint file, taking into account the checkpoint
        format version.


def close(self):
    raise NotImplementedError

Closes the current writer.

    Please see documentation in ``iobase.Sink`` for an example.

Returns:
        An object representing the writes that were performed by the current
        writer.


def _constraint_for_fixed_slice_element(self, builder: expressions.Builder, slice_element: ElementDefinition) -> Optional[expressions.Builder]:
    if not slice_element.HasField('fixed'):
        return None
    fixed_choice = cast(Any, slice_element).fixed
    fixed_message = getattr(fixed_choice, fixed_choice.WhichOneof('choice'))
    return self._constraint_for_slice_element(slice_element, builder, fixed_message, _MessageType.FIXED)

Builds an expression representing the slice constraint's fixed value.

    Builds a FHIRPath expression representing the constraint imposed by the
    given slice's fixed value. Returns `None` if the `slice_element` does
    not define a fixed value. Given a message for a pattern like:

    "fixedCodeableConcept" : {
        "coding" : [{
        "system" : "http://example.org/canonical",
        "code" : "a-code",
        }]
    }

    Builds a constraint like:

    coding[0].system = "http://example.org/canonical" and code = "a-code"
        and display.empty() and version.empty() and userSelected.empty()
    ).exists()

    See more information on slices with fixed values at the following links:
    https://build.fhir.org/profiling.html#discriminator
    https://build.fhir.org/profiling-examples.html#contacts

Args:
        builder: An expression builder to the slice element.
        slice_element: The element definition describing the slice constraint.

Returns:
        An expression builder defining the constraint described by
        `slice_element`'s `fixed` field or `None` if the `slice_element` does
        not have a `fixed` field.


def quant_noise(module: nn.Module, p: float, block_size: int):
    if p <= 0:
        return module
    if not isinstance(module, (nn.Linear, nn.Embedding, nn.Conv2d)):
        raise NotImplementedError('Module unsupported for quant_noise.')
    is_conv = module.weight.ndim == 4
    if not is_conv:
        if module.weight.size(1) % block_size != 0:
            raise AssertionError('Input features must be a multiple of block sizes')
    elif module.kernel_size == (1, 1):
        if module.in_channels % block_size != 0:
            raise AssertionError('Input channels must be a multiple of block sizes')
    else:
        k = module.kernel_size[0] * module.kernel_size[1]
        if k % block_size != 0:
            raise AssertionError('Kernel size must be a multiple of block size')

    def _forward_pre_hook(mod, input):
        if mod.training:
            if not is_conv:
                weight = mod.weight
                in_features = weight.size(1)
                out_features = weight.size(0)
                mask = torch.zeros(in_features // block_size * out_features, device=weight.device)
                mask.bernoulli_(p)
                mask = mask.repeat_interleave(block_size, -1).view(-1, in_features)
            else:
                weight = mod.weight
                in_channels = mod.in_channels
                out_channels = mod.out_channels
                if mod.kernel_size == (1, 1):
                    mask = torch.zeros(int(in_channels // block_size * out_channels), device=weight.device)
                    mask.bernoulli_(p)
                    mask = mask.repeat_interleave(block_size, -1).view(-1, in_channels)
                else:
                    mask = torch.zeros(weight.size(0), weight.size(1), device=weight.device)
                    mask.bernoulli_(p)
                    mask = mask.unsqueeze(2).unsqueeze(3).repeat(1, 1, mod.kernel_size[0], mod.kernel_size[1])
            mask = mask.to(torch.bool)
            s = 1 / (1 - p)
            mod.weight.data = s * weight.masked_fill(mask, 0)
    module.register_forward_pre_hook(_forward_pre_hook)
    return module

From:
    https://github.com/facebookresearch/fairseq/blob/dd0079bde7f678b0cd0715cbd0ae68d661b7226d/fairseq/modules/quant_noise.py

    Wraps modules and applies quantization noise to the weights for subsequent quantization with Iterative Product
    Quantization as described in "Training with Quantization Noise for Extreme Model Compression"

Args:
        - module: nn.Module
        - p: amount of Quantization Noise
        - block_size: size of the blocks for subsequent quantization with iPQ

    Remarks:
        - Module weights must have the right sizes wrt the block size
        - Only Linear, Embedding and Conv2d modules are supported for the moment
        - For more detail on how to quantize by blocks with convolutional weights, see "And the Bit Goes Down:
        Revisiting the Quantization of Neural Networks"
        - We implement the simplest form of noise here as stated in the paper which consists in randomly dropping
        blocks


def __init__(self, vocab_file=None, tokenizer_file=None, do_lower_case=True, unk_token='[UNK]', sep_token='[SEP]', pad_token='[PAD]', cls_token='[CLS]', mask_token='[MASK]', tokenize_chinese_chars=True, strip_accents=None, **kwargs):
    super().__init__(vocab_file, tokenizer_file=tokenizer_file, do_lower_case=do_lower_case, unk_token=unk_token, sep_token=sep_token, pad_token=pad_token, cls_token=cls_token, mask_token=mask_token, tokenize_chinese_chars=tokenize_chinese_chars, strip_accents=strip_accents, **kwargs)
    normalizer_state = json.loads(self.backend_tokenizer.normalizer.__getstate__())
    if normalizer_state.get('lowercase', do_lower_case) != do_lower_case or normalizer_state.get('strip_accents', strip_accents) != strip_accents:
        normalizer_class = getattr(normalizers, normalizer_state.pop('type'))
        normalizer_state['lowercase'] = do_lower_case
        normalizer_state['strip_accents'] = strip_accents
        self.backend_tokenizer.normalizer = normalizer_class(**normalizer_state)
    vocab = self.backend_tokenizer.get_vocab()
    self.backend_tokenizer.pre_tokenizer = PreTokenizer.custom(JiebaPreTokenizer(vocab))
    self.do_lower_case = do_lower_case

Construct a "fast" RoFormer tokenizer (backed by HuggingFace's *tokenizers* library).

    [`RoFormerTokenizerFast`] is almost identical to [`BertTokenizerFast`] and runs end-to-end tokenization:
    punctuation splitting and wordpiece. There are some difference between them when tokenizing Chinese.

    This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should
    refer to this superclass for more information regarding those methods.

Example:
    ```python
    >>> from transformers import RoFormerTokenizerFast

    >>> tokenizer = RoFormerTokenizerFast.from_pretrained("junnyu/roformer_chinese_base")
    >>> tokenizer.tokenize("今天天气非常好。")
    ['今', '天', '天', '气', '非常', '好', '。']
    ```


def softplus(x):
    if any_symbolic_tensors((x,)):
        return Softplus().symbolic_call(x)
    return backend.nn.softplus(x)

Softplus activation function.

    It is defined as `f(x) = log(exp(x) + 1)`, where `log` is the natural
    logarithm and `exp` is the exponential function.

Args:
        x: Input tensor.

Returns:
        A tensor with the same shape as `x`.

Example:
    >>> x = keras.ops.convert_to_tensor([-0.555, 0.0, 0.555])
    >>> keras.ops.softplus(x)
    array([0.45366603, 0.6931472, 1.008666], dtype=float32)


def __init__(self, feature_extractor, tokenizer):
    super().__init__(feature_extractor, tokenizer)
    self.current_processor = self.feature_extractor
    self._in_target_context_manager = False

Constructs a MusicGen processor which wraps an EnCodec feature extractor and a T5 tokenizer into a single processor
    class.

    [`MusicgenProcessor`] offers all the functionalities of [`EncodecFeatureExtractor`] and [`TTokenizer`]. See
    [`~MusicgenProcessor.__call__`] and [`~MusicgenProcessor.decode`] for more information.

Args:
        feature_extractor (`EncodecFeatureExtractor`):
            An instance of [`EncodecFeatureExtractor`]. The feature extractor is a required input.
        tokenizer (`T5Tokenizer`):
            An instance of [`T5Tokenizer`]. The tokenizer is a required input.


def forward(self, index: int, output: torch.Tensor, multi_stage_features: List[torch.Tensor], multi_stage_positional_embeddings: List[torch.Tensor], attention_mask: Optional[torch.Tensor]=None, query_embeddings: Optional[torch.Tensor]=None, output_attentions: Optional[bool]=False):
    level_index = index % self.num_feature_levels
    attention_mask[torch.where(attention_mask.sum(-1) == attention_mask.shape[-1])] = False
    output, cross_attn_weights = self.cross_attn(output, multi_stage_features[level_index], memory_mask=attention_mask, memory_key_padding_mask=None, pos=multi_stage_positional_embeddings[level_index], query_pos=query_embeddings)
    output, self_attn_weights = self.self_attn(output, output_mask=None, output_key_padding_mask=None, query_pos=query_embeddings)
    output = self.ffn(output)
    outputs = (output,)
    if output_attentions:
        outputs += (self_attn_weights, cross_attn_weights)
    return outputs

Args:
    index (`int`): index of the layer in the Transformer decoder.
    output (`torch.FloatTensor`): the object queries of shape `(N, batch, hidden_dim)`
    multi_stage_features (`List[torch.Tensor]`): the multi-scale features from the pixel decoder.
    multi_stage_positional_embeddings (`List[torch.Tensor]`):
        positional embeddings for the multi_stage_features
    attention_mask (`torch.FloatTensor`): attention mask for the masked cross attention layer
    query_embeddings (`torch.FloatTensor`, *optional*):
        position embeddings that are added to the queries and keys in the self-attention layer.
    output_attentions (`bool`, *optional*):
        Whether or not to return the attentions tensors of all attention layers. See `attentions` under
        returned tensors for more detail.


def _unique_urls(urls: Iterable[str]) -> Iterable[str]:
    seen: Set[str] = set()
    for url in urls:
        if url not in seen:
            seen.add(url)
            yield url

Filters URLs to remove duplicates.

Args:
        urls: The URLs to filter.

Yields:
        The URLs filtered to only those without duplicates.


def broadcast_recv_v2(shape, dtype, group_size, group_key, instance_key, communication_hint='auto', timeout=0):
    return gen_collective_ops.collective_bcast_recv_v2(T=dtype, group_size=group_size, group_key=group_key, instance_key=instance_key, shape=shape, communication_hint=communication_hint.lower(), timeout_seconds=timeout)

Receives a broadcasts tensor, across devices.

Args:
        shape: an int tensor.  Shape of the tensor to be received.
        dtype: Type of the tensor to be received.
        group_size: an int32 tensor.  One plus the number of receiving tensors, i.e.
        the total number of devices participating.  Each tensor must reside on a
        different device.
        group_key: an int32 tensor identifying the group of devices.
        instance_key: an int32 tensor identifying the participating group of Ops.
        communication_hint: preferred collective communication.  The implementation
        may fall back to another mechanism.  Options include `auto`, `ring`, and
        `nccl`.
        timeout: If set to a non zero, set a completion timeout to detect staleness.
        If the timer goes off, a DeadlineExceededError is raised.
        The timeout value in seconds. This feature is experimental.

Returns:
        An Op implementing the broadcast receive.


def get_text_features(self, input_ids, attention_mask=None, position_ids=None, params: Optional[dict]=None, dropout_rng: jax.random.PRNGKey=None, train=False):
    if position_ids is None:
        position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_ids.shape)
    if attention_mask is None:
        attention_mask = jnp.ones_like(input_ids)
    rngs = {}
    if dropout_rng is not None:
        rngs['dropout'] = dropout_rng

    def _get_features(module, input_ids, attention_mask, position_ids, deterministic):
        text_outputs = module.text_model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, deterministic=deterministic)
        pooled_output = text_outputs[1]
        text_features = module.text_projection(pooled_output)
        return text_features
    return self.module.apply({'params': params or self.params}, jnp.array(input_ids, dtype='i4'), jnp.array(attention_mask, dtype='i4'), jnp.array(position_ids, dtype='i4'), not train, method=_get_features, rngs=rngs)

Args:
    input_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`):
        Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you
        provide it.

        Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
        [`PreTrainedTokenizer.__call__`] for details.

        [What are input IDs?](../glossary#input-ids)

Returns:
    text_features (`jnp.ndarray` of shape `(batch_size, output_dim`): The text embeddings obtained by applying
    the projection layer to the pooled output of [`FlaxCLIPTextModel`].

Example:
    ```python
    >>> from transformers import AutoTokenizer, FlaxCLIPModel

    >>> model = FlaxCLIPModel.from_pretrained("openai/clip-vit-base-patch32")
    >>> tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32")

    >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="np")
    >>> text_features = model.get_text_features(**inputs)
    ```


def hinge(y_true, y_pred):
    y_pred = tensor_conversion.convert_to_tensor_v2_with_dispatch(y_pred)
    y_true = math_ops.cast(y_true, y_pred.dtype)
    y_true = _maybe_convert_labels(y_true)
    return backend.mean(math_ops.maximum(1.0 - y_true * y_pred, 0.0), axis=-1)

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

    `loss = mean(maximum(1 - y_true * y_pred, 0), axis=-1)`

    Standalone usage:

    >>> y_true = np.random.choice([-1, 1], size=(2, 3))
    >>> y_pred = np.random.random(size=(2, 3))
    >>> loss = tf.keras.losses.hinge(y_true, y_pred)
    >>> assert loss.shape == (2,)
    >>> assert np.array_equal(
    ...     loss.numpy(),
    ...     np.mean(np.maximum(1. - y_true * y_pred, 0.), axis=-1))

Args:
        y_true: The ground truth values. `y_true` values are expected to be -1 or 1.
        If binary (0 or 1) labels are provided they will be converted to -1 or 1.
        shape = `[batch_size, d0, .. dN]`.
        y_pred: The predicted values. shape = `[batch_size, d0, .. dN]`.

Returns:
        Hinge loss values. shape = `[batch_size, d0, .. dN-1]`.


def get_flat_tensor_shapes(element_spec):
    return [spec.shape for spec in get_flat_tensor_specs(element_spec)]

Returns a list `tf.TensorShapes`s for the element tensor representation.

Args:
        element_spec: A nested structure of `tf.TypeSpec` objects representing to
        element type specification.

Returns:
        A list `tf.TensorShapes`s for the element tensor representation.


def create_distributed_mesh(mesh_dims: Union[List[Tuple[str, int]], Dict[str, int]], mesh_name: str='', local_devices: Optional[List[Union[tf_device.DeviceSpec, str]]]=None, device_type: Optional[str]=None, use_xla_spmd: bool=layout.USE_XLA_SPMD) -> layout.Mesh:
    if isinstance(mesh_dims, dict):
        mesh_dims = list(mesh_dims.items())
    dim_names, shape = zip(*mesh_dims)
    if not accelerator_util.is_initialized():
        raise ValueError('Accelerators are uninitialized, please run dtensor.initialize_accelerator_system() first.')
    if device_type and device_type.upper() == 'TPU':
        if local_devices is not None:
            raise ValueError(f'Do not specify devices for {device_type.upper()} meshes. Using a partial list of devices for {device_type.upper()} is not supported.')
    device_specs, device_type = _make_device_specs(local_devices, device_type)
    if device_type.upper() in ['CPU', 'GPU']:
        local_spec = tf_device.DeviceSpec(job=config.job_name(), replica=0, task=config.client_id())
        device_specs = [local_spec.make_merged_spec(d) for d in device_specs]
        num_global_devices = len(device_specs) * config.num_clients()
        if np.prod(shape) != num_global_devices:
            raise ValueError(f'Global number of devices ({len(device_specs)} per client * {config.num_clients()} clients = {num_global_devices}) must be equal to total size of the mesh of shape {shape}')
        global_device_ids = np.arange(num_global_devices).reshape(shape)
        flattened = np.ravel(global_device_ids).tolist()
        start_idx = len(device_specs) * config.client_id()
        local_device_ids = flattened[start_idx:start_idx + len(device_specs)]
        mesh = layout.Mesh(dim_names=dim_names, global_device_ids=global_device_ids, local_device_ids=local_device_ids, local_devices=device_specs, mesh_name=mesh_name, use_xla_spmd=use_xla_spmd)
        _print_context(num_global_devices, config.num_clients(), config.client_id(), device_type, mesh)
        return mesh
    if device_type.upper() == 'TPU':
        mesh = tpu_util.create_tpu_mesh(mesh_dim_names=dim_names, mesh_shape=shape, mesh_name=mesh_name, use_xla_spmd=use_xla_spmd)
        _print_context(config.num_global_devices(device_type), config.num_clients(), config.client_id(), device_type, mesh)
        return mesh
    raise ValueError(f'Device type {device_type} is not CPU, GPU or TPU')

Creates a distributed mesh.

    This is similar to `create_mesh`, but with a different set of arguments to
    create a mesh that spans evenly across a multi-client DTensor cluster.

    For CPU and GPU meshes, users can choose to use fewer local devices than what
    is available `local_devices`.

    For TPU, only meshes that uses all TPU cores is supported by the DTensor
    runtime.

Args:
        mesh_dims: A dict of dim_name: dim_size, or a list of (dim_name, dim_size)
        tuples. e.g. `{'x' : 4, 'y' : 1}` or `[('x', 4), ('y', 1)]`.
        mesh_name: Name of the created mesh. Defaults to ''.
        local_devices: String representations of devices to use. This is the device
        part of tf.DeviceSpec, e.g. 'CPU:0'. Defaults to all available local
        logical devices.
        device_type: Type of device to build the mesh for. Defaults to 'CPU'.
        Supported values are 'CPU', 'GPU', 'TPU'.6
        use_xla_spmd: Boolean when True, will use XLA SPMD instead of DTensor SPMD.

Returns:
        A mesh that spans evenly across all DTensor clients in the cluster.


def _generate_signatures(signature_functions: dict[str, Callable[..., Any]], object_map: object_identity.ObjectIdentityDictionary, defaults=None):
    signatures = {}
    for signature_key, function in sorted(signature_functions.items()):
        if function.graph.captures:
            argument_inputs = function.graph.inputs[:-len(function.graph.captures)]
        else:
            argument_inputs = function.graph.inputs
        mapped_inputs, exterior_argument_placeholders = _map_function_arguments_to_created_inputs(argument_inputs, signature_key, function.name, defaults)
        kwarg_names = list(sorted(object_map[function].function.structured_input_signature[1].keys()))
        outputs = object_map[function](**{kwarg_name: mapped_input for kwarg_name, mapped_input in zip(kwarg_names, mapped_inputs)})
        signatures[signature_key] = signature_def_utils.build_signature_def(_tensor_dict_to_tensorinfo(exterior_argument_placeholders), _tensor_dict_to_tensorinfo(outputs), method_name=signature_constants.PREDICT_METHOD_NAME, defaults=defaults.get(signature_key, None))
    return signatures

Validates and calls `signature_functions` in the exported graph.

Args:
        signature_functions: A dictionary mapping string keys to concrete TensorFlow
        functions (e.g. from `signature_serialization.canonicalize_signatures`)
        which will be used to generate SignatureDefs.
        object_map: A dictionary that contains mappings from signature functions to
        concrete functions in the exported graph.
        defaults: A dictionary mapping signature_key to dictionary of
        user_specified_name to Tensor representing default values.

Returns:
        Each function in the `signature_functions` dictionary is called with
        placeholder Tensors, generating a function call operation and output
        Tensors. The placeholder Tensors, the function call operation, and the
        output Tensors from the function call are part of the default Graph.

        This function then returns a dictionary with the same structure as
        `signature_functions`, with the concrete functions replaced by SignatureDefs
        implicitly containing information about how to call each function from a
        TensorFlow 1.x Session / the C++ Loader API. These SignatureDefs reference
        the generated placeholders and Tensor outputs by name.

        The caller is expected to include the default Graph set while calling this
        function as a MetaGraph in a SavedModel, including the returned
        SignatureDefs as part of that MetaGraph.


def get_extra_inputs():
    g = ops.get_default_graph()
    if isinstance(g, _FuncGraph):
        return g.extra_inputs
    else:
        return []

Returns the captured input tensors by the function.

Returns:
        If the default graph is being used to define a function, the
        returned list of tensors are those accessed inside the function body
        but defined outside the function body so far. Otherwise, returns an
        empty list.


def get_sequence_dense_tensor(self, transformation_cache, state_manager):
    pass

Returns a `TensorSequenceLengthPair`.

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


def _get_imports_for_module(module: str, output_package: str, symbols_by_module: Mapping[str, set[_Entrypoint]], generated_imports_by_module: Mapping[str, set[str]], file_prefixes_to_strip: Sequence[str], module_prefix: str, use_lazy_loading: bool, subpackage_rewrite: Optional[str]) -> str:
    content = ''
    symbol_imports = list(symbols_by_module[module])
    symbol_imports = sorted(symbol_imports, key=lambda s: f'{s.exported_symbol.file_name}:{s.name}')
    generated_imports = sorted(generated_imports_by_module[module])
    for imp in generated_imports:
        if subpackage_rewrite:
            imp = imp.replace(output_package, subpackage_rewrite)
        last_dot = imp.rfind('.')
        if use_lazy_loading:
            content += f"  '{imp[last_dot + 1:]}': ('', '{imp}'),\n"
        else:
            content += f'from {imp[:last_dot]} import {imp[last_dot + 1:]}\n'
    for s in symbol_imports:
        content += f'{s.get_import(file_prefixes_to_strip, module_prefix, use_lazy_loading=use_lazy_loading)}\n'
    return content

Returns the imports for a module.

Args:
        module: The module to get imports for.
        output_package: The package to use for the imports.
        symbols_by_module: The symbols that should be exposed by each module.
        generated_imports_by_module: The sub-modules that should be exposed by each
        module.
        file_prefixes_to_strip: The prefixes to strip from the file names of the
        imports.
        module_prefix: A prefix to add to the non-generated imports.
        use_lazy_loading: Whether to use lazy loading or not.
        subpackage_rewrite: The subpackage to use for the imports.


def preferred_version(api_name: str, key: str=None) -> str:
    api_url = 'https://discovery.googleapis.com/discovery/v1/apis?name=%s&key=%s&preferred=true' % (api_name, key or '')
    print('DISCOVERY FETCH:', api_url)
    api_info = json.load(request.urlopen(api_url))
    return api_info['items'][0]['version']

Helper to get default API version.

Args:
        api_name: The API endpoint name, for example dfareporting.
        key: Optional Google API Key: https://cloud.google.com/docs/authentication/api-keys

    Returns (str):
        The API version.

Raises:
        HttpError: If the wrong API values are specified.


def _all_gather(self, input_tensor: core.TensorLike, options: Optional[collective_util.Options]) -> core.Tensor:
    instance_key = self._next_instance_key()
    options = self._options.merge(options)
    ordering_token = self._get_ordering_token()
    with ops.device(self._device):
        return collective_ops.all_gather_v2(input_tensor, self._group_size, self._group_key, instance_key, communication_hint=options.implementation.value, timeout=options.timeout_seconds, ordering_token=ordering_token)

All-gather a dense tensor.

Args:
        input_tensor: a dense tensor. It must have the same shape on all replicas.
        options: an optional tf.distribute.experimental.CommunicationOptions. If
        provided, it overrides the default options.

Returns:
        The reduced tensor.


def forward(self, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.FloatTensor]=None, token_type_ids: Optional[torch.LongTensor]=None, position_ids: Optional[torch.LongTensor]=None, head_mask: Optional[torch.FloatTensor]=None, inputs_embeds: Optional[torch.FloatTensor]=None, encoder_hidden_states: Optional[torch.FloatTensor]=None, encoder_attention_mask: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, **kwargs) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    if labels is not None:
        use_cache = False
    outputs = self.roberta(input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, past_key_values=past_key_values, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    sequence_output = outputs[0]
    prediction_scores = self.lm_head(sequence_output)
    lm_loss = None
    if labels is not None:
        lm_loss = self.loss_function(prediction_scores, labels, vocab_size=self.config.vocab_size, **kwargs)
    if not return_dict:
        output = (prediction_scores,) + outputs[2:]
        return (lm_loss,) + output if lm_loss is not None else output
    return CausalLMOutputWithCrossAttentions(loss=lm_loss, logits=prediction_scores, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, cross_attentions=outputs.cross_attentions)

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

Example:
    ```python
    >>> from transformers import AutoTokenizer, RobertaForCausalLM, RobertaConfig
    >>> import torch

    >>> tokenizer = AutoTokenizer.from_pretrained("FacebookAI/roberta-base")
    >>> config = RobertaConfig.from_pretrained("FacebookAI/roberta-base")
    >>> config.is_decoder = True
    >>> model = RobertaForCausalLM.from_pretrained("FacebookAI/roberta-base", config=config)
    >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
    >>> outputs = model(**inputs)
    >>> prediction_logits = outputs.logits
    ```


def HandleCurlError(e, logger=None):
    if not logger:
        logger = logging
    code = e.args[0]
    msg = e.args[1]
    if code in (pycurl.E_UNSUPPORTED_PROTOCOL, pycurl.E_URL_MALFORMAT, pycurl.E_SSL_ENGINE_NOTFOUND, pycurl.E_SSL_ENGINE_SETFAILED, pycurl.E_SSL_CACERT_BADFILE):
        raise error.ConfigurationError(msg)
    if code in (pycurl.E_FAILED_INIT, pycurl.E_COULDNT_CONNECT, pycurl.E_PARTIAL_FILE, pycurl.E_WRITE_ERROR, pycurl.E_READ_ERROR, pycurl.E_OPERATION_TIMEOUTED, pycurl.E_SSL_CONNECT_ERROR, pycurl.E_COULDNT_RESOLVE_PROXY, pycurl.E_COULDNT_RESOLVE_HOST, pycurl.E_GOT_NOTHING):
        logger.debug('Possibly transient error: %s', msg)
        return
    if code in (pycurl.E_SSL_PEER_CERTIFICATE,):
        raise error.SourceUnavailable(msg)
    raise error.Error(msg)

Handle a curl exception.

    See http://curl.haxx.se/libcurl/c/libcurl-errors.html for a list of codes.

Args:
        e: pycurl.error
        logger: logger object

Raises:
        ConfigurationError:
        PermissionDenied:
        SourceUnavailable:
        Error:


def __init__(self, col, row, is_non_singular=None, is_self_adjoint=None, is_positive_definite=None, is_square=None, name='LinearOperatorToeplitz'):
    parameters = dict(col=col, row=row, is_non_singular=is_non_singular, is_self_adjoint=is_self_adjoint, is_positive_definite=is_positive_definite, is_square=is_square, name=name)
    with ops.name_scope(name, values=[row, col]):
        self._row = linear_operator_util.convert_nonref_to_tensor(row, name='row')
        self._col = linear_operator_util.convert_nonref_to_tensor(col, name='col')
        self._check_row_col(self._row, self._col)
        if is_square is False:
            raise ValueError('Only square Toeplitz operators currently supported.')
        is_square = True
        super(LinearOperatorToeplitz, self).__init__(dtype=self._row.dtype, is_non_singular=is_non_singular, is_self_adjoint=is_self_adjoint, is_positive_definite=is_positive_definite, is_square=is_square, parameters=parameters, name=name)

Initialize a `LinearOperatorToeplitz`.

Args:
        col: Shape `[B1,...,Bb, N]` `Tensor` with `b >= 0` `N >= 0`.
        The first column of the operator. Allowed dtypes: `float16`, `float32`,
        `float64`, `complex64`, `complex128`. Note that the first entry of
        `col` is assumed to be the same as the first entry of `row`.
        row: Shape `[B1,...,Bb, N]` `Tensor` with `b >= 0` `N >= 0`.
        The first row of the operator. Allowed dtypes: `float16`, `float32`,
        `float64`, `complex64`, `complex128`. Note that the first entry of
        `row` is assumed to be the same as the first entry of `col`.
        is_non_singular:  Expect that this operator is non-singular.
        is_self_adjoint:  Expect that this operator is equal to its hermitian
        transpose.  If `diag.dtype` is real, this is auto-set to `True`.
        is_positive_definite:  Expect that this operator is positive definite,
        meaning the quadratic form `x^H A x` has positive real part for all
        nonzero `x`.  Note that we do not require the operator to be
        self-adjoint to be positive-definite.  See:
        https://en.wikipedia.org/wiki/Positive-definite_matrix#Extension_for_non-symmetric_matrices
        is_square:  Expect that this operator acts like square [batch] matrices.
        name: A name for this `LinearOperator`.


def merge(value_list: List[Any], merge_fn: Optional[Callable[[KeyPath, Any, Any], Any]]=None) -> Any:
    if not isinstance(value_list, list):
        raise TypeError('value_list should be a list')
    if not value_list:
        return None
    new_value = canonicalize(value_list[0], sparse_list_as_dict=True)
    for value in value_list[1:]:
        if value is None:
            continue
        new_value = merge_tree(new_value, canonicalize(value, sparse_list_as_dict=True), merge_fn)

    def _listify_dict_equivalent(p, v):
        del p
        if isinstance(v, dict):
            v = try_listify_dict_with_int_keys(v, True)[0]
        return v
    return transform(new_value, _listify_dict_equivalent)

Merge a list of hierarchical values.

    Example::

        original = {
        'a': 1,
        'b': 2,
        'c': {
            'd': 'foo',
            'e': 'bar'
        }
        }
        patch =  {
        'b': 3,
        'd': [1, 2, 3],
        'c': {
            'e': 'bar2',
            'f': 10
        }
        }
        output = pg.utils.merge([original, patch])
        assert output == {
        'a': 1,
        # b is updated.
        'b': 3,
        'c': {
            'd': 'foo',
            # e is updated.
            'e': 'bar2',
            # f is added.
            'f': 10
        },
        # d is inserted.
        'd': [1, 2, 3]
        })

Args:
        value_list: A list of hierarchical values to merge. Later value will be
        treated as updates if it's a dict or otherwise a replacement of former
        value. The merge process will keep input values intact.
        merge_fn: A function to handle value merge that will be called for updated
        or added keys. If a branch is added/updated, the root of branch will be
        passed to merge_fn. the signature of function is: `(path, left_value,
        right_value) -> final_value` If a key is only present in src dict,
        old_value is MISSING_VALUE; If a key is only present in dest dict,
        new_value is MISSING_VALUE; otherwise both new_value and old_value are
        filled. If final_value is MISSING_VALUE for a path, it will be removed
        from its parent collection.

Returns:
        A merged value.

Raises:
        TypeError: If `value_list` is not a list.
        KeyError: If new key is found while not allowed.


def __init__(self, loss_tensor, fail_on_nan_loss=True):
    self._loss_tensor = loss_tensor
    self._fail_on_nan_loss = fail_on_nan_loss

Initializes a `NanTensorHook`.

Args:
        loss_tensor: `Tensor`, the loss tensor.
        fail_on_nan_loss: `bool`, whether to raise exception when loss is NaN.


def extract_variable_info(kwargs: Any) -> Tuple[str, Tuple[int, ...], dtypes.DType, Callable[[], Any], Optional[int]]:

    def get_restore_uid(initial_value: Callable[..., Any]) -> int | None:
        return getattr(initial_value, 'restore_uid', None)
    if isinstance(kwargs['initial_value'], functools.partial) and ('shape' in kwargs['initial_value'].keywords or kwargs['initial_value'].args):
        if 'shape' in kwargs['initial_value'].keywords:
            shape = kwargs['initial_value'].keywords['shape']
        else:
            shape = kwargs['initial_value'].args[0]
        return (kwargs['name'], shape, kwargs['initial_value'].keywords.get('dtype', kwargs['dtype']), kwargs['initial_value'].func, get_restore_uid(kwargs['initial_value'].func))
    elif 'shape' not in kwargs or kwargs['shape'] is None or (not callable(kwargs['initial_value'])):
        raise ValueError('Unable to extract initializer function and shape from {}. Please either pass a function that expects a shape and dtype as the initial value for your variable or functools.partial object with the shape and dtype kwargs set. This is needed so that we can initialize the shards of the ShardedVariable locally.'.format(kwargs['initial_value']))
    else:
        return (kwargs['name'], kwargs['shape'], kwargs['dtype'], kwargs['initial_value'], get_restore_uid(kwargs['initial_value']))

Extracts the variable creation attributes from the kwargs.

Args:
        kwargs: a dict of keyword arguments that were passed to a variable creator
        scope.

Returns:
        A tuple of variable name, shape, dtype, initialization function,
        restore_uid.


def diff(a, n=1, axis=-1):
    return Diff(n=n, axis=axis)(a)

Calculate the n-th discrete difference along the given axis.

    The first difference is given by `out[i] = a[i+1] - a[i]` along
    the given axis, higher differences are calculated by using `diff`
    recursively.

Args:
        a: Input tensor.
        n: The number of times values are differenced. Defaults to `1`.
        axis: Axis to compute discrete difference(s) along.
            Defaults to `-1`.(last axis).

Returns:
        Tensor of diagonals.

Example:
    >>> from keras.src import ops
    >>> x = ops.convert_to_tensor([1, 2, 4, 7, 0])
    >>> ops.diff(x)
    array([ 1,  2,  3, -7])
    >>> ops.diff(x, n=2)
    array([  1,   1, -10])

    >>> x = ops.convert_to_tensor([[1, 3, 6, 10], [0, 5, 6, 8]])
    >>> ops.diff(x)
    array([[2, 3, 4],
        [5, 1, 2]])
    >>> ops.diff(x, axis=0)
    array([[-1,  2,  0, -2]])


def _validate_inputs(self, graph_def, input_tensors):
    self._save_conversion_params_metric(graph_def)
    self._quant_mode = QuantizationMode(self.optimizations, self.target_spec, self.representative_dataset, graph_def, self._experimental_disable_per_channel, self.experimental_new_dynamic_range_quantizer, self._experimental_low_bit_qat, self._experimental_full_integer_quantization_bias_type, self._experimental_variable_quantization, self._experimental_strict_qdq)
    self._validate_inference_input_output_types(self._quant_mode)
    if not self._is_unknown_shapes_allowed():
        for tensor in input_tensors:
            shape_list = tensor.shape.as_list()
            if None in shape_list[1:]:
                raise ValueError("None is only supported in the 1st dimension. Tensor '{0}' has invalid shape '{1}'.".format(_get_tensor_name(tensor), shape_list))
            elif shape_list and shape_list[0] is None:
                shape = tensor.shape.as_list()
                shape[0] = 1
                tensor.set_shape(shape)
    if self._trackable_obj is None or not hasattr(self._trackable_obj, 'graph_debug_info'):
        self._debug_info = _get_debug_info(_build_debug_info_func(self._funcs[0].graph), graph_def)
    else:
        self._debug_info = _get_debug_info(_convert_debug_info_func(self._trackable_obj.graph_debug_info), graph_def)

Validate the input parameters.

Args:
        graph_def: The TensorFlow GraphDef.
        input_tensors: List of input tensors.

Raises:
        ValueError: Input shape is not specified. Invalid quantization parameters.


def enable_graph_building_optimization(fn: _F) -> _F:

    def wrapper(*args, **kwargs):
        if flags.config().graph_building_optimization.value():
            return fn(*args, **kwargs)
        flags.config().graph_building_optimization.reset(True)
        try:
            return fn(*args, **kwargs)
        finally:
            flags.config().graph_building_optimization.reset(False)
    return wrapper

Decorator for enabling graph_building_optimization on a test.

    This function returns a decorator intended to be applied to test methods in
    a `tf.test.TestCase` class. Doing so will enable graph_building_optimization,
    execute the test, then reset the feature flag to its default value.

Example:
    class MyTest(test.TestCase):

        @enable_graph_building_optimization
        def testFoo(self):
        ...

Args:
        fn: the function to be wrapped.

Returns:
        The wrapped function.


def lu_reconstruct(lower_upper, perm, validate_args=False, name=None):
    with ops.name_scope(name or 'lu_reconstruct'):
        lower_upper = ops.convert_to_tensor(lower_upper, dtype_hint=dtypes.float32, name='lower_upper')
        perm = ops.convert_to_tensor(perm, dtype_hint=dtypes.int32, name='perm')
        assertions = lu_reconstruct_assertions(lower_upper, perm, validate_args)
        if assertions:
            with ops.control_dependencies(assertions):
                lower_upper = array_ops.identity(lower_upper)
                perm = array_ops.identity(perm)
        shape = array_ops.shape(lower_upper)
        lower = set_diag(band_part(lower_upper, num_lower=-1, num_upper=0), array_ops.ones(shape[:-1], dtype=lower_upper.dtype))
        upper = band_part(lower_upper, num_lower=0, num_upper=-1)
        x = math_ops.matmul(lower, upper)
        if lower_upper.shape is None or lower_upper.shape.rank is None or lower_upper.shape.rank != 2:
            batch_size = math_ops.reduce_prod(shape[:-2])
            d = shape[-1]
            x = array_ops.reshape(x, [batch_size, d, d])
            perm = array_ops.reshape(perm, [batch_size, d])
            perm = map_fn.map_fn(array_ops.invert_permutation, perm)
            batch_indices = array_ops.broadcast_to(math_ops.range(batch_size)[:, array_ops.newaxis], [batch_size, d])
            x = array_ops.gather_nd(x, array_ops_stack.stack([batch_indices, perm], axis=-1))
            x = array_ops.reshape(x, shape)
        else:
            x = array_ops.gather(x, array_ops.invert_permutation(perm))
        x.set_shape(lower_upper.shape)
        return x

The reconstruct one or more matrices from their LU decomposition(s).

Args:
        lower_upper: `lu` as returned by `tf.linalg.lu`, i.e., if `matmul(P,
        matmul(L, U)) = X` then `lower_upper = L + U - eye`.
        perm: `p` as returned by `tf.linag.lu`, i.e., if `matmul(P, matmul(L, U)) =
        X` then `perm = argmax(P)`.
        validate_args: Python `bool` indicating whether arguments should be checked
        for correctness.
        Default value: `False` (i.e., don't validate arguments).
        name: Python `str` name given to ops managed by this object.
        Default value: `None` (i.e., 'lu_reconstruct').

Returns:
        x: The original input to `tf.linalg.lu`, i.e., `x` as in,
        `lu_reconstruct(*tf.linalg.lu(x))`.

    #### Examples

    ```python
    import numpy as np
    import tensorflow as tf
    import tensorflow_probability as tfp

    x = [[[3., 4], [1, 2]],
        [[7., 8], [3, 4]]]
    x_reconstructed = tf.linalg.lu_reconstruct(*tf.linalg.lu(x))
    tf.assert_near(x, x_reconstructed)
    # ==> True
    ```


def get_next_as_list(self, name=None):
    del name
    with ops.device(self._worker):
        return self._format_data_list_with_options(self._iterator.get_next())

Get next element from the underlying iterator.

    Runs the iterator get_next() within a device scope. Since this doesn't use
    get_next_as_optional(), it is considerably faster than get_next_as_list(),
    but it raises EOFError if any of the device doesn't get any data.

Args:
        name: not used.

Returns:
        A list consisting of the next data from each device.


def encode(self, input_audio, start_level=0, end_level=None, bs_chunks=1):
    audio_chunks = torch.chunk(input_audio, bs_chunks, dim=0)
    music_tokens_list = []
    for chunk_i in audio_chunks:
        music_tokens_i = self._encode(chunk_i, start_level=start_level, end_level=end_level)
        music_tokens_list.append(music_tokens_i)
    music_tokens = [torch.cat(music_tokens_level, dim=0) for music_tokens_level in zip(*music_tokens_list)]
    return music_tokens

Transforms the `input_audio` to a discrete representation made out of `music_tokens`.

Args:
        input_audio (`torch.Tensor`):
            Raw audio which will be encoded to its discrete representation using the codebook. The closest `code`
            form the codebook will be computed for each sequence of samples.
        start_level (`int`, *optional*, defaults to 0):
            Level at which the encoding process will start. Default to 0.
        end_level (`int`, *optional*):
            Level at which the encoding process will start. Default to None.
        bs_chunks (int, *optional*, defaults to 1):
            Number of chunks of raw audio to process at the same time.


def show_path(from_op, tensors, sources):
    if isinstance(from_op, tensor_lib.Tensor):
        from_op = from_op.op
    if not isinstance(tensors, list):
        tensors = [tensors]
    final_ops = [_as_operation(tensor) for tensor in tensors]
    visited_ops = set((x.op for x in sources))
    ops_to_visit = list(final_ops)
    some_op_output = {}
    while ops_to_visit:
        op = ops_to_visit.pop()
        if op in visited_ops:
            continue
        visited_ops.add(op)
        if op == from_op:
            path_op = op
            path = [path_op]
            while path_op not in final_ops:
                path_op = some_op_output[path_op]
                path.append(path_op)
            return ' <- '.join(('%s (%s)' % (x.name, x.type) for x in reversed(path)))
        else:
            for inp in graph_inputs(op):
                if inp not in visited_ops and inp not in sources:
                    some_op_output[inp] = op
                    ops_to_visit.append(inp)
    return '??'

Find one path from `from_op` to any of `tensors`, ignoring `sources`.

Args:
        from_op: A `tf.Operation`.
        tensors: A `tf.Operation`, a `tf.Tensor`, or a list thereof.
        sources: A list of `tf.Tensor`.

Returns:
        A python string containing the path, or "??" if none is found.


def import_event(tensor, name=None):
    return gen_summary_ops.import_event(_summary_state.writer._resource, tensor, name=name)

Writes a `tf.compat.v1.Event` binary proto.

    This can be used to import existing event logs into a new summary writer sink.
    Please note that this is lower level than the other summary functions and
    will ignore the `tf.summary.should_record_summaries` setting.

Args:
        tensor: A `tf.Tensor` of type `string` containing a serialized
        `tf.compat.v1.Event` proto.
        name: A name for the operation (optional).

Returns:
        The created `tf.Operation`.


def __init__(self, group_key_start=1):
    self._group_key = group_key_start
    self._instance_key_table = {}
    self._lock = threading.Lock()
    self._known_groups = {}

Initializes the object.

Args:
        group_key_start: the starting integer of group key.


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

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

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

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


def compute_predicted_aligned_error(logits: torch.Tensor, max_bin: int=31, no_bins: int=64, **kwargs) -> Dict[str, torch.Tensor]:
    boundaries = torch.linspace(0, max_bin, steps=no_bins - 1, device=logits.device)
    aligned_confidence_probs = torch.nn.functional.softmax(logits, dim=-1)
    predicted_aligned_error, max_predicted_aligned_error = _calculate_expected_aligned_error(alignment_confidence_breaks=boundaries, aligned_distance_error_probs=aligned_confidence_probs)
    return {'aligned_confidence_probs': aligned_confidence_probs, 'predicted_aligned_error': predicted_aligned_error, 'max_predicted_aligned_error': max_predicted_aligned_error}

Computes aligned confidence metrics from logits.

Args:
        logits: [*, num_res, num_res, num_bins] the logits output from
        PredictedAlignedErrorHead.
        max_bin: Maximum bin value
        no_bins: Number of bins

Returns:
        aligned_confidence_probs: [*, num_res, num_res, num_bins] the predicted
        aligned error probabilities over bins for each residue pair.
        predicted_aligned_error: [*, num_res, num_res] the expected aligned distance
        error for each pair of residues.
        max_predicted_aligned_error: [*] the maximum predicted error possible.


def strip_strings(model):
    model.description = None
    for subgraph in model.subgraphs:
        subgraph.name = None
        for tensor in subgraph.tensors:
            tensor.name = None
    model.signatureDefs = None

Strips all nonessential strings from the model to reduce model size.

    We remove the following strings:
    (find strings by searching ":string" in the tensorflow lite flatbuffer schema)
    1. Model description
    2. SubGraph name
    3. Tensor names
    We retain OperatorCode custom_code and Metadata name.

Args:
        model: The model from which to remove nonessential strings.


def _split_bytecode(bytecode: list[opcodes.Opcode], processed_blocks: set[Block], python_version) -> list[Block]:
    targets = {op.target for op in bytecode if op.target}
    blocks = []
    code = []
    prev_block: Block = None
    i = 0
    while i < len(bytecode):
        op = bytecode[i]
        if python_version >= (3, 12) and isinstance(op, opcodes.SEND):
            if code:
                prev_block = Block(code)
                blocks.append(prev_block)
                code = []
            new_blocks, i = _preprocess_async_for_and_yield(i, bytecode, prev_block, processed_blocks)
            blocks.extend(new_blocks)
            prev_block = blocks[-1]
            continue
        code.append(op)
        if op.no_next() or op.does_jump() or op.pops_block() or (op.next is None) or (op.next in targets and (not isinstance(op.next, opcodes.GET_ANEXT) or python_version < (3, 12))):
            prev_block = Block(code)
            blocks.append(prev_block)
            code = []
        i += 1
    return blocks

Given a sequence of bytecodes, return basic blocks.

    This will split the code at "basic block boundaries". These occur at
    every instruction that is jumped to, and after every instruction that jumps
    somewhere else (or returns / aborts).

Args:
        bytecode: A list of instances of opcodes.Opcode. (E.g. returned from
        opcodes.dis())

Returns:
        A list of _Block instances.


def floating_point_ops(self, input_dict: Dict[str, Union[torch.Tensor, Any]], exclude_embeddings: bool=True) -> int:
    return 6 * self.estimate_tokens(input_dict) * self.num_parameters(exclude_embeddings=exclude_embeddings)

Get number of (optionally, non-embeddings) floating-point operations for the forward and backward passes of a
    batch with this transformer model. Default approximation neglects the quadratic dependency on the number of
    tokens (valid if `12 * d_model << sequence_length`) as laid out in [this
    paper](https://huggingface.co/papers/2001.08361) section 2.1. Should be overridden for transformers with parameter
    re-use e.g. Albert or Universal Transformers, or if doing long-range modeling with very high sequence lengths.

Args:
        batch_size (`int`):
            The batch size for the forward pass.

        sequence_length (`int`):
            The number of tokens in each line of the batch.

        exclude_embeddings (`bool`, *optional*, defaults to `True`):
            Whether or not to count embedding and softmax operations.

Returns:
        `int`: The number of floating-point operations.


def execute(self):
    for name_context, spec in zip(self._map_task.name_contexts, self._map_task.operations):
        op = create_operation(name_context, spec, self._counter_factory, None, self._state_sampler, test_shuffle_source=self._test_shuffle_source, test_shuffle_sink=self._test_shuffle_sink)
        self._ops.append(op)
        if hasattr(op.spec, 'input'):
            producer, output_index = op.spec.input
            self._ops[producer].add_receiver(op, output_index)
        if hasattr(op.spec, 'inputs'):
            for producer, output_index in op.spec.inputs:
                self._ops[producer].add_receiver(op, output_index)
    for ix, op in reversed(list(enumerate(self._ops))):
        _LOGGER.debug('Starting op %d %s', ix, op)
        op.start()
    for op in self._ops:
        op.finish()

Executes all the operation_specs.Worker* instructions in a map task.

    We update the map_task with the execution status, expressed as counters.

Raises:
        RuntimeError: if we find more than on read instruction in task spec.
        TypeError: if the spec parameter is not an instance of the recognized
        operation_specs.Worker* classes.


def convert_n_to_tensor_or_composite(values, dtype=None, name=None) -> list[Union[EagerTensor, SymbolicTensor, composite_tensor.CompositeTensor, type(None)]]:
    return internal_convert_n_to_tensor_or_composite(values=values, dtype=dtype, name=name, as_ref=False)

Converts `values` to a list of `Output` or `CompositeTensor` objects.

    Any `CompositeTensor` objects in `values` are returned unmodified.

Args:
        values: A list of `None`, `CompositeTensor``, or objects that can be
        consumed by `convert_to_tensor()`.
        dtype: (Optional.) The required `DType` of the returned `Tensor`s or
        `CompositeTensor`s.
        name: (Optional.) A name prefix to used when a new `Tensor` is created, in
        which case element `i` will be given the name `name + '_' + i`.

Returns:
        A list of `Tensor` and/or `CompositeTensor` objects.

Raises:
        TypeError: If no conversion function is registered for an element in
        `values`.
        RuntimeError: If a registered conversion function returns an invalid
        value.