def _prepare_4d_causal_attention_mask_with_cache_position(attention_mask: torch.Tensor, sequence_length: int, target_length: int, dtype: torch.dtype, cache_position: torch.Tensor, batch_size: int, config: MoshiDepthConfig, past_key_values: Cache):
    if attention_mask is not None and attention_mask.dim() == 4:
        causal_mask = attention_mask
    else:
        min_dtype = torch.finfo(dtype).min
        causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device)
        diagonal_attend_mask = torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1)
        text_config = config.get_text_config()
        if getattr(text_config, 'use_sliding_window', True) and text_config.sliding_window is not None:
            if not isinstance(past_key_values, SlidingWindowCache) or sequence_length > target_length:
                sliding_attend_mask = torch.arange(target_length, device=cache_position.device) <= cache_position.reshape(-1, 1) - text_config.sliding_window
                diagonal_attend_mask.bitwise_or_(sliding_attend_mask)
        causal_mask *= diagonal_attend_mask
        causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
        if attention_mask is not None:
            causal_mask = causal_mask.clone()
            if attention_mask.shape[-1] > target_length:
                attention_mask = attention_mask[:, :target_length]
            mask_length = attention_mask.shape[-1]
            padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(causal_mask.device)
            padding_mask = padding_mask == 0
            causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(padding_mask, min_dtype)
    return causal_mask

Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
    `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.

Args:
        attention_mask (`torch.Tensor`):
            A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`.
        sequence_length (`int`):
            The sequence length being processed.
        target_length (`int`):
            The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet.
        dtype (`torch.dtype`):
            The dtype to use for the 4D attention mask.
        cache_position (`torch.Tensor`):
            Indices depicting the position of the input sequence tokens in the sequence.
        batch_size (`torch.Tensor`):
            Batch size.
        config (`MoshiDepthConfig`):
            The model's configuration class
        past_key_values (`Cache`):
            The cache class that is being used currently to generate


def forward(self, hidden_features):
    hidden_features = hidden_features.transpose(-1, -2)
    if self.head_aggregation == 'use_last':
        hidden_features = hidden_features[..., -1]
    elif self.head_aggregation == 'max_pool':
        hidden_features = hidden_features.max(dim=-1).values
    elif self.head_aggregation == 'avg_pool':
        hidden_features = hidden_features.mean(dim=-1)
    if self.flatten:
        hidden_features = self.flatten(hidden_features)
    hidden_features = self.dropout(hidden_features)
    hidden_features = self.projection(hidden_features)
    if self.distribution_output is None and self.output_range is not None:
        hidden_features = torch.sigmoid(hidden_features) * (self.output_range[1] - self.output_range[0]) + self.output_range[0]
    return hidden_features

Args:
    hidden_features (`torch.Tensor` of shape `(batch_size x num_patch x d_model)` in `flatten` mode
        or `(batch_size x n_vars x num_patch x d_model)` in `common_channel`/`mix_channel` mode.): Input hidden
        features.

Returns:
    `torch.Tensor` of shape `(batch_size x num_targets)`.


def __init__(self, content: str, type: Union[str, language_v1.Document.Type]='PLAIN_TEXT', language_hint: Optional[str]=None, encoding: Optional[str]='UTF8', from_gcs: bool=False):
    self.content = content
    self.type = type
    self.encoding = encoding
    self.language_hint = language_hint
    self.from_gcs = from_gcs

Represents the input to :class:`AnnotateText` transform.

Args:
        content (str): The content of the input or the Google Cloud Storage URI
        where the file is stored.
        type (`Union[str, google.cloud.language_v1.Document.Type]`): Text type.
        Possible values are `HTML`, `PLAIN_TEXT`. The default value is
        `PLAIN_TEXT`.
        language_hint (`Optional[str]`): The language of the text. If not specified,
        language will be automatically detected. Values should conform to
        ISO-639-1 standard.
        encoding (`Optional[str]`): Text encoding. Possible values are: `NONE`,
        `UTF8`, `UTF16`, `UTF32`. The default value is `UTF8`.
        from_gcs (bool): Whether the content should be interpret as a Google Cloud
        Storage URI. The default value is :data:`False`.


def __init__(self, name, default_name=None, values=None) -> None:
    self._name_scope = name_scope(name, default_name, values, skip_on_eager=False)
    self._name = default_name if name is None else name

Initialize the context manager.

Args:
        name: The name argument that is passed to the op function.
        default_name: The default name to use if the `name` argument is `None`.
        values: The list of `Tensor` arguments that are passed to the op function.

Raises:
        TypeError: if `default_name` is passed in but not a string.


def repeat_ranges(params, splits, repeats):
    splits_checks = [check_ops.assert_non_negative(splits, message="Input argument 'splits' must be non-negative"), check_ops.assert_integer(splits, message=f"Input argument 'splits' must be integer, but got {splits.dtype} instead")]
    repeats_checks = [check_ops.assert_non_negative(repeats, message="Input argument 'repeats' must be non-negative"), check_ops.assert_integer(repeats, message=f"Input argument 'repeats' must be integer, but got {repeats.dtype} instead")]
    splits = control_flow_ops.with_dependencies(splits_checks, splits)
    repeats = control_flow_ops.with_dependencies(repeats_checks, repeats)
    if repeats.shape.ndims != 0:
        repeated_starts = repeat(splits[:-1], repeats, axis=0)
        repeated_limits = repeat(splits[1:], repeats, axis=0)
    else:
        repeated_splits = repeat(splits, repeats, axis=0)
        n_splits = array_ops.shape(repeated_splits, out_type=repeats.dtype)[0]
        repeated_starts = repeated_splits[:n_splits - repeats]
        repeated_limits = repeated_splits[repeats:]
    one = array_ops.ones((), repeated_starts.dtype)
    offsets = gen_ragged_math_ops.ragged_range(repeated_starts, repeated_limits, one)
    return array_ops.gather(params, offsets.rt_dense_values)

Repeats each range of `params` (as specified by `splits`) `repeats` times.

    Let the `i`th range of `params` be defined as
    `params[splits[i]:splits[i + 1]]`.  Then this function returns a tensor
    containing range 0 repeated `repeats[0]` times, followed by range 1 repeated
    `repeats[1]`, ..., followed by the last range repeated `repeats[-1]` times.

Args:
        params: The `Tensor` whose values should be repeated.
        splits: A splits tensor indicating the ranges of `params` that should be
        repeated. Elements should be non-negative integers.
        repeats: The number of times each range should be repeated. Supports
        broadcasting from a scalar value. Elements should be non-negative
        integers.

Returns:
        A `Tensor` with the same rank and type as `params`.

    #### Example:

    >>> print(repeat_ranges(
    ...     params=tf.constant(['a', 'b', 'c']),
    ...     splits=tf.constant([0, 2, 3]),
    ...     repeats=tf.constant(3)))
    tf.Tensor([b'a' b'b' b'a' b'b' b'a' b'b' b'c' b'c' b'c'],
        shape=(9,), dtype=string)


def create_model(fingerprint_input, model_settings, model_architecture, is_training, runtime_settings=None):
    if model_architecture == 'single_fc':
        return create_single_fc_model(fingerprint_input, model_settings, is_training)
    elif model_architecture == 'conv':
        return create_conv_model(fingerprint_input, model_settings, is_training)
    elif model_architecture == 'low_latency_conv':
        return create_low_latency_conv_model(fingerprint_input, model_settings, is_training)
    elif model_architecture == 'low_latency_svdf':
        return create_low_latency_svdf_model(fingerprint_input, model_settings, is_training, runtime_settings)
    elif model_architecture == 'tiny_conv':
        return create_tiny_conv_model(fingerprint_input, model_settings, is_training)
    elif model_architecture == 'tiny_embedding_conv':
        return create_tiny_embedding_conv_model(fingerprint_input, model_settings, is_training)
    else:
        raise Exception('model_architecture argument "' + model_architecture + '" not recognized, should be one of "single_fc", "conv",' + ' "low_latency_conv, "low_latency_svdf",' + ' "tiny_conv", or "tiny_embedding_conv"')

Builds a model of the requested architecture compatible with the settings.

    There are many possible ways of deriving predictions from a spectrogram
    input, so this function provides an abstract interface for creating different
    kinds of models in a black-box way. You need to pass in a TensorFlow node as
    the 'fingerprint' input, and this should output a batch of 1D features that
    describe the audio. Typically this will be derived from a spectrogram that's
    been run through an MFCC, but in theory it can be any feature vector of the
    size specified in model_settings['fingerprint_size'].

    The function will build the graph it needs in the current TensorFlow graph,
    and return the tensorflow output that will contain the 'logits' input to the
    softmax prediction process. If training flag is on, it will also return a
    placeholder node that can be used to control the dropout amount.

    See the implementations below for the possible model architectures that can be
    requested.

Args:
        fingerprint_input: TensorFlow node that will output audio feature vectors.
        model_settings: Dictionary of information about the model.
        model_architecture: String specifying which kind of model to create.
        is_training: Whether the model is going to be used for training.
        runtime_settings: Dictionary of information about the runtime.

Returns:
        TensorFlow node outputting logits results, and optionally a dropout
        placeholder.

Raises:
        Exception: If the architecture type isn't recognized.


def _preprocess_conv3d_input(x, data_format):
    tf_data_format = 'NDHWC'
    if data_format == 'channels_first':
        if not _has_nchw_support():
            x = array_ops.transpose(x, (0, 2, 3, 4, 1))
        else:
            tf_data_format = 'NCDHW'
    return (x, tf_data_format)

Transpose and cast the input before the conv3d.

Args:
        x: input tensor.
        data_format: string, `"channels_last"` or `"channels_first"`.

Returns:
        A tensor.


def compact(self, accumulator, *args, **kwargs):
    return accumulator

Optionally returns a more compact representation of the accumulator.

    This is called before an accumulator is sent across the wire, and can
    be useful in cases where values are buffered or otherwise lazily
    kept unprocessed when added to the accumulator.  Should return an
    equivalent, though possibly modified, accumulator.

    By default returns the accumulator unmodified.

Args:
        accumulator: the current accumulator
        *args: Additional arguments and side inputs.
        **kwargs: Additional arguments and side inputs.


def close(self, cancel_pending_enqueues=False, name=None):
    if name is None:
        name = '%s_BarrierClose' % self._name
    return gen_data_flow_ops.barrier_close(self._barrier_ref, cancel_pending_enqueues=cancel_pending_enqueues, name=name)

Closes this barrier.

    This operation signals that no more new key values will be inserted in the
    given barrier. Subsequent InsertMany operations with new keys will fail.
    InsertMany operations that just complement already existing keys with other
    components, will continue to succeed. Subsequent TakeMany operations will
    continue to succeed if sufficient elements remain in the barrier. Subsequent
    TakeMany operations that would block will fail immediately.

    If `cancel_pending_enqueues` is `True`, all pending requests to the
    underlying queue will also be canceled, and completing of already
    started values is also not acceptable anymore.

Args:
        cancel_pending_enqueues: (Optional.) A boolean, defaulting to
        `False` (described above).
        name: Optional name for the op.

Returns:
        The operation that closes the barrier.


def WriteExecution(self, execution):
    debug_event = debug_event_pb2.DebugEvent(execution=execution)
    self._EnsureTimestampAdded(debug_event)
    _pywrap_debug_events_writer.WriteExecution(self._dump_root, debug_event)

Write a Execution proto with the writer.

Args:
        execution: An Execution proto, describing a TensorFlow op or graph
        execution event.


def _get_cuda_compute_capabilities_or_die() -> list[str]:
    try:
        nvidia_smi = _find_executable_or_die('nvidia-smi')
        nvidia_smi_proc = subprocess.run([nvidia_smi, '--query-gpu=compute_cap', '--format=csv,noheader'], capture_output=True, check=True, text=True)
        capabilities = sorted(set(nvidia_smi_proc.stdout.strip().split('\n')))
        logging.info('Found CUDA compute capabilities: %s', capabilities)
        return capabilities
    except (RuntimeError, subprocess.CalledProcessError) as e:
        logging.info('Could not find nvidia-smi, or nvidia-smi command failed. Please pass capabilities directly using --cuda_compute_capabilities.')
        raise e

Finds compute capabilities via nvidia-smi or rasies exception.

Returns:
        list of unique, sorted strings representing compute capabilities:

Raises:
        RuntimeError: if path to nvidia-smi couldn't be found.
        subprocess.CalledProcessError: if nvidia-smi process failed.


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

Returns boolean features, `True` in the NaN elements of the
    [`EventSet`][temporian.EventSet].

    Note that for `int` and `bool` this will always be `False` since those types
    don't support NaNs. It only makes actual sense to use on `float` (or
    `tp.float32`) features.

    See also `evset.notnan()`.

Example:
        ```python
        >>> a = tp.event_set(
        ...     timestamps=[1, 2, 3],
        ...     features={"M":[np.nan, 5., np.nan], "N":  [-1, 0, 5]},
        ... )
        >>> b = a.isnan()
        >>> b
        indexes: ...
                'M': [ True False True]
                'N': [False False False]
        ...

        >>> # Count nans
        >>> b["M"].cast(int).cumsum()
        indexes: ...
                timestamps: [1. 2. 3.]
                'M': [1 1 2]
        ...

        ```

Returns:
        EventSet with boolean features.


def stack_inputs(self, stack_indices=None, tile_variants=False):
    if stack_indices is None:
        stack_indices = range(len(self._inputs))
    length = self.pfor.loop_len_vector
    for i in stack_indices:
        inp = self._inputs[i]
        is_variant = inp.t.dtype == dtypes.variant
        if not inp.is_stacked:
            self._inputs[i] = _stack(inp.t, length)
            if tile_variants and is_variant:
                self._inputs[i] = wrap(_tile_variant_with_length(self._inputs[i].t, length), True)
        elif not tile_variants and is_variant:
            self._inputs[i] = wrap(_untile_variant(self._inputs[i].t), True)

Stacks unstacked inputs at `stack_indices`.

Args:
        stack_indices: indices of inputs at which stacking is done. If None,
        stacking is done at all indices.
        tile_variants: If True, affected indices which have a variant dtype will
        be tiled after this operation to match the expected shape of a
        vectorized tensor. Variants generally need to be un-tiled when they are
        inputs to operations and tiled when returned.


def LSTMCell(weights, m_prev, c_prev, x, pad):
    xm = array_ops.concat([x, m_prev], 1)
    xmw = math_ops.matmul(xm, weights)
    in_value, in_gate, forget_gate, out_gate = array_ops.split(value=xmw, num_or_size_splits=4, axis=1)
    in_value = math_ops.tanh(in_value)
    in_gate = math_ops.sigmoid(in_gate)
    forget_gate = math_ops.sigmoid(forget_gate)
    out_gate = math_ops.sigmoid(out_gate)
    c_next = Clip(Clip(forget_gate * c_prev) + Clip(in_gate * in_value))
    m_next = Clip(out_gate * c_next)
    c_next = c_prev * pad + c_next * (1.0 - pad)
    m_next = m_prev * pad + m_next * (1.0 - pad)
    return (m_next, c_next)

Unrolls a single LSTM cell with clipped activations forward by one step.

Args:
        weights: Weight matrix with shape LSTMCellWeightsShape.
        m_prev: Previous m states with shape [batch_size, num_nodes].
        c_prev: Previous c states with shape [batch_size, num_nodes].
        x: Input with shape [batch_size, num_inputs].
        pad: Padding with shape [batch_size, 1].  Each padding value is either
        0 or 1, where 1 indicates padding; i.e. the input is shorter than the
        sequence length, and the (m, c) states should simply be passed through
        from the previous states.

Returns:
        The next (m, c) states, each with shape [batch_size, num_nodes].


def resize(self, image: np.ndarray, size: Dict[str, int], resample: PILImageResampling=PILImageResampling.BICUBIC, data_format: Optional[Union[str, ChannelDimension]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None, **kwargs) -> np.ndarray:
    size = get_size_dict(size, default_to_square=True, param_name='size')
    if 'height' not in size or 'width' not in size:
        raise ValueError(f'The `size` argument must contain `height` and `width` keys. Got {size.keys()}')
    return resize(image, size=(size['height'], size['width']), resample=resample, data_format=data_format, input_data_format=input_data_format, **kwargs)

Resize an image to (size["height"], size["width"]).

Args:
        image (`np.ndarray`):
            Image to resize.
        size (`Dict[str, int]`):
            Size of the output image.
        resample (`PILImageResampling`, *optional*, defaults to `PIL.Image.BICUBIC`):
            Resampling filter to use when resiizing the image.
        data_format (`str` or `ChannelDimension`, *optional*):
            The channel dimension format of the image. If not provided, it will be the same as the input image.
        input_data_format (`str` or `ChannelDimension`, *optional*):
            The channel dimension format of the input image. If not provided, it will be inferred.


def sym_descendants(self, where: Optional[Callable[[Any], bool]]=None, option: DescendantQueryOption=DescendantQueryOption.ALL, include_self: bool=False) -> List[Any]:
    descendants = []
    where = where or (lambda x: True)

    def visit(k, v, p):
        del k, p
        if not where(v):
            return TraverseAction.ENTER
        if not include_self and self is v:
            return TraverseAction.ENTER
        if option == DescendantQueryOption.IMMEDIATE:
            descendants.append(v)
            return TraverseAction.CONTINUE
        leaf_descendants = []
        if isinstance(v, Symbolic):
            leaf_descendants = v.sym_descendants(where, option)
        if option is DescendantQueryOption.ALL or not leaf_descendants:
            descendants.append(v)
        descendants.extend(leaf_descendants)
        return TraverseAction.CONTINUE
    traverse(self, visit)
    return descendants

Returns all descendants of specific classes.

Args:
        where: Optional callable object as the filter of descendants to return.
        option: Descendant query options, indicating whether all matched,
        immediate matched or only the matched leaf nodes will be returned.
        include_self: If True, `self` will be included in the query, otherwise
        only strict descendants are included.

Returns:
        A list of objects that match the descendant_cls.


def _init_from_proto(self, context_def, import_scope=None):
    assert isinstance(context_def, control_flow_pb2.WhileContextDef)
    g = ops.get_default_graph()
    self._name = ops.prepend_name_scope(context_def.context_name, import_scope)
    if context_def.maximum_iterations_name:
        self._maximum_iterations = g.as_graph_element(ops.prepend_name_scope(context_def.maximum_iterations_name, import_scope))
    else:
        self._maximum_iterations = None
    self._parallel_iterations = context_def.parallel_iterations
    self._back_prop = context_def.back_prop
    self._swap_memory = context_def.swap_memory
    self._pivot_for_pred = g.as_graph_element(ops.prepend_name_scope(context_def.pivot_for_pred_name, import_scope))
    self._pivot_for_body = g.as_graph_element(ops.prepend_name_scope(context_def.pivot_for_body_name, import_scope))
    self._pivot = g.as_graph_element(ops.prepend_name_scope(context_def.pivot_name, import_scope))
    self._loop_exits = [g.as_graph_element(ops.prepend_name_scope(exit_name, import_scope)) for exit_name in context_def.loop_exit_names]
    self._loop_enters = [g.as_graph_element(ops.prepend_name_scope(enter_name, import_scope)) for enter_name in context_def.loop_enter_names]
    super(WhileContext, self).__init__(values_def=context_def.values_def, import_scope=import_scope)
    if import_scope:
        for tensor_name in self._values:
            op = g.as_graph_element(tensor_name).op
            if util.IsLoopEnter(op):
                op._set_attr('frame_name', attr_value_pb2.AttrValue(s=compat.as_bytes(self.name)))
    self._graph = ops.get_default_graph()

Creates a new `WhileContext` from protocol buffer.

Args:
        context_def: `WhileContextDef` protocol buffer.
        import_scope: Optional `string`. Name scope to add.


def __init__(self, config: GPT2Config):
    super().__init__()
    self.summary_type = getattr(config, 'summary_type', 'last')
    if self.summary_type == 'attn':
        raise NotImplementedError
    self.summary = nn.Identity()
    if hasattr(config, 'summary_use_proj') and config.summary_use_proj:
        if hasattr(config, 'summary_proj_to_labels') and config.summary_proj_to_labels and (config.num_labels > 0):
            num_classes = config.num_labels
        else:
            num_classes = config.hidden_size
        self.summary = nn.Linear(config.hidden_size, num_classes)
    activation_string = getattr(config, 'summary_activation', None)
    self.activation: Callable = get_activation(activation_string) if activation_string else nn.Identity()
    self.first_dropout = nn.Identity()
    if hasattr(config, 'summary_first_dropout') and config.summary_first_dropout > 0:
        self.first_dropout = nn.Dropout(config.summary_first_dropout)
    self.last_dropout = nn.Identity()
    if hasattr(config, 'summary_last_dropout') and config.summary_last_dropout > 0:
        self.last_dropout = nn.Dropout(config.summary_last_dropout)

Compute a single vector summary of a sequence hidden states.

Args:
        config ([`GPT2Config`]):
            The config used by the model. Relevant arguments in the config class of the model are (refer to the actual
            config class of your model for the default values it uses):

            - **summary_type** (`str`) -- The method to use to make this summary. Accepted values are:

                - `"last"` -- Take the last token hidden state (like XLNet)
                - `"first"` -- Take the first token hidden state (like Bert)
                - `"mean"` -- Take the mean of all tokens hidden states
                - `"cls_index"` -- Supply a Tensor of classification token position (GPT/GPT-2)
                - `"attn"` -- Not implemented now, use multi-head attention

            - **summary_use_proj** (`bool`) -- Add a projection after the vector extraction.
            - **summary_proj_to_labels** (`bool`) -- If `True`, the projection outputs to `config.num_labels` classes
            (otherwise to `config.hidden_size`).
            - **summary_activation** (`Optional[str]`) -- Set to `"tanh"` to add a tanh activation to the output,
            another string or `None` will add no activation.
            - **summary_first_dropout** (`float`) -- Optional dropout probability before the projection and activation.
            - **summary_last_dropout** (`float`)-- Optional dropout probability after the projection and activation.


def _eventual_warn_about_too_long_sequence(self, ids: List[int], max_length: Optional[int], verbose: bool):
    if max_length is None and len(ids) > self.model_max_length and verbose and (self.model_max_length != 0):
        if not self.deprecation_warnings.get('sequence-length-is-longer-than-the-specified-maximum', False):
            logger.warning(f'Token indices sequence length is longer than the specified maximum sequence length for this model ({len(ids)} > {self.model_max_length}). Running this sequence through the model will result in indexing errors')
        self.deprecation_warnings['sequence-length-is-longer-than-the-specified-maximum'] = True

Depending on the input and internal state we might trigger a warning about a sequence that is too long for its
    corresponding model

Args:
        ids (`List[str]`): The ids produced by the tokenization
        max_length (`int`, *optional*): The max_length desired (does not trigger a warning if it is set)
        verbose (`bool`): Whether or not to print more information and warnings.


def _get_index(self, data: _instance_base.Instance | ConcreteValue) -> int | None:
    if isinstance(data, ConcreteValue):
        return self.ctx.convert.value_to_constant(data, (int, type(None)))
    elif isinstance(data, _instance_base.Instance):
        if data.cls != self.ctx.convert.int_type:
            raise abstract_utils.ConversionError()
        else:
            return None
    else:
        raise abstract_utils.ConversionError()

Helper function for getslice_slot that extracts int or None from data.

    If data is an Instance of int, None is returned.

Args:
        data: The object to extract from. Usually a ConcreteValue or an Instance.

Returns:
        The value (an int or None) of the index.

Raises:
        abstract_utils.ConversionError: If the data could not be converted.


def sample_points_using_uncertainty(self, logits: torch.Tensor, uncertainty_function, num_points: int, oversample_ratio: int, importance_sample_ratio: float) -> torch.Tensor:
    num_boxes = logits.shape[0]
    num_points_sampled = int(num_points * oversample_ratio)
    point_coordinates = torch.rand(num_boxes, num_points_sampled, 2, device=logits.device)
    point_logits = sample_point(logits, point_coordinates, align_corners=False)
    point_uncertainties = uncertainty_function(point_logits)
    num_uncertain_points = int(importance_sample_ratio * num_points)
    num_random_points = num_points - num_uncertain_points
    idx = torch.topk(point_uncertainties[:, 0, :], k=num_uncertain_points, dim=1)[1]
    shift = num_points_sampled * torch.arange(num_boxes, dtype=torch.long, device=logits.device)
    idx += shift[:, None]
    point_coordinates = point_coordinates.view(-1, 2)[idx.view(-1), :].view(num_boxes, num_uncertain_points, 2)
    if num_random_points > 0:
        point_coordinates = torch.cat([point_coordinates, torch.rand(num_boxes, num_random_points, 2, device=logits.device)], dim=1)
    return point_coordinates

This function is meant for sampling points in [0, 1] * [0, 1] coordinate space based on their uncertainty. The
    uncertainty is calculated for each point using the passed `uncertainty function` that takes points logit
    prediction as input.

Args:
        logits (`float`):
            Logit predictions for P points.
        uncertainty_function:
            A function that takes logit predictions for P points and returns their uncertainties.
        num_points (`int`):
            The number of points P to sample.
        oversample_ratio (`int`):
            Oversampling parameter.
        importance_sample_ratio (`float`):
            Ratio of points that are sampled via importance sampling.

Returns:
        point_coordinates (`torch.Tensor`):
            Coordinates for P sampled points.


def start_server(self):
    persists_shell_cmd = self._get_persisting_command()
    self.log.debug('Snippet server for package %s is using protocol %d.%d', self.package, _PROTOCOL_MAJOR_VERSION, _PROTOCOL_MINOR_VERSION)
    option_str = self._get_instrument_options_str()
    cmd = _LAUNCH_CMD.format(shell_cmd=persists_shell_cmd, user=self._get_user_command_string(), snippet_package=self.package, instrument_options=option_str)
    self._proc = self._run_adb_cmd(cmd)
    self._server_start_stdout = []
    line = self._read_protocol_line()
    match = re.match('^SNIPPET START, PROTOCOL ([0-9]+) ([0-9]+)$', line)
    if not match or int(match.group(1)) != _PROTOCOL_MAJOR_VERSION:
        raise errors.ServerStartProtocolError(self._device, line)
    line = self._read_protocol_line()
    match = re.match('^SNIPPET SERVING, PORT ([0-9]+)$', line)
    if not match:
        message = _SNIPPET_SERVER_START_ERROR_DEBUG_TIP.format(instrumentation_result=line, server_start_stdout='\n'.join(self._server_start_stdout))
        raise errors.ServerStartProtocolError(self._device, message)
    self.device_port = int(match.group(1))

Starts the server on the remote device.

    This function starts the snippet server with adb command, checks the
    protocol version of the server, parses device port from the server
    output and sets it to self.device_port.

Raises:
        errors.ServerStartProtocolError: if the protocol reported by the server
        startup process is unknown.
        errors.ServerStartError: if failed to start the server or process the
        server output.


def _VerifyValues(self, tensor_in_sizes, filter_in_sizes, stride, padding, data_type, data_format='NHWC'):
    total_size_1 = 1
    total_size_2 = 1
    for s in tensor_in_sizes:
        total_size_1 *= s
    for s in filter_in_sizes:
        total_size_2 *= s
    x1 = np.array([f * 1.0 for f in range(1, total_size_1 + 1)], dtype=data_type).reshape(tensor_in_sizes)
    x2 = np.array([f * 1.0 for f in range(1, total_size_2 + 1)], dtype=data_type).reshape(filter_in_sizes)
    with self.session() as sess:
        if data_type == np.float32:
            tolerance = 0.0001
        else:
            self.assertEqual(data_type, np.float64)
            tolerance = 1e-08
        t1 = array_ops.placeholder(shape=tensor_in_sizes, dtype=data_type)
        t2 = array_ops.placeholder(shape=filter_in_sizes, dtype=data_type)
        native_t1 = t1
        strides = [1, stride, stride, 1]
        if data_format == 'NCHW':
            native_t1 = array_ops.transpose(t1, [0, 3, 1, 2])
            strides = [1, 1, stride, stride]
        with self.test_scope():
            conv_native = nn_ops.depthwise_conv2d_native(native_t1, t2, strides=strides, data_format=data_format, padding=padding)
        if data_format == 'NCHW':
            conv_native = array_ops.transpose(conv_native, [0, 2, 3, 1])
        with ops.device('CPU'):
            conv_interface = ReferenceDepthwiseConv2D(t1, t2, strides=[1, stride, stride, 1], padding=padding)
        native_result = sess.run(conv_native, {t1: x1, t2: x2})
        interface_result = sess.run(conv_interface, {t1: x1, t2: x2})
    print('data_type:', data_type, 'max diff = ', np.amax(np.absolute(native_result - interface_result)))
    self.assertAllClose(np.ravel(native_result), np.ravel(interface_result), rtol=tolerance)

Verifies the output values of the convolution function.

Args:
        tensor_in_sizes: Input tensor dimensions in
        [batch, input_rows, input_cols, input_depth].
        filter_in_sizes: Filter tensor dimensions in
        [filter_rows, filter_cols, input_depth, depth_multiplier].
        stride: Stride.
        padding: Padding type.
        data_type: The data type to use.
        data_format: The data_format of the input. "NHWC" or "NCHW".


def _composition_must_be_self_adjoint(operators):
    if len(operators) == 1 and operators[0].is_self_adjoint:
        return True
    if linear_operator_util.is_aat_form(operators):
        return True
    return False

Runs some checks to see if composition operators must be SA.

Args:
        operators: List of LinearOperators.

Returns:
        True if the composition must be SA. False if it is not SA OR if we did not
        determine whether the composition is SA.


def input(self):
    if not self._inbound_nodes:
        raise AttributeError('Layer ' + self.name + ' is not connected, no input to return.')
    return self._get_node_attribute_at_index(0, 'input_tensors', 'input')

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 build_alibi_tensor(attention_mask: jnp.ndarray, num_heads: int, dtype: Optional[jnp.dtype]=jnp.float32):
    batch_size, seq_length = attention_mask.shape
    closest_power_of_2 = 2 ** math.floor(math.log2(num_heads))
    base = jnp.array(2 ** (-2 ** (-(math.log2(closest_power_of_2) - 3))), dtype=jnp.float32)
    powers = jnp.arange(1, 1 + closest_power_of_2, dtype=jnp.float32)
    slopes = jax.lax.pow(base, powers)
    if closest_power_of_2 != num_heads:
        extra_base = jnp.array(2 ** (-2 ** (-(math.log2(2 * closest_power_of_2) - 3))), dtype=jnp.float32)
        num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2)
        extra_powers = jnp.arange(1, 1 + 2 * num_remaining_heads, 2, dtype=jnp.float32)
        slopes = jnp.cat([slopes, jax.lax.pow(extra_base, extra_powers)], axis=0)
    arange_tensor = ((attention_mask.cumsum(axis=-1) - 1) * attention_mask)[:, None, :]
    alibi = slopes[..., None] * arange_tensor
    alibi = jnp.expand_dims(alibi, axis=2)
    return jnp.asarray(alibi, dtype)

Flax implementation of the BLOOM Alibi tensor. BLOOM Alibi tensor is not causal as the original paper mentions, it
    relies on a translation invariance of softmax for quick implementation: with l being a tensor, and a fixed value
    `softmax(l+a) = softmax(l)`. Based on
    https://github.com/ofirpress/attention_with_linear_biases/blob/a35aaca144e0eb6b789dfcb46784c4b8e31b7983/fairseq/models/transformer.py#L742
    Link to paper: https://huggingface.co/papers/2108.12409

Args:
        attention_mask (`jnp.ndarray`):
            Token-wise attention mask, this should be of shape `(batch_size, max_seq_len)`.
        num_heads (`int`):
            Number of attention heads.
        dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`):
            The data type (dtype) of the output tensor.

    Returns: Alibi tensor of shape `(batch_size * num_heads, 1, max_seq_len)`.


def set_floatx(value):
    global _FLOATX
    accepted_dtypes = {'bfloat16', 'float16', 'float32', 'float64'}
    if value not in accepted_dtypes:
        raise ValueError(f'Unknown `floatx` value: {value}. Expected one of {accepted_dtypes}')
    _FLOATX = str(value)

Set the default float dtype.

    Note: It is not recommended to set this to `"float16"` for training,
    as this will likely cause numeric stability issues.
    Instead, mixed precision, which leverages
    a mix of `float16` and `float32`. It can be configured by calling
    `keras.mixed_precision.set_dtype_policy('mixed_float16')`.

Args:
        value: String; `'bfloat16'`, `'float16'`, `'float32'`, or `'float64'`.

Example:
    >>> keras.config.floatx()
    'float32'

    >>> keras.config.set_floatx('float64')
    >>> keras.config.floatx()
    'float64'

    >>> # Set it back to float32
    >>> keras.config.set_floatx('float32')

Raises:
        ValueError: In case of invalid value.


def is_flaky(max_attempts: int=5, wait_before_retry: Optional[float]=None, description: Optional[str]=None):

    def decorator(test_func_ref):

        @functools.wraps(test_func_ref)
        def wrapper(*args, **kwargs):
            retry_count = 1
            while retry_count < max_attempts:
                try:
                    return test_func_ref(*args, **kwargs)
                except Exception as err:
                    logger.error(f'Test failed with {err} at try {retry_count}/{max_attempts}.')
                    if wait_before_retry is not None:
                        time.sleep(wait_before_retry)
                    retry_count += 1
            return test_func_ref(*args, **kwargs)
        return unittest.skipUnless(_run_flaky_tests, 'test is flaky')(wrapper)
    return decorator

To decorate flaky tests. They will be retried on failures.

    Please note that our push tests use `pytest-rerunfailures`, which prompts the CI to rerun certain types of
    failed tests. More specifically, if the test exception contains any substring in `FLAKY_TEST_FAILURE_PATTERNS`
    (in `.circleci/create_circleci_config.py`), it will be rerun. If you find a recurrent pattern of failures,
    expand `FLAKY_TEST_FAILURE_PATTERNS` in our CI configuration instead of using `is_flaky`.

Args:
        max_attempts (`int`, *optional*, defaults to 5):
            The maximum number of attempts to retry the flaky test.
        wait_before_retry (`float`, *optional*):
            If provided, will wait that number of seconds before retrying the test.
        description (`str`, *optional*):
            A string to describe the situation (what / where / why is flaky, link to GH issue/PR comments, errors,
            etc.)


def forward(self, **super_kwargs) -> CausalLMOutputWithPast:
    return super().forward(**super_kwargs)

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

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

    >>> model = BitNetForCausalLM.from_pretrained("microsoft/bitnet-b1.58-2B-4T")
    >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/bitnet-b1.58-2B-4T")

    >>> prompt = f'<|begin_of_text|>User: Hey, are you conscious? Can you talk to me?<|eot_id|>Assistant: '
    >>> inputs = tokenizer(prompt, return_tensors="pt")

    >>> # Generate
    >>> generate_ids = model.generate(inputs.input_ids, max_length=100)
    >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
    "User: Hey, are you conscious? Can you talk to me?Assistant: No, I'm not conscious. I'm an artificial intelligence designed to assist with information and tasks. How can I help you today?"
    ```


def run_benchmark(self, dataset, num_elements, iters=1, warmup=True, apply_default_optimizations=False, session_config=None):
    options = options_lib.Options()
    options.experimental_optimization.apply_default_optimizations = apply_default_optimizations
    dataset = dataset.with_options(options)
    dataset = dataset.skip(num_elements - 1)
    if context.executing_eagerly():
        median_duration = self._run_eager_benchmark(iterable=dataset, iters=iters, warmup=warmup)
        return median_duration / float(num_elements)
    iterator = dataset_ops.make_initializable_iterator(dataset)
    next_element = iterator.get_next()
    op = nest.flatten(next_element)[0].op
    median_duration = self._run_graph_benchmark(iterable=op, iters=iters, warmup=warmup, session_config=session_config, initializer=iterator.initializer)
    return median_duration / float(num_elements)

Benchmarks the dataset.

    Runs the dataset `iters` times. In each iteration, the benchmark measures
    the time it takes to go through `num_elements` elements of the dataset.

Args:
        dataset: Dataset to benchmark.
        num_elements: Number of dataset elements to iterate through each benchmark
        iteration.
        iters: Number of times to repeat the timing.
        warmup: If true, warms up the session caches by running an untimed run.
        apply_default_optimizations: Determines whether default optimizations
        should be applied.
        session_config: A ConfigProto protocol buffer with configuration options
        for the session. Applicable only for benchmarking in graph mode.

Returns:
        A float, representing the per-element wall time of the dataset in seconds.
        This is the median time (with respect to `iters`) it takes for the dataset
        to go through `num_elements` elements, divided by `num_elements.`


def assert_no_garbage_created(f: _F) -> _F:

    def decorator(self: 'TensorFlowTestCase', **kwargs):
        """Sets DEBUG_SAVEALL, runs the test, and checks for new garbage."""
        gc.disable()
        previous_debug_flags = gc.get_debug()
        gc.set_debug(gc.DEBUG_UNCOLLECTABLE)
        gc.collect()
        previous_garbage = len(gc.garbage)
        result = f(self, **kwargs)
        gc.collect()
        new_garbage = len(gc.garbage)
        if new_garbage > previous_garbage:
            for i, obj in enumerate(gc.garbage[previous_garbage:]):
                if getattr(obj, '__module__', '') == 'ast':
                    new_garbage -= 3
        if new_garbage > previous_garbage:
            logging.error("The decorated test created work for Python's garbage collector, likely due to a reference cycle. New objects in cycle(s):")
            for i, obj in enumerate(gc.garbage[previous_garbage:]):
                try:
                    logging.error('Object %d of %d', i, len(gc.garbage) - previous_garbage)

                    def _safe_object_str(obj) -> str:
                        return '<%s %d>' % (obj.__class__.__name__, id(obj))
                    logging.error('  Object type: %s', _safe_object_str(obj))
                    logging.error('  Referrer types: %s', ', '.join([_safe_object_str(ref) for ref in gc.get_referrers(obj)]))
                    logging.error('  Referent types: %s', ', '.join([_safe_object_str(ref) for ref in gc.get_referents(obj)]))
                    logging.error('  Object attribute names: %s', dir(obj))
                    logging.error('  Object __str__:')
                    logging.error(obj)
                    logging.error('  Object __repr__:')
                    logging.error(repr(obj))
                except Exception:
                    logging.error('(Exception while printing object)')
        if new_garbage > previous_garbage:
            for i in range(previous_garbage, new_garbage):
                if _find_reference_cycle(gc.garbage, i):
                    break
        self.assertEqual(previous_garbage, new_garbage)
        gc.set_debug(previous_debug_flags)
        gc.enable()
        return result
    return decorator

Test method decorator to assert that no garbage has been created.

    Note that this decorator sets DEBUG_SAVEALL, which in some Python interpreters
    cannot be un-set (i.e. will disable garbage collection for any other unit
    tests in the same file/shard).

Args:
        f: The function to decorate.

Returns:
        The decorated function.


def rgb_to_grayscale(images, data_format=None):
    if any_symbolic_tensors((images,)):
        return RGBToGrayscale(data_format=data_format).symbolic_call(images)
    return backend.image.rgb_to_grayscale(images, data_format=data_format)

Convert RGB images to grayscale.

    This function converts RGB images to grayscale images. It supports both
    3D and 4D tensors.

Args:
        images: Input image or batch of images. Must be 3D or 4D.
        data_format: A string specifying the data format of the input tensor.
            It can be either `"channels_last"` or `"channels_first"`.
            `"channels_last"` corresponds to inputs with shape
            `(batch, height, width, channels)`, while `"channels_first"`
            corresponds to inputs with shape `(batch, channels, height, width)`.
            If not specified, the value will default to
            `keras.config.image_data_format`.

Returns:
        Grayscale image or batch of grayscale images.

Example:
    >>> import numpy as np
    >>> from keras import ops
    >>> x = np.random.random((2, 4, 4, 3))
    >>> y = ops.image.rgb_to_grayscale(x)
    >>> y.shape
    (2, 4, 4, 1)

    >>> x = np.random.random((4, 4, 3)) # Single RGB image
    >>> y = ops.image.rgb_to_grayscale(x)
    >>> y.shape
    (4, 4, 1)

    >>> x = np.random.random((2, 3, 4, 4))
    >>> y = ops.image.rgb_to_grayscale(x, data_format="channels_first")
    >>> y.shape
    (2, 1, 4, 4)


def __init__(self, axis=-1, **kwargs):
    super().__init__(**kwargs)
    self.axis = axis
    self.supports_masking = True
    self._build_at_init()

Softmax activation layer.

    Formula:
    ``` python
    exp_x = exp(x - max(x))
    f(x) = exp_x / sum(exp_x)
    ```

Example:
    >>> softmax_layer = keras.layers.Softmax()
    >>> input = np.array([1.0, 2.0, 1.0])
    >>> result = softmax_layer(input)
    >>> result
    [0.21194157, 0.5761169, 0.21194157]

Args:
        axis: Integer, or list of Integers, axis along which the softmax
            normalization is applied.
        **kwargs: Base layer keyword arguments, such as `name` and `dtype`.

    Call arguments:
        inputs: The inputs (logits) to the softmax layer.
        mask: A boolean mask of the same shape as `inputs`. The mask
            specifies 1 to keep and 0 to mask. Defaults to `None`.

Returns:
        Softmaxed output with the same shape as `inputs`.


def sparse_subtract(x1, x2):
    if isinstance(x2, tf.SparseTensor):
        return tf.sparse.add(x1, tf.sparse.map_values(tf.negative, x2))
    else:
        return tf.sparse.add(x1, tf.negative(x2))

Subtraction for `tf.SparseTensor`s.

    Either `x1` or `x2` or both can be `tf.SparseTensor`s.

Args:
        x1: fist tensor to add.
        x2: second tensor to add.

Returns:
        The sum of `x1` and `x2`, which is a `tf.SparseTensor` if and only if
        both `x1` or `x2` are `tf.SparseTensor`s.


def node_info(self, args, screen_info=None):
    _ = screen_info
    parsed = self._arg_parsers['node_info'].parse_args(args)
    node_name, unused_slot = debug_graphs.parse_node_or_tensor_name(parsed.node_name)
    if not self._debug_dump.node_exists(node_name):
        output = cli_shared.error('There is no node named "%s" in the partition graphs' % node_name)
        _add_main_menu(output, node_name=None, enable_list_tensors=True, enable_node_info=False, enable_list_inputs=False, enable_list_outputs=False)
        return output
    lines = ['Node %s' % node_name]
    font_attr_segs = {0: [(len(lines[-1]) - len(node_name), len(lines[-1]), 'bold')]}
    lines.append('')
    lines.append('  Op: %s' % self._debug_dump.node_op_type(node_name))
    lines.append('  Device: %s' % self._debug_dump.node_device(node_name))
    output = debugger_cli_common.RichTextLines(lines, font_attr_segs=font_attr_segs)
    inputs = self._exclude_denylisted_ops(self._debug_dump.node_inputs(node_name))
    ctrl_inputs = self._exclude_denylisted_ops(self._debug_dump.node_inputs(node_name, is_control=True))
    output.extend(self._format_neighbors('input', inputs, ctrl_inputs))
    recs = self._exclude_denylisted_ops(self._debug_dump.node_recipients(node_name))
    ctrl_recs = self._exclude_denylisted_ops(self._debug_dump.node_recipients(node_name, is_control=True))
    output.extend(self._format_neighbors('recipient', recs, ctrl_recs))
    if parsed.attributes:
        output.extend(self._list_node_attributes(node_name))
    if parsed.dumps:
        output.extend(self._list_node_dumps(node_name))
    if parsed.traceback:
        output.extend(self._render_node_traceback(node_name))
    _add_main_menu(output, node_name=node_name, enable_node_info=False)
    return output

Command handler for node_info.

    Query information about a given node.

Args:
        args: Command-line arguments, excluding the command prefix, as a list of
        str.
        screen_info: Optional dict input containing screen information such as
        cols.

Returns:
        Output text lines as a RichTextLines object.


def load_custom_generate(self, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]]=None, trust_remote_code: Optional[bool]=None, **kwargs) -> Callable:
    is_local_code = os.path.exists(pretrained_model_name_or_path)
    has_custom_generate_folder = True
    if is_local_code:
        if not os.path.exists(os.path.join(pretrained_model_name_or_path, 'custom_generate/generate.py')):
            has_custom_generate_folder = False
    elif not file_exists(pretrained_model_name_or_path, 'custom_generate/generate.py'):
        has_custom_generate_folder = False
    if not has_custom_generate_folder:
        raise OSError(f"`{pretrained_model_name_or_path}` does not contain a `custom_generate` subdirectory with a `generate.py` file, can't load the custom generate function.")
    error_message = f'The repository `{pretrained_model_name_or_path}` contains custom generation code that will override the default `generate` method.'
    resolve_trust_remote_code(trust_remote_code, pretrained_model_name_or_path, has_local_code=is_local_code, has_remote_code=not is_local_code, error_message=error_message)
    check_python_requirements(pretrained_model_name_or_path, requirements_file='custom_generate/requirements.txt', **kwargs)
    module = get_cached_module_file(pretrained_model_name_or_path, module_file='custom_generate/generate.py', **kwargs)
    custom_generate_function = get_class_in_module('generate', module)
    return custom_generate_function

Loads and returns a custom generate function, given a model repo.

Args:
        pretrained_model_name_or_path (`str` or `os.PathLike`):
            Can be either:
                - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
                - A path to a *directory* containing model weights saved using
                [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
        trust_remote_code (`bool`, *optional*):
            Whether or not to allow for custom models defined on the Hub in their own modeling files. This option
            should only be set to `True` for repositories you trust and in which you have read the code, as it will
            execute code present on the Hub on your local machine.
        **kwargs:
            Additional keyword arguments for remote code loading.

Raises:
        OSError: If `pretrained_model_name_or_path` does not contain a `custom_generate` subdirectory.

Returns:
        A callable that can be used to generate text.


def __new__(cls, *args) -> 'InvokeExpressionNode':
    if not args:
        return super().__new__(cls)
    _, identifier, parent_node = args
    if identifier == 'reference' and isinstance(parent_node.return_type, _fhir_path_data_types.ReferenceStructureDataType):
        return super().__new__(InvokeReferenceNode)
    return super().__new__(InvokeExpressionNode)

Creates a new InvokeExpressionNode node or one of its subclasses.

    Creates either an InvokeExpressionNode or InvokeReferenceNode, a subclass of
    InvokeExpressionNode. The InvokeReferenceNode is returned when a field named
    'reference' is invoked against a FHIR Reference resource. Database backends
    have special behavior for reference nodes. This reference-specific node type
    allows them to define visitors to implement their reference-specific logic.

Args:
        *args: The args passed to `__init__`.

Returns:
        A new InvokeExpressionNode of the appropriate type.


def forward(self, input_ids: torch.LongTensor, spkr_id: torch.Tensor, lang_id: torch.Tensor) -> Tuple[torch.Tensor]:
    hidden_states = self.unit_embedding(input_ids).transpose(1, 2)
    spkr = self.speaker_embedding(spkr_id).transpose(1, 2)
    lang = self.language_embedding(lang_id).transpose(1, 2)
    log_dur_pred = self.dur_predictor(hidden_states.transpose(1, 2))
    dur_out = torch.clamp(torch.round(torch.expm1(log_dur_pred)).long(), min=1)
    if hidden_states.size(0) == 1:
        hidden_states = torch.repeat_interleave(hidden_states, dur_out.view(-1), dim=2)
    else:
        if hidden_states.shape[0] > 1 and self.training:
            logger.warning('`self.training=True` and you use batching. You lose parallelism during the hifigan\n                               forward pass because the samples are interleaved.')
        hidden_states = [torch.repeat_interleave(hidden_state, duration, dim=-1).transpose(0, 1) for hidden_state, duration in zip(hidden_states, dur_out)]
        hidden_states = nn.utils.rnn.pad_sequence(hidden_states, batch_first=True).transpose(1, 2)
    spkr = spkr.repeat(1, 1, hidden_states.shape[-1])
    lang = lang.repeat(1, 1, hidden_states.shape[-1])
    hidden_states = torch.cat([lang, hidden_states, spkr], dim=1)
    hidden_states = self.hifi_gan(hidden_states)
    unit_lengths = self._get_dur_output_lengths(input_ids, dur_out)
    lengths = self._get_output_hifigan_lengths(unit_lengths)
    return (hidden_states, lengths)

Args:
    input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
        Indices of input sequence tokens in the vocabulary.

        Indices can be obtained using [`SeamlessM4TTextToUnitForConditionalGeneration`]. [What are input
        IDs?](../glossary#input-ids)
    spkr_id (`int`, *optional*):
        The id of the speaker used for speech synthesis. Must be lower than `config.vocoder_num_spkrs`.
    tgt_lang (`str`, *optional*):
        The language id to use as target language for translation.


def convert_fields_for_spec(fields, field_values):
    _convert_fields(fields, field_values, context=_ConversionContext.SPEC)

Type-checks and converts field values for a TypeSpec (in place).

    This is similar to `convert_fields`, except that we expect a `TypeSpec` for
    tensor-like types.  In particular, if the `value_type` of a field is
    `tf.Tensor` or a `CompositeTensor` subclass, then the corresponding value in
    `fields` is expected to contain a `TypeSpec` (rather than a value described by
    that `TypeSpec`).

Args:
        fields: A list of `ExtensionTypeField` objects.
        field_values: A `dict` mapping field names to values.  Must contain an entry
        for each field.  I.e., `set(field_values.keys())` must be equal to
        `set([f.name for f in fields])`.

Raises:
        ValueError: If the keys of `field_values` do not match the names of
        the fields in `fields`.
        TypeError: If any value in `field_values` does not have the type indicated
        by the corresponding `ExtensionTypeField` object.


def __init__(self, name, state):
    stl.base.ParameterizedObject.__init__(self, name)
    self.state = state
    self.resolved_params = []

State specified in trasition spec and resolved.

Attributes:
        state: Original parameterized state.
        resolved_params: List of parameter values resolved. The order of values is
        same to that of parameters.


def num_records_produced(self, name=None):
    if self._reader_ref.dtype == dtypes.resource:
        return gen_io_ops.reader_num_records_produced_v2(self._reader_ref, name=name)
    else:
        return gen_io_ops.reader_num_records_produced(self._reader_ref, name=name)

Returns the number of records this reader has produced.

    This is the same as the number of Read executions that have
    succeeded.

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

Returns:
        An int64 Tensor.


def set_shard_dimensions(self, shard_dimensions):
    if len(shard_dimensions) != self.number_of_tuple_elements:
        raise ValueError(f'shard_dimensions is {str(shard_dimensions)}, but must be a list of length {self.number_of_tuple_elements}')
    for policy, dimension in zip(self._sharding_policies, shard_dimensions):
        policy.set_shard_dimension(dimension)
    self._validate()

Sets the shard_dimension of each element of the queue.

    shard_dimensions must be a list of length
    self.number_of_tuple_elements, and each element must be
    convertible to a Dimension compatible with self.tuple_shapes.

Args:
        shard_dimensions: the dimensions of each queue element.

Raises:
        ValueError: if shard_dimensions is not of length
        self.number_of_tuple_elements; or an element of
        shard_dimensions cannot be converted to a Dimension; or an
        element of shard_dimensions is a Dimension that is out of
        range for the corresponding tuple element shape.


def segment_ids_to_row_splits(segment_ids, num_segments=None, out_type=None, name=None):
    from tensorflow.python.ops import bincount_ops
    if out_type is None:
        if isinstance(segment_ids, tensor.Tensor):
            out_type = segment_ids.dtype
        elif isinstance(num_segments, tensor.Tensor):
            out_type = num_segments.dtype
        else:
            out_type = dtypes.int64
    else:
        out_type = dtypes.as_dtype(out_type)
    with ops.name_scope(name, 'SegmentIdsToRaggedSplits', [segment_ids]) as name:
        segment_ids = ragged_util.convert_to_int_tensor(segment_ids, 'segment_ids', dtype=dtypes.int32)
        segment_ids.shape.assert_has_rank(1)
        if num_segments is not None:
            num_segments = ragged_util.convert_to_int_tensor(num_segments, 'num_segments', dtype=dtypes.int32)
            num_segments.shape.assert_has_rank(0)
        row_lengths = bincount_ops.bincount(segment_ids, minlength=num_segments, maxlength=num_segments, dtype=out_type)
        splits = array_ops.concat([[0], math_ops.cumsum(row_lengths)], axis=0)
        if num_segments is not None:
            const_num_segments = tensor_util.constant_value(num_segments)
            if const_num_segments is not None:
                splits.set_shape(tensor_shape.TensorShape([const_num_segments + 1]))
        return splits

Generates the RaggedTensor `row_splits` corresponding to a segmentation.

    Returns an integer vector `splits`, where `splits[0] = 0` and
    `splits[i] = splits[i-1] + count(segment_ids==i)`.  Example:

    >>> print(tf.ragged.segment_ids_to_row_splits([0, 0, 0, 2, 2, 3, 4, 4, 4]))
    tf.Tensor([0 3 3 5 6 9], shape=(6,), dtype=int64)

Args:
        segment_ids: A 1-D integer Tensor.
        num_segments: A scalar integer indicating the number of segments.  Defaults
        to `max(segment_ids) + 1` (or zero if `segment_ids` is empty).
        out_type: The dtype for the return value.  Defaults to `segment_ids.dtype`,
        or `tf.int64` if `segment_ids` does not have a dtype.
        name: A name prefix for the returned tensor (optional).

Returns:
        A sorted 1-D integer Tensor, with `shape=[num_segments + 1]`.


def _encode_fhir_path_builder(self, builder: expressions.Builder) -> Optional[str]:
    try:
        sql_expression = self._bq_interpreter.encode(builder)
    except Exception as e:
        self._error_reporter.report_fhir_path_error(self._abs_path_invocation(builder), str(builder), self._error_message_for_exception(e))
        return None
    return sql_expression

Returns a Standard SQL translation of the constraint `fhir_path_expression`.

    If an error is encountered during encoding, the associated error reporter
    will be notified, and this method will return `None`.

Args:
        builder: Builder containing the information to be encoded to Standard SQL.

Returns:
        A Standard SQL encoding of the constraint `fhir_path_expression` upon
        successful completion. The SQL will evaluate to a single boolean
        indicating whether the constraint is satisfied.


def cuda(self) -> Rotation:
    if self._rot_mats is not None:
        return Rotation(rot_mats=self._rot_mats.cuda(), quats=None)
    elif self._quats is not None:
        return Rotation(rot_mats=None, quats=self._quats.cuda(), normalize_quats=False)
    else:
        raise ValueError('Both rotations are None')

Analogous to the cuda() method of torch Tensors

Returns:
        A copy of the Rotation in CUDA memory


def test_sample_paths_2d(self, random_type, seed):
    mu = np.array([0.2, 0.7])
    a = np.array([[0.4, 0.1], [0.3, 0.2]])
    b = np.array([[0.33, -0.03], [0.21, 0.5]])

    def drift_fn(t, x):
        return mu * tf.sqrt(t) * tf.ones_like(x, dtype=t.dtype)

    def vol_fn(t, x):
        del x
        return (a * t + b) * tf.ones([2, 2], dtype=t.dtype)
    num_samples = 10000
    times = np.array([0.1, 0.21, 0.32, 0.43, 0.55])
    x0 = np.array([0.1, -1.1])
    paths = self.evaluate(euler_sampling.sample(dim=2, drift_fn=drift_fn, volatility_fn=vol_fn, times=times, num_samples=num_samples, initial_state=x0, time_step=0.01, random_type=random_type, seed=seed))
    self.assertAllClose(paths.shape, (num_samples, 5, 2), atol=0)
    means = np.mean(paths, axis=0)
    times = np.reshape(times, [-1, 1])
    expected_means = x0 + 2.0 / 3.0 * mu * np.power(times, 1.5)
    self.assertAllClose(means, expected_means, rtol=0.01, atol=0.01)

Tests path properties for 2-dimentional Ito process.

    We construct the following Ito processes.

    dX_1 = mu_1 sqrt(t) dt + s11 dW_1 + s12 dW_2
    dX_2 = mu_2 sqrt(t) dt + s21 dW_1 + s22 dW_2

    mu_1, mu_2 are constants.
    s_ij = a_ij t + b_ij

    For this process expected value at time t is (x_0)_i + 2/3 * mu_i * t^1.5.

Args:
        random_type: Random number type defined by tff.math.random.RandomType
        enum.
        seed: Random seed.


def enable_cpu_offload(self, accelerator_id: Optional[int]=0, **kwargs):
    if is_accelerate_available():
        from accelerate import cpu_offload_with_hook
    else:
        raise ImportError('`enable_model_cpu_offload` requires `accelerate`.')
    gpu_id = kwargs.get('gpu_id', 0)
    if gpu_id != 0:
        warnings.warn('The argument `gpu_id` is deprecated and will be removed in version 4.54.0 of Transformers. Please use `accelerator_id` instead.', FutureWarning)
        accelerator_id = gpu_id
    device_type = 'cuda'
    if is_torch_accelerator_available():
        device_type = torch.accelerator.current_accelerator().type
    device = torch.device(f'{device_type}:{accelerator_id}')
    torch_accelerator_module = getattr(torch, device_type)
    if self.device.type != 'cpu':
        self.to('cpu')
        torch_accelerator_module.empty_cache()
    self.semantic.input_embeds_layer, _ = cpu_offload_with_hook(self.semantic.input_embeds_layer, device)
    hook = None
    for cpu_offloaded_model in [self.semantic, self.coarse_acoustics, self.fine_acoustics]:
        _, hook = cpu_offload_with_hook(cpu_offloaded_model, device, prev_module_hook=hook)
    self.fine_acoustics_hook = hook
    _, hook = cpu_offload_with_hook(self.codec_model, device, prev_module_hook=hook)
    self.codec_model_hook = hook

Offloads all sub-models to CPU using accelerate, reducing memory usage with a low impact on performance. This
    method moves one whole sub-model at a time to the accelerator when it is used, and the sub-model remains in accelerator until the next sub-model runs.

Args:
        accelerator_id (`int`, *optional*, defaults to 0):
            accelerator id on which the sub-models will be loaded and offloaded. This argument is deprecated.
        kwargs (`dict`, *optional*):
            additional keyword arguments:
                `gpu_id`: accelerator id on which the sub-models will be loaded and offloaded.


def distribute_tensor(tensor, layout):
    from keras.src.distribution import TensorLayout
    if isinstance(layout, TensorLayout):
        layout = layout.backend_layout
    if jax_utils.is_in_jax_tracing_scope():
        return jax.lax.with_sharding_constraint(tensor, layout)
    if isinstance(tensor, jax.Array):
        if isinstance(layout, jax.sharding.Sharding) and tensor.sharding.is_equivalent_to(layout, ndim=len(tensor.shape)):
            return tensor
        elif isinstance(layout, jax_layout.Layout):
            current_layout = getattr(tensor, 'layout', None)
            if current_layout == layout:
                return tensor
    return jax.device_put(tensor, layout)

Distribute the tensor based on the layout.

    Note that this function can be used both in eager context, or within a
    jitted function.

Args:
        tensor: `jax.Array` that need to be distributed.
        layout: `TensorLayout` for the created variable, or a
            JAX-supported layout instance
            (e.g. `jax.experimental.layout.Layout`, `jax.sharding.Sharding`).

Returns:
        Distributed value.


def forward(self, pixel_values: Optional[torch.Tensor]=None, bool_masked_pos: Optional[torch.BoolTensor]=None, head_mask: Optional[torch.Tensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, interpolate_pos_encoding: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[tuple, BaseModelOutput]:
    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
    if pixel_values is None:
        raise ValueError('You have to specify pixel_values')
    head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
    embedding_output = self.embeddings(pixel_values, bool_masked_pos=bool_masked_pos, interpolate_pos_encoding=interpolate_pos_encoding)
    encoder_outputs = self.encoder(embedding_output, head_mask=head_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    sequence_output = encoder_outputs[0]
    sequence_output = self.layernorm(sequence_output)
    if not return_dict:
        head_outputs = (sequence_output,)
        return head_outputs + encoder_outputs[1:]
    return BaseModelOutput(last_hidden_state=sequence_output, hidden_states=encoder_outputs.hidden_states, attentions=encoder_outputs.attentions)

bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`, *optional*):
    Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).

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

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

    >>> image_processor = AutoImageProcessor.from_pretrained("facebook/vit-msn-small")
    >>> model = ViTMSNModel.from_pretrained("facebook/vit-msn-small")
    >>> inputs = image_processor(images=image, return_tensors="pt")
    >>> with torch.no_grad():
    ...     outputs = model(**inputs)
    >>> last_hidden_states = outputs.last_hidden_state
    ```


def get_conversion_metadata(model_buffer):
    model_object = flatbuffer_utils.convert_bytearray_to_object(model_buffer)
    if not model_object or not model_object.metadata:
        return None
    for meta in model_object.metadata:
        if meta.name.decode('utf-8') == CONVERSION_METADATA_FIELD_NAME:
            metadata_buf = model_object.buffers[meta.buffer].data.tobytes()
            return conversion_metadata_fb.ConversionMetadataT.InitFromObj(conversion_metadata_fb.ConversionMetadata.GetRootAsConversionMetadata(metadata_buf, 0))
    return None

Read conversion metadata from a tflite model.

Args:
        model_buffer: A tflite model.

Returns:
        The conversion metadata or None if it is not populated.


def stateless_random_shuffle(input_tensor, seed, name=None):
    with tf.compat.v1.name_scope(name, default_name='stateless_random_shuffle', values=[input_tensor, seed]):
        input_tensor = tf.convert_to_tensor(input_tensor, name='input_tensor')
        seed = tf.convert_to_tensor(seed, name='random_seed')
        uniforms = tf.random.stateless_uniform(shape=[tf.shape(input_tensor)[0]], seed=seed, dtype=tf.float64)
    return tf.gather(input_tensor, tf.argsort(uniforms, stable=True, axis=0))

Produces stateless random shuffle of the 1st dimension of an input Tensor.

    This is a stateless version of `tf.random_shuffle`. If run twice with the same
    seed, produces the same result.

    Example
    ```python
    identity_shuffle = tf.range(100)
    random_shuffle = stateless_random_shuffle(identity_shuffle, seed=(42, 2))
    ```

Args:
        input_tensor: float32, float64, int32 or int64 1-D Tensor.
        seed: int32 or int64 Tensor of shape [2].
        name: Python `str` name prefixed to ops created by this function.

Returns:
        A Tensor of the same shape and dtype as `input_tensor`.


def match_next_flag(tt_flags, pos):
    match = _FLAG_DOUBLE_QUOTE_PAT.match(tt_flags, pos)
    if match:
        return (match, True)
    match = _FLAG_SINGLE_QUOTE_PAT.match(tt_flags, pos)
    if match:
        return (match, True)
    match = _FLAG_NO_QUOTE_PAT.match(tt_flags, pos)
    if match:
        return (match, True)
    match = _FLAG_NO_EQUAL_PAT.match(tt_flags, pos)
    if match:
        return (match, False)
    return (None, False)

Returns the match for the next TensorTracer flag.

Args:
        tt_flags: a string that contains the flags.
        pos: where in flags to start the search.

Returns:
        A pair where the first element is the regular-expression
        match found and the second element indicates if the match
        has a value.


def _broadcast_normalize_batch_in_training(x, gamma, beta, reduction_axes, epsilon=0.001):
    mean, var = nn.moments(x, reduction_axes, None, None, False)
    target_shape = []
    for axis in range(ndim(x)):
        if axis in reduction_axes:
            target_shape.append(1)
        else:
            target_shape.append(array_ops.shape(x)[axis])
    target_shape = array_ops_stack.stack(target_shape)
    broadcast_mean = array_ops.reshape(mean, target_shape)
    broadcast_var = array_ops.reshape(var, target_shape)
    if gamma is None:
        broadcast_gamma = None
    else:
        broadcast_gamma = array_ops.reshape(gamma, target_shape)
    if beta is None:
        broadcast_beta = None
    else:
        broadcast_beta = array_ops.reshape(beta, target_shape)
    normed = nn.batch_normalization(x, broadcast_mean, broadcast_var, broadcast_beta, broadcast_gamma, epsilon)
    return (normed, mean, var)

Non-fused, broadcast version of `normalize_batch_in_training`.

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

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


def _test_runner(test_id, test_env):
    global _running_in_worker, _env
    _running_in_worker = True
    _env = test_env
    test = unittest.defaultTestLoader.loadTestsFromName(test_id)
    runner = unittest.TextTestRunner()
    result = runner.run(test)
    failures = result.failures + result.expectedFailures + result.errors
    if failures:
        ret = _TestResult(status='failure', message=failures[0][1])
    elif result.skipped:
        ret = _TestResult(status='skipped', message=result.skipped[0][1])
    else:
        ret = _TestResult(status='ok', message=None)
    if ret.message:
        print(ret.message)
    return ret

Executes the test with the given test_id.

    This is a simple wrapper around TestRunner to be used with
    multi_process_runner. Similar to test.main(), but it executes only one test
    specified by test_id and returns whether the test succeeds. If the test fails,
    the function prints failures and errors to stdout.

Args:
        test_id: TestCase.id()
        test_env: a TestEnvironment object.

Returns:
        A boolean indicates whether the test succeeds.


def __init__(self, state, value):
    stl.base.NamedObject.__init__(self, state)
    self.value = value
    self.param_values = []

State instance defined in a state transition spec.

    A state instance is defined either in "pre_states", "post_states", or
    "error_states" in a state transition spec. Resolve() turns this into
    state.StateValue resolved.

Attributes:
        value: Current value of state of |name|.
        param_values: List of unresolved parameters of state of |name|.


def log_get(recipe_id=[], timezone='America/Los_Angeles', days=1):
    body = {'resourceNames': ['projects/%s' % UI_PROJECT], 'filter': '       logName="projects/%s/logs/StarThinker"        AND labels.version="%s"        AND labels.layer="JOB"     ' % (UI_PROJECT, LOG_VERSION), 'orderBy': 'timestamp desc', 'pageSize': 1000}
    if recipe_id:
        if isinstance(recipe_id, str):
            recipe_id = [recipe_id]
        body['filter'] += ' AND ( %s )' % ' OR '.join(('operation.id="%s"' % r for r in recipe_id))
    for entry in API_StackDriver(Configuration(service=UI_SERVICE, project=UI_PROJECT), 'service', iterate=True).entries().list(body=body).execute():
        yield entry

Returns last actionable job run for a specific recipe or all recipes.

    Pulls status entries from StackDriver in reverse order.  A single recipe may
    be run multiple times for multiple tasks at different hours, do not
    assume a JOB_END means a recipe is complete.  Only way to ensure a recipe is
    complete
    is to compare all tasks run against all tasks in recipe ( not done by log
    code).

    Args: - recipe_id ( string or list ) - Optional, if provided returns a single
    record for a single job. - timezone ( string ) - The local timezone to cast
    all record times into.

Returns:
        - ( iterator ) - Each log entry.


def _try_put(self, item):
    try:
        self._event_queue.put(item)
    except QueueClosedError:
        self._internal_close()
        if self._worker.failure_exc_info:
            _, exception, _ = self._worker.failure_exc_info
            raise exception from None

Attempts to enqueue an item to the event queue.

    If the queue is closed, this will close the EventFileWriter and reraise the
    exception that caused the queue closure, if one exists.

Args:
        item: the item to enqueue


def recipe_bigquery_storage(config, auth_read, bucket, auth_write, path, dataset, table, schema):
    bigquery(config, {'auth': auth_read, 'from': {'bucket': bucket, 'path': path}, 'to': {'auth': auth_write, 'dataset': dataset, 'table': table}, 'schema': schema})

Move using bucket and path prefix.

Args:
        auth_read (authentication) - Credentials used for reading data.
        bucket (string) - Google cloud bucket.
        auth_write (authentication) - Credentials used for writing data.
        path (string) - Path prefix to read from, no * required.
        dataset (string) - Existing BigQuery dataset.
        table (string) - Table to create from this query.
        schema (json) - Schema provided in JSON list format or empty list.


def _time_step(time, output_ta_t, state):
    if in_graph_mode:
        input_t = tuple((ta.read(time) for ta in input_ta))
        for input_, shape in zip(input_t, inputs_got_shape):
            input_.set_shape(shape[1:])
    else:
        input_t = tuple((ta[time.numpy()] for ta in input_ta))
    input_t = nest.pack_sequence_as(structure=inputs, flat_sequence=input_t)
    call_cell = lambda: cell(input_t, state)
    if sequence_length is not None:
        output, new_state = _rnn_step(time=time, sequence_length=sequence_length, min_sequence_length=min_sequence_length, max_sequence_length=max_sequence_length, zero_output=zero_output, state=state, call_cell=call_cell, state_size=state_size, skip_conditionals=True)
    else:
        output, new_state = call_cell()
    output = nest.flatten(output)
    if in_graph_mode:
        output_ta_t = tuple((ta.write(time, out) for ta, out in zip(output_ta_t, output)))
    else:
        for ta, out in zip(output_ta_t, output):
            ta[time.numpy()] = out
    return (time + 1, output_ta_t, new_state)

Take a time step of the dynamic RNN.

Args:
        time: int32 scalar Tensor.
        output_ta_t: List of `TensorArray`s that represent the output.
        state: nested tuple of vector tensors that represent the state.

Returns:
        The tuple (time + 1, output_ta_t with updated flow, new_state).


def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, layer_head_mask: torch.Tensor, is_index_masked=None, is_index_global_attn=None, is_global_attn=None, output_attentions=False):
    residual = hidden_states
    attn_outputs = self.self_attn(hidden_states=hidden_states, attention_mask=attention_mask, layer_head_mask=layer_head_mask, is_index_masked=is_index_masked, is_index_global_attn=is_index_global_attn, is_global_attn=is_global_attn, output_attentions=output_attentions)
    hidden_states = attn_outputs[0]
    hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
    hidden_states = residual + hidden_states
    hidden_states = self.self_attn_layer_norm(hidden_states)
    residual = hidden_states
    hidden_states = self.activation_fn(self.fc1(hidden_states))
    hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
    hidden_states = self.fc2(hidden_states)
    hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
    hidden_states = residual + hidden_states
    hidden_states = self.final_layer_norm(hidden_states)
    if hidden_states.dtype == torch.float16 and (torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any()):
        clamp_value = torch.finfo(hidden_states.dtype).max - 1000
        hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
    return (hidden_states,) + attn_outputs[1:]

Args:
    hidden_states (`torch.FloatTensor`): input to the layer of shape *(batch, seq_len, embed_dim)*
    attention_mask (`torch.FloatTensor`): attention mask of size
        *(batch, 1, tgt_len, src_len)* where padding elements are indicated by very large negative values.
    layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size
        *(encoder_attention_heads,)*.


def _to_int(operand: List[WorkSpaceMessage]) -> Optional[int]:
    if not operand:
        return None
    if len(operand) > 1:
        raise ValueError('Expected single int result but got multiple items.')
    if not fhir_types.is_integer(operand[0].message):
        raise ValueError('Expected single int but got a non-int value.')
    return proto_utils.get_value_at_field(operand[0].message, 'value')

Converts an evaluation result to an int value or None.

Args:
        operand: an expression operand result to convert to int.

Returns:
        the int value, or None if the operand was empty.

Raises:
        ValueError if it is not an empty result or a single, int value.


def __init__(self, tensor_proto, initialized=True):
    self._tensor_proto = tensor_proto
    self._initialized = initialized

Constructor.

Args:
        tensor_proto: the `TensorProto` object that cannot be represented as a
        `np.ndarray` object.
        initialized: (`bool`) whether the Tensor is initialized.


def _maybe_save_assets(write_fn, assets_to_add=None):
    asset_filename_map = {}
    if assets_to_add is None:
        tf_logging.info('No assets to save.')
        return asset_filename_map
    for asset_tensor in assets_to_add:
        asset_source_filepath = _asset_path_from_tensor(asset_tensor)
        if not asset_source_filepath:
            raise ValueError(f'Asset filepath tensor {asset_tensor} in is invalid.')
        asset_filename = get_asset_filename_to_add(asset_source_filepath, asset_filename_map)
        write_fn(asset_filename, asset_tensor)
        asset_filename_map[asset_filename] = asset_source_filepath
    tf_logging.info('Assets added to graph.')
    return asset_filename_map

Saves assets to the meta graph.

Args:
        write_fn: A function callback that writes assets into meta graph.
        assets_to_add: The list where the asset paths are setup.

Returns:
        A dict of asset basenames for saving to the original full path to the asset.

Raises:
        ValueError: Indicating an invalid filepath tensor.


def mask_to_rgb(self, image: np.ndarray, palette: Optional[List[Tuple[int, int]]]=None, data_format: Optional[Union[str, ChannelDimension]]=None) -> np.ndarray:
    return mask_to_rgb(image, palette=palette, data_format=data_format)

Converts a segmentation map to RGB format.

Args:
        image (`np.ndarray`):
            Segmentation map with dimensions (height, width) where pixel values represent the class index.
        palette (`List[Tuple[int, int]]`, *optional*, defaults to `None`):
            Palette to use to convert the mask to RGB format. If unset, the mask is duplicated across the channel
            dimension.
        data_format (`ChannelDimension` or `str`, *optional*):
            The channel dimension format for the output image. If unset, the channel dimension format of the input
            image is used. Can be one of:
            - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
            - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.

Returns:
        `np.ndarray`: The mask in RGB format.


def add_session_log(self, session_log, global_step=None):
    event = event_pb2.Event(session_log=session_log)
    self._add_event(event, global_step)

Adds a `SessionLog` protocol buffer to the event file.

    This method wraps the provided session in an `Event` protocol buffer
    and adds it to the event file.

Args:
        session_log: A `SessionLog` protocol buffer.
        global_step: Number. Optional global step value to record with the
        summary.


def parse_step(step_name):
    prefix = 'step'
    step_name = step_name.lower().replace(' ', '_')
    step_name = step_name[len(prefix):] if prefix and step_name.startswith(prefix) else step_name
    return step_name.strip(':_')

Replaces white spaces and removes 'Step:' label

Args:
        step_name(str): step name passed in metric ParDo

Returns:
        lower case step name without namespace and step label


def _create_controller_info_record(self, controller_module_name):
    module = self._controller_modules[controller_module_name]
    controller_info = None
    try:
        controller_info = module.get_info(copy.copy(self._controller_objects[controller_module_name]))
    except AttributeError:
        logging.warning('No optional debug info found for controller %s. To provide it, implement `get_info`.', controller_module_name)
    try:
        yaml.dump(controller_info)
    except TypeError:
        logging.warning('The info of controller %s in class "%s" is not YAML serializable! Coercing it to string.', controller_module_name, self._class_name)
        controller_info = str(controller_info)
    return records.ControllerInfoRecord(self._class_name, module.MOBLY_CONTROLLER_CONFIG_NAME, controller_info)

Creates controller info record for a particular controller type.

    Info is retrieved from all the controller objects spawned from the
    specified module, using the controller module's `get_info` function.

Args:
        controller_module_name: string, the name of the controller module
        to retrieve info from.

Returns:
        A records.ControllerInfoRecord object.


def relative_batch_tokens_ids_to_midi(self, tokens: np.ndarray, beatstep: np.ndarray, beat_offset_idx: int=0, bars_per_batch: int=2, cutoff_time_idx: int=12):
    beat_offset_idx = 0 if beat_offset_idx is None else beat_offset_idx
    notes = self.relative_batch_tokens_ids_to_notes(tokens=tokens, beat_offset_idx=beat_offset_idx, bars_per_batch=bars_per_batch, cutoff_time_idx=cutoff_time_idx)
    midi = self.notes_to_midi(notes, beatstep, offset_sec=beatstep[beat_offset_idx])
    return midi

Converts tokens to Midi. This method calls `relative_batch_tokens_ids_to_notes` method to convert batch tokens
    to notes then uses `notes_to_midi` method to convert them to Midi.

Args:
        tokens (`numpy.ndarray`):
            Denotes tokens which alongside beatstep will be converted to Midi.
        beatstep (`np.ndarray`):
            We get beatstep from feature extractor which is also used to get Midi.
        beat_offset_idx (`int`, *optional*, defaults to 0):
            Denotes beat offset index for each note in generated Midi.
        bars_per_batch (`int`, *optional*, defaults to 2):
            A parameter to control the Midi output generation.
        cutoff_time_idx (`int`, *optional*, defaults to 12):
            Denotes the cutoff time index for each note in generated Midi.


def _create_all_weights(self, params):
    raise NotImplementedError

Creates and sets all optimizer weights.

Args:
        params: list or tuple of `Variable` objects that will be minimized
        using this optimizer.

Returns:
        Specific weight values that are used in `get_updates`


def __init__(self, config: GraniteMoeHybridConfig):
    super(GraniteMoeHybridMoE, self).__init__()
    self.input_size = config.hidden_size
    self.hidden_size = config.intermediate_size
    self.activation = ACT2FN[config.hidden_act]
    self.input_linear = GraniteMoeHybridParallelExperts(config.num_local_experts, self.input_size, self.hidden_size * 2)
    self.output_linear = GraniteMoeHybridParallelExperts(config.num_local_experts, self.hidden_size, self.input_size)
    self.router = GraniteMoeHybridTopKGating(input_size=self.input_size, num_experts=config.num_local_experts, top_k=config.num_experts_per_tok)

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

Args:
        config:
            Configuration object with model hyperparameters.


def _MaybeColocateWith(inputs):
    if not inputs:
        yield
    else:
        with ops.colocate_with(inputs[0]), _MaybeColocateWith(inputs[1:]):
            yield

A context manager for (maybe) colocating with a list of input tensors.

Args:
        inputs: A list of `Tensor` or `Operation` objects.

Returns:
        A context manager.


def make_initializable_iterator(dataset: DatasetV1, shared_name=None) -> iterator_ops.Iterator:
    try:
        return dataset._make_initializable_iterator(shared_name)
    except AttributeError:
        return DatasetV1Adapter(dataset)._make_initializable_iterator(shared_name)

Creates an iterator for elements of `dataset`.

    Note: The returned iterator will be in an uninitialized state,
    and you must run the `iterator.initializer` operation before using it:

    ```python
    dataset = ...
    iterator = tf.compat.v1.data.make_initializable_iterator(dataset)
    # ...
    sess.run(iterator.initializer)
    ```

Args:
        dataset: A `tf.data.Dataset`.
        shared_name: (Optional.) If non-empty, the returned iterator will be shared
        under the given name across multiple sessions that share the same devices
        (e.g. when using a remote server).

Returns:
        A `tf.data.Iterator` for elements of `dataset`.

Raises:
        RuntimeError: If eager execution is enabled.

    @compatibility(TF2)
    This is a legacy API for consuming dataset elements and should only be used
    during transition from TF 1 to TF 2. Note that using this API should be
    a transient state of your code base as there are in general no guarantees
    about the interoperability of TF 1 and TF 2 code.

    In TF 2 datasets are Python iterables which means you can consume their
    elements using `for elem in dataset: ...` or by explicitly creating iterator
    via `iterator = iter(dataset)` and fetching its elements via
    `values = next(iterator)`.
    @end_compatibility


def _cmp_rel(self, state, op_name, x, y):
    ret = self.ctx.program.NewVariable()
    leftover_x = self.ctx.program.NewVariable()
    leftover_y = self.ctx.program.NewVariable()
    op_not_eq = op_name not in ('EQ', 'NE')
    reported = False
    for b1 in x.bindings:
        for b2 in y.bindings:
            op = getattr(slots, op_name)
            try:
                err = False
                val = compare.cmp_rel(self.ctx, op, b1.data, b2.data)
            except compare.CmpTypeError:
                val = None
                if state.node.HasCombination([b1, b2]):
                    err = True
                    reported = True
                    self.ctx.errorlog.unsupported_operands(self.frames, op, x, y)
            if val is None:
                if op_not_eq and isinstance(b1.data, abstract.Class) and err:
                    ret.AddBinding(self.ctx.convert.unsolvable, {b1, b2}, state.node)
                elif isinstance(b1.data, abstract.SequenceLength):
                    ret.AddBinding(self.ctx.convert.bool_values[val], {b1, b2}, state.node)
                else:
                    leftover_x.PasteBinding(b1, state.node)
                    leftover_y.PasteBinding(b2, state.node)
            else:
                ret.AddBinding(self.ctx.convert.bool_values[val], {b1, b2}, state.node)
    if leftover_x.bindings:
        op = f'__{op_name.lower()}__'
        report_errors = op_not_eq and (not bool(ret.bindings)) and (not reported)
        state, leftover_ret = vm_utils.call_binary_operator(state, op, leftover_x, leftover_y, report_errors=report_errors, ctx=self.ctx)
        ret.PasteVariable(leftover_ret, state.node)
    return (state, ret)

Implementation of relational operators CMP_(LT|LE|EQ|NE|GE|GT).

Args:
        state: Initial FrameState.
        op_name: An operator name, e.g., "EQ".
        x: A variable of the lhs value.
        y: A variable of the rhs value.

Returns:
        A tuple of the new FrameState and the return variable.


def _assign_stablehlo_quantization_config_or_populate_default(self, args):
    if self.experimental_stablehlo_quantizer_config is not None and Optimize.DEFAULT not in self.optimizations:
        args['quantization_config'] = self.experimental_stablehlo_quantizer_config
    elif Optimize.DEFAULT in self.optimizations and self.representative_dataset:
        if len(self._saved_model_exported_names) != 1:
            raise ValueError('StableHLO quantizer is only supported when converting from a SavedModel with one signature key.')
        signature_key = self._saved_model_exported_names[0]
        tfrecord_file_path = tempfile.mkstemp(suffix='.tfrecord', prefix=signature_key)[1]
        rd.TfRecordRepresentativeDatasetSaver({signature_key: tfrecord_file_path}).save({signature_key: self.representative_dataset()})
        quantization_config = qc.QuantizationConfig(static_range_ptq_preset=qc.StaticRangePtqPreset(representative_datasets=[qc.RepresentativeDatasetConfig(tf_record=qc.TfRecordFile(path=tfrecord_file_path))], enable_per_channel_quantized_weight=True, enable_full_int_quantization=True), pipeline_config=qc.PipelineConfig(unpack_quantized_types=False))
        args['quantization_config'] = quantization_config
    else:
        raise ValueError('StableHLO quantizer only supports static-range and weight-only PTQ.')

Assigns `QuantizationConfig` to `args` or populate default.

Args:
        args: Dictionary of argument names and associated values.


def assert_type_v2(tensor, tf_type, message=None, name=None):
    assert_type(tensor=tensor, tf_type=tf_type, message=message, name=name)

Asserts that the given `Tensor` is of the specified type.

    This can always be checked statically, so this method returns nothing.

Example:
    >>> a = tf.Variable(1.0)
    >>> tf.debugging.assert_type(a, tf_type= tf.float32)

    >>> b = tf.constant(21)
    >>> tf.debugging.assert_type(b, tf_type=tf.bool)
    Traceback (most recent call last):
    ...
    TypeError: ...

    >>> c = tf.SparseTensor(indices=[[0, 0], [1, 2]], values=[1, 2],
    ...  dense_shape=[3, 4])
    >>> tf.debugging.assert_type(c, tf_type= tf.int32)

Args:
        tensor: A `Tensor`, `SparseTensor` or `tf.Variable` .
        tf_type: A tensorflow type (`dtypes.float32`, `tf.int64`, `dtypes.bool`,
        etc).
        message: A string to prefix to the default message.
        name:  A name for this operation. Defaults to "assert_type"

Raises:
        TypeError: If the tensor's data type doesn't match `tf_type`.


def from_fn(cls, dna_spec: DNASpec, generator_fn: Callable[['DecisionPoint'], Union[List[int], float, str, 'DNA']]) -> 'DNA':
    if not isinstance(dna_spec, DNASpec):
        raise TypeError(f"Argument 'dna_spec' should be DNASpec type. Encountered {dna_spec}.")
    if dna_spec.is_space:
        children = []
        for child_spec in dna_spec.elements:
            children.append(DNA.from_fn(child_spec, generator_fn))
        if len(children) == 1:
            return children[0]
        dna = DNA(None, children)
    elif dna_spec.is_categorical:
        assert isinstance(dna_spec, DecisionPoint), dna_spec
        decision = generator_fn(dna_spec)
        if isinstance(decision, DNA):
            dna = decision
        else:
            if len(decision) != dna_spec.num_choices:
                raise ValueError(f'Number of DNA child values does not match the number of choices. Child values: {decision!r}, Choices: {dna_spec.num_choices}, Location: {dna_spec.location.path}.')
            children = []
            for i, choice in enumerate(decision):
                choice_location = utils.KeyPath(i, dna_spec.location)
                if not isinstance(choice, int):
                    raise ValueError(f'Choice value should be int. Encountered: {choice}, Location: {choice_location.path}.')
                if choice >= len(dna_spec.candidates):
                    raise ValueError(f'Choice out of range. Value: {choice}, Candidates: {len(dna_spec.candidates)}, Location: {choice_location.path}.')
                child_dna = DNA.from_fn(dna_spec.candidates[choice], generator_fn)
                children.append(DNA(choice, [child_dna]))
            dna = DNA(None, children)
    else:
        assert isinstance(dna_spec, DecisionPoint), dna_spec
        decision = generator_fn(dna_spec)
        if isinstance(decision, DNA):
            dna = decision
        else:
            dna = DNA(decision)
    dna_spec.validate(dna)
    return dna

Generate a DNA with user generator function.

Args:
        dna_spec: The DNASpec for the DNA.
        generator_fn: A callable object with signature:

        `(decision_point) -> decision`

        The decision_point is a `Choices` object or a `Float` object.
        The returned decision should be:

        * a list of integer or a DNA object for a `Choices` decision point.
        When a DNA is returned, it will be used as the DNA for the entire
        sub-tree, hence `generate_fn` will not be called on sub-decision
        points.
        * a float or a DNA object for a Float decision point.
        * a string or a DNA object for a CustomDecisionPoint.

Returns:
        A DNA generated from the user function.


def regularized_evolution(mutator=mutators.Uniform(), population_size: int=100, tournament_size: int=10, seed: Optional[int]=None):
    if tournament_size < 2:
        raise ValueError(f'`tournament_size` must be no less than 2. Encountered: {tournament_size}')
    if population_size < tournament_size:
        raise ValueError(f'The value of `population_size` ({population_size}) must be no less than the value of `tournament_size` ({tournament_size}).')
    return base.Evolution(selectors.Random(tournament_size, seed=seed) >> selectors.Top(1) >> mutator, population_init=(pg.geno.Random(seed=seed), population_size), population_update=selectors.Last(population_size))

Regularized Evolution algorithm.

    https://www.aaai.org/ojs/index.php/AAAI/article/view/4405.

Args:
        mutator: Mutator to use.
        population_size: Population size. Must be larger than tournament size.
        tournament_size: Tournament size.
        seed: Random seed. If None, the current system time is used as seed.

Returns:
        An `Evolution` object.


def Save(obj: _Serializable, filename: Path, compress: bool=False, open_function=open) -> None:
    with open_function(filename, 'wb') as fi:
        if compress:
            with gzip.GzipFile(filename='', mode='wb', fileobj=fi, mtime=1.0) as zfi:
                zfi.write(Encode(obj))
        else:
            fi.write(Encode(obj))

Saves a serializable object to a file.

Args:
        obj: The object to serialize.
        filename: filename to write to.
        compress: if True, the data will be compressed using gzip. The given
        filename will be used, unaltered.
        open_function: The function to use to open files. Defaults to the builtin
        open() function.


def slice_update(inputs, start_indices, updates):
    if any_symbolic_tensors((inputs, start_indices, updates)):
        return SliceUpdate().symbolic_call(inputs, start_indices, updates)
    return backend.core.slice_update(inputs, start_indices, updates)

Update an input by slicing in a tensor of updated values.

    At a high level, this operation does
    `inputs[start_indices: start_indices + updates.shape] = updates`.
    Assume inputs is a tensor of shape `(D0, D1, ..., Dn)`,
    `start_indices` must be a list/tuple of n integers, specifying the starting
    indices. `updates` must have the same rank as `inputs`, and the size of each
    dim must not exceed `Di - start_indices[i]`. For example, if we have 2D
    inputs `inputs = np.zeros((5, 5))`, and we want to update the intersection
    of last 2 rows and last 2 columns as 1, i.e.,
    `inputs[3:, 3:] = np.ones((2, 2))`, then we can use the code below:

    ```python
    inputs = np.zeros((5, 5))
    start_indices = [3, 3]
    updates = np.ones((2, 2))
    inputs = keras.ops.slice_update(inputs, start_indices, updates)
    ```

Args:
        inputs: A tensor, the tensor to be updated.
        start_indices: A list/tuple of shape `(inputs.ndim,)`, specifying
            the starting indices for updating.
        updates: A tensor, the new values to be put to `inputs` at `indices`.
            `updates` must have the same rank as `inputs`.

Returns:
        A tensor, has the same shape and dtype as `inputs`.