def matches(pcoll, regex, group=0):
    regex = Regex._regex_compile(regex)

    def _process(element):
        m = regex.match(element)
        if m:
            yield m.group(group)
    return pcoll | FlatMap(_process)

Returns the matches (group 0 by default) if zero or more characters at the
    beginning of string match the regular expression. To match the entire
    string, add "$" sign at the end of regex expression.

    Group can be integer value or a string value.

Args:
        regex: the regular expression string or (re.compile) pattern.
        group: (optional) name/number of the group, it can be integer or a string
        value. Defaults to 0, meaning the entire matched string will be
        returned.


def ConvertOutputToUnicode(self, buf):
    if isinstance(buf, str):
        buf = buf.encode(self._encoding)
    return str(buf, self._encoding, 'replace')

Converts a console output string buf to unicode.

    Mainly used for testing. Allows test comparisons in unicode while ensuring
    that unicode => encoding => unicode works.

Args:
        buf: The console output string to convert.

Returns:
        The console output string buf converted to unicode.


def if_false(self, predicate: Callable[..., bool]):
    return Conditional(predicate, None, self)

Conditionally applies current operation when predicate returns False.

Args:
        predicate: The predicate that takes the outputs from the previous
        operation as input, with optional keyword arguments `global_state` and
        `step`. Returns False if current operation needs to be enabled.
        Otherwise no operation will be performed.

Returns:
        A conditional operation.


def report_proto(self):
    if self._report_proto:
        return self._report_proto
    else:
        raise ValueError('Call to report_proto must be done after tracing.Report proto only exists for trace_mode=[summary|full_tensor_summary]')

Getter for tensor_tracer.proto object for summary and full_tensor_summary modes.

Returns:
        A tensor_tracer.proto object.

Raises:
        ValueError if called before tracing happens, or when trace mode is not
        summary or full_tensor_summary.


def __init__(self, c_in, c_hidden, no_heads, starting=True, inf=1000000000.0):
    super().__init__()
    self.c_in = c_in
    self.c_hidden = c_hidden
    self.no_heads = no_heads
    self.starting = starting
    self.inf = inf
    self.layer_norm = LayerNorm(self.c_in)
    self.linear = EsmFoldLinear(c_in, self.no_heads, bias=False, init='normal')
    self.mha = EsmFoldAttention(self.c_in, self.c_in, self.c_in, self.c_hidden, self.no_heads)

Args:
    c_in:
        Input channel dimension
    c_hidden:
        Overall hidden channel dimension (not per-head)
    no_heads:
        Number of attention heads


def forward(self, hidden_state):
    output_tensor = hidden_state
    output_tensor = self.conv1(self.snake1(output_tensor))
    output_tensor = self.conv2(self.snake2(output_tensor))
    padding = (hidden_state.shape[-1] - output_tensor.shape[-1]) // 2
    if padding > 0:
        hidden_state = hidden_state[..., padding:-padding]
    output_tensor = hidden_state + output_tensor
    return output_tensor

Forward pass through the residual unit.

Args:
        hidden_state (`torch.Tensor` of shape `(batch_size, channels, time_steps)`):
            Input tensor .

Returns:
        output_tensor (`torch.Tensor` of shape `(batch_size, channels, time_steps)`):
            Input tensor after passing through the residual unit.


def GetLogicalLines(self):
    self._StartNewLine()
    return self._logical_lines

Fetch the result of the tree walk.

    Note: only call this after visiting the whole tree.

Returns:
        A list of LogicalLine objects.


def parse_cluster_spec(cluster_spec, cluster, verbose=False):
    job_strings = cluster_spec.split(',')
    if not cluster_spec:
        raise ValueError('Empty cluster_spec string')
    for job_string in job_strings:
        job_def = cluster.job.add()
        if job_string.count('|') != 1:
            raise ValueError("Not exactly one instance of '|' in cluster_spec")
        job_name = job_string.split('|')[0]
        if not job_name:
            raise ValueError('Empty job_name in cluster_spec')
        job_def.name = job_name
        if verbose:
            logging.info('Added job named "%s"', job_name)
        job_tasks = job_string.split('|')[1].split(';')
        for i in range(len(job_tasks)):
            if not job_tasks[i]:
                raise ValueError('Empty task string at position %d' % i)
            job_def.tasks[i] = job_tasks[i]
            if verbose:
                logging.info('  Added task "%s" to job "%s"', job_tasks[i], job_name)

Parse content of cluster_spec string and inject info into cluster protobuf.

Args:
        cluster_spec: cluster specification string, e.g.,
            "local|localhost:2222;localhost:2223"
        cluster: cluster protobuf.
        verbose: If verbose logging is requested.

Raises:
        ValueError: if the cluster_spec string is invalid.


def __init__(self, weight_shape: Sequence[int], bias_size: Optional[int]=None, activation_fn: Optional[ops.Operation]=None, use_biasadd: bool=True) -> None:
    self.bias_size = bias_size
    self.activation_fn = activation_fn
    self.use_biasadd = use_biasadd
    self.filters = np.random.uniform(low=-1.0, high=1.0, size=weight_shape)
    if bias_size is not None:
        self.bias = np.random.uniform(low=-1.0, high=1.0, size=bias_size)

Initializes a MatmulModel.

Args:
        weight_shape: Shape of the weight tensor.
        bias_size: If None, do not use bias. Else, use given size as bias.
        activation_fn: The activation function to be used. No activation
        function if None.
        use_biasadd: If True, use BiasAdd for adding bias, else use AddV2.


def _ConstructAndTestGradient(self, image_shape, kernel_shape, strides, rates, padding, use_gpu, dtype=dtypes.float32):
    assert image_shape[3] == kernel_shape[2]
    np.random.seed(1)
    image = np.random.random_sample(image_shape).astype(np.float32)
    kernel = np.random.random_sample(kernel_shape).astype(np.float32)
    strides = [1] + strides + [1]
    rates = [1] + rates + [1]
    image_tensor = constant_op.constant(image, shape=image_shape, name='input', dtype=dtype)
    kernel_tensor = constant_op.constant(kernel, shape=kernel_shape, name='filter', dtype=dtype)

    def compute_dilation2d(image_tensor, kernel_tensor):
        return nn_ops.dilation2d(image_tensor, kernel_tensor, strides=strides, rates=rates, padding=padding, name='dilation2d')
    with test_util.device(use_gpu=use_gpu):
        with self.cached_session():
            err1 = gradient_checker_v2.max_error(*gradient_checker_v2.compute_gradient(lambda x: compute_dilation2d(x, kernel_tensor), [image_tensor]))
            err2 = gradient_checker_v2.max_error(*gradient_checker_v2.compute_gradient(lambda x: compute_dilation2d(image_tensor, x), [kernel_tensor]))
            err = max(err1, err2)
    print('Dilation gradient error = %f' % err)
    if dtype == dtypes.bfloat16:
        self.assertLess(err, 4.0)
    else:
        self.assertLess(err, 0.0001)

Verifies the gradients of the dilation function.

Args:
        image_shape: Input shape, [batch, in_height, in_width, channels].
        kernel_shape: Filter shape, [filter_height, filter_width, channels].
        strides: Output strides, specified as [stride_height, stride_width].
        rates: Atrous rates, specified as [rate_height, rate_width].
        padding: Padding type.
        use_gpu: Whether we are running on GPU.


def check_destinations(destinations):
    if isinstance(destinations, (resource_variable_ops.BaseResourceVariable, tensor_lib.Tensor)):
        return bool(destinations.device)
    return bool(destinations)

Checks whether `destinations` is not empty.

Args:
        destinations: a `DistributedValues`, variable, or string object.

Returns:
        Boolean which is True if `destinations` is not empty.


def get_model_class_for_feature(feature: str, framework: str='pt') -> Type:
    task = FeaturesManager.feature_to_task(feature)
    FeaturesManager._validate_framework_choice(framework)
    if framework == 'pt':
        task_to_automodel = FeaturesManager._TASKS_TO_AUTOMODELS
    else:
        task_to_automodel = FeaturesManager._TASKS_TO_TF_AUTOMODELS
    if task not in task_to_automodel:
        raise KeyError(f'Unknown task: {feature}. Possible values are {list(FeaturesManager._TASKS_TO_AUTOMODELS.values())}')
    return task_to_automodel[task]

Attempts to retrieve an AutoModel class from a feature name.

Args:
        feature (`str`):
            The feature required.
        framework (`str`, *optional*, defaults to `"pt"`):
            The framework to use for the export.

Returns:
        The AutoModel class corresponding to the feature.


def get_tensor_by_name(self, name) -> tensor_lib.Tensor:
    if not isinstance(name, str):
        raise TypeError('Tensor names are strings (or similar), not %s.' % type(name).__name__)
    tensor = cast(tensor_lib.Tensor, self.as_graph_element(name, allow_tensor=True, allow_operation=False))
    return tensor

Returns the `Tensor` with the given `name`.

    This method may be called concurrently from multiple threads.

Args:
        name: The name of the `Tensor` to return.

Returns:
        The `Tensor` with the given `name`.

Raises:
        TypeError: If `name` is not a string.
        KeyError: If `name` does not correspond to a tensor in this graph.


def _ValidateFractionalAvgPoolResult(self, input_tensor, pooling_ratio, pseudo_random, overlapping):
    with self.cached_session() as sess:
        p, r, c = nn_ops.fractional_avg_pool_v2(input_tensor, pooling_ratio, pseudo_random, overlapping, seed=self._SEED)
        actual, row_seq, col_seq = self.evaluate([p, r, c])
        expected = self._GetExpectedFractionalAvgPoolResult(input_tensor, row_seq, col_seq, overlapping)
        self.assertShapeEqual(expected, p)
        self.assertAllClose(expected, actual)

Validate FractionalAvgPool's result against expected.

    Expected result is computed given input_tensor, and pooling region defined
    by row_seq and col_seq.

Args:
        input_tensor: A tensor or numpy ndarray.
        pooling_ratio: A list or tuple of length 4, first and last element be 1.
        pseudo_random: Use pseudo random method to generate pooling sequence.
        overlapping: Use overlapping when pooling.

Returns:
        None


def _init_init_op(self, init_op=USE_DEFAULT, init_feed_dict=None):
    if init_op is Supervisor.USE_DEFAULT:
        init_op = self._get_first_op_from_collection(ops.GraphKeys.INIT_OP)
        if init_op is None:
            init_op = variables.global_variables_initializer()
            ops.add_to_collection(ops.GraphKeys.INIT_OP, init_op)
    self._init_op = init_op
    self._init_feed_dict = init_feed_dict

Initializes init_op.

Args:
        init_op: `Operation` to initialize the variables. If set to USE_DEFAULT,
        create an op that initializes all variables and tables.
        init_feed_dict: A dictionary that maps `Tensor` objects to feed values.
        This feed dictionary will be used when `init_op` is evaluated.


def _compute_dynamic_ntk_parameters(config: Optional[PretrainedConfig]=None, device: Optional['torch.device']=None, seq_len: Optional[int]=None, **rope_kwargs) -> tuple['torch.Tensor', float]:
    if config is not None and len(rope_kwargs) > 0:
        raise ValueError(f'Unexpected arguments: `**rope_kwargs` and `config` are mutually exclusive in `_compute_dynamic_ntk_parameters`, got `rope_kwargs`={rope_kwargs} and `config`={config}')
    if len(rope_kwargs) > 0:
        base = rope_kwargs['base']
        dim = rope_kwargs['dim']
        max_position_embeddings = rope_kwargs['max_position_embeddings']
        factor = rope_kwargs['factor']
    elif config is not None:
        base = config.rope_theta
        partial_rotary_factor = config.partial_rotary_factor if hasattr(config, 'partial_rotary_factor') else 1.0
        head_dim = getattr(config, 'head_dim', config.hidden_size // config.num_attention_heads)
        dim = int(head_dim * partial_rotary_factor)
        max_position_embeddings = config.max_position_embeddings
        factor = config.rope_scaling['factor']
    attention_factor = 1.0
    seq_len = seq_len if seq_len is not None and seq_len > max_position_embeddings else max_position_embeddings
    base = base * (factor * seq_len / max_position_embeddings - (factor - 1)) ** (dim / (dim - 2))
    inv_freq = 1.0 / base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
    return (inv_freq, attention_factor)

Computes the inverse frequencies with NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla

Args:
        config ([`~transformers.PretrainedConfig`]):
            The model configuration.
        device (`torch.device`):
            The device to use for initialization of the inverse frequencies.
        seq_len (`int`, *optional*):
            The current sequence length, used to update the dynamic RoPE at inference time.
        rope_kwargs (`Dict`, *optional*):
            BC compatibility with the previous RoPE class instantiation, will be removed in v4.45.

Returns:
        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
        post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).


def __init__(self, supported_ops=None, supported_types=None, experimental_select_user_tf_ops=None, experimental_supported_backends=None):
    if supported_ops is None:
        supported_ops = {OpsSet.TFLITE_BUILTINS}
    self.supported_ops = supported_ops
    if supported_types is None:
        supported_types = set()
    self.supported_types = supported_types
    if experimental_select_user_tf_ops is None:
        experimental_select_user_tf_ops = set()
    self.experimental_select_user_tf_ops = experimental_select_user_tf_ops
    self.experimental_supported_backends = experimental_supported_backends
    self._experimental_custom_op_registerers = []
    self._experimental_supported_accumulation_type = None

Specification of target device used to optimize the model.

Attributes:
        supported_ops: Experimental flag, subject to change. Set of `tf.lite.OpsSet`
        options, where each option represents a set of operators supported by the
        target device. (default {tf.lite.OpsSet.TFLITE_BUILTINS}))
        supported_types: Set of `tf.dtypes.DType` data types supported on the target
        device. If initialized, optimization might be driven by the smallest type
        in this set. (default set())
        experimental_select_user_tf_ops: Experimental flag, subject to change. Set
        of user's TensorFlow operators' names that are required in the TensorFlow
        Lite runtime. These ops will be exported as select TensorFlow ops in the
        model (in conjunction with the tf.lite.OpsSet.SELECT_TF_OPS flag). This is
        an advanced feature that should only be used if the client is using TF ops
        that may not be linked in by default with the TF ops that are provided
        when using the SELECT_TF_OPS path. The client is responsible for linking
        these ops into the target runtime.
        experimental_supported_backends: Experimental flag, subject to change. Set
        containing names of supported backends. Currently only "GPU" is supported,
        more options will be available later.


def _parse_logline_timestamp(t):
    date, time = t.split(' ')
    month, day = date.split('-')
    h, m, s = time.split(':')
    s, ms = s.split('.')
    return (month, day, h, m, s, ms)

Parses a logline timestamp into a tuple.

Args:
        t: Timestamp in logline format.

Returns:
        An iterable of date and time elements in the order of month, day, hour,
        minute, second, microsecond.


def extended_timestamp(self):
    return self._extended_timestamp

Extended timestamp, possibly with an index suffix.

    The index suffix, e.g., "-1", is for disambiguating multiple dumps of the
    same tensor with the same timestamp, which can occur if the dumping events
    are spaced by shorter than the temporal resolution of the timestamps.

Returns:
        (`str`) The extended timestamp.


def parse_readable_size_str(size_str):
    size_str = size_str.strip()
    if size_str.endswith('B'):
        size_str = size_str[:-1]
    if size_str.isdigit():
        return int(size_str)
    elif size_str.endswith('k'):
        return int(float(size_str[:-1]) * 1024)
    elif size_str.endswith('M'):
        return int(float(size_str[:-1]) * 1048576)
    elif size_str.endswith('G'):
        return int(float(size_str[:-1]) * 1073741824)
    else:
        raise ValueError('Failed to parsed human-readable byte size str: "%s"' % size_str)

Convert a human-readable str representation to number of bytes.

    Only the units "kB", "MB", "GB" are supported. The "B character at the end
    of the input `str` may be omitted.

Args:
        size_str: (`str`) A human-readable str representing a number of bytes
        (e.g., "0", "1023", "1.1kB", "24 MB", "23GB", "100 G".

Returns:
        (`int`) The parsed number of bytes.

Raises:
        ValueError: on failure to parse the input `size_str`.


def __init__(self, element_spec):
    super().__init__()
    self._element_spec = element_spec

Type specification for `tf.experimental.Optional`.

    For instance, `tf.OptionalSpec` can be used to define a tf.function that takes
    `tf.experimental.Optional` as an input argument:

    >>> @tf.function(input_signature=[tf.OptionalSpec(
    ...   tf.TensorSpec(shape=(), dtype=tf.int32, name=None))])
    ... def maybe_square(optional):
    ...   if optional.has_value():
    ...     x = optional.get_value()
    ...     return x * x
    ...   return -1
    >>> optional = tf.experimental.Optional.from_value(5)
    >>> print(maybe_square(optional))
    tf.Tensor(25, shape=(), dtype=int32)

Attributes:
        element_spec: A (nested) structure of `TypeSpec` objects that represents the
        type specification of the optional element.


def __init__(self, config: AriaTextConfig):
    super().__init__(self)
    self.intermediate_size = config.intermediate_size * config.moe_num_shared_experts

Shared Expert MLP for shared experts.

    Unlike routed experts, shared experts process all tokens without routing.
    This class reconfigures the intermediate size in comparison to the LlamaMLP.

Args:
        config (`AriaTextConfig`): Configuration object for the Aria language model.


def get(self):
    raise NotImplementedError()

Get the current tracking value.

Returns:
        The current tracked value, the type of which depends on the specific
        tracker implementation.


def _einsum_v1_parse_and_resolve_equation(equation, input_shapes):
    equation = equation.replace(' ', '')
    match = re.match('^([a-zA-Z,.]+)(->[a-zA-Z.]*)?$', equation)
    if not match:
        raise ValueError(f'Indices have incorrect format. Received: {equation}.')
    input_axis_labels = match.group(1).split(',')
    output_axis_labels = match.group(2)[2:] if match.group(2) else None
    if len(input_shapes) != len(input_axis_labels):
        raise ValueError(f'Got {len(input_shapes)} arguments for equation "{equation}", expecting {len(input_axis_labels)}.')
    ellipsis_axes = ''
    if '...' in equation:
        unused = ''.join((c for c in string.ascii_letters if c not in ''.join(input_axis_labels)))
        for i, ax in enumerate(input_axis_labels):
            if '...' in ax:
                parts = ax.split('...')
                if len(parts) != 2:
                    raise ValueError(f'Unable to resolve ellipsis. Excess number found: {len(parts) - 1} vs 1.')
                if input_shapes[i].ndims is None:
                    raise ValueError('Unable to statically infer ellipsis axes. The input shapes has a dynamic dimensionality.')
                n = input_shapes[i].ndims - len(''.join(parts))
                if n < 0:
                    raise ValueError('Ellipses lengths do not match.')
                if len(unused) < n:
                    raise ValueError('Unable to resolve ellipsis, too many distinct labels.')
                replace_axes = unused[-n:] if n > 0 else ''
                input_axis_labels[i] = input_axis_labels[i].replace('...', replace_axes)
                if len(replace_axes) > len(ellipsis_axes):
                    ellipsis_axes = replace_axes
        if any(('.' in ax for ax in input_axis_labels)):
            raise ValueError(f'Period "." found outside of ellipsis in input {input_axis_labels}.')
        if output_axis_labels is not None:
            output_axis_labels = output_axis_labels.replace('...', ellipsis_axes)
            if '.' in output_axis_labels:
                raise ValueError(f'Period "." found outside of ellipsis in output {output_axis_labels}.')
    if output_axis_labels is None:
        axis_labels = set(''.join(input_axis_labels)) - set(ellipsis_axes)
        indices = ''.join(sorted(axis_labels))
        counts = {ax: 0 for ax in indices}
        for axes_ in input_axis_labels:
            for ax in axes_:
                if ax not in ellipsis_axes:
                    counts[ax] += 1
        output_axis_labels = ellipsis_axes + ''.join(sorted((ax for ax in axis_labels if counts[ax] == 1)))
    return (input_axis_labels, output_axis_labels)

Helper for einsum() that splits/resolves inputs & outputs.

Args:
        equation: Equation string given as argument to einsum().
        input_shapes: List of the shapes of all inputs given to einsum()

Returns:
        input_axis_labels, output_axis_labels where:
        input_axis_labels: List of length len(input_shapes) of strings
        representing the character label for each dimension of each given input,
        resolving any broadcast (...) axes,
        output_axis_labels: A string of character labels for each axes of output
        tensor, filling in missing output subscripts and broadcast axes.

Raises:
        ValueError: If equation is in the uncorrect format, incorrect number of
        inputs given or broadcast axes "..." or output axes could not be resolved.


def train_validation_split(arrays, validation_split):

    def _can_split(t):
        tensor_types = _get_tensor_types()
        return isinstance(t, tensor_types) or t is None
    flat_arrays = nest.flatten(arrays)
    unsplitable = [type(t) for t in flat_arrays if not _can_split(t)]
    if unsplitable:
        raise ValueError('`validation_split` is only supported for Tensors or NumPy arrays, found following types in the input: {}'.format(unsplitable))
    if all((t is None for t in flat_arrays)):
        return (arrays, arrays)
    first_non_none = None
    for t in flat_arrays:
        if t is not None:
            first_non_none = t
            break
    batch_dim = int(first_non_none.shape[0])
    split_at = int(math.floor(batch_dim * (1.0 - validation_split)))
    if split_at == 0 or split_at == batch_dim:
        raise ValueError('Training data contains {batch_dim} samples, which is not sufficient to split it into a validation and training set as specified by `validation_split={validation_split}`. Either provide more data, or a different value for the `validation_split` argument.'.format(batch_dim=batch_dim, validation_split=validation_split))

    def _split(t, start, end):
        if t is None:
            return t
        return t[start:end]
    train_arrays = nest.map_structure(functools.partial(_split, start=0, end=split_at), arrays)
    val_arrays = nest.map_structure(functools.partial(_split, start=split_at, end=batch_dim), arrays)
    return (train_arrays, val_arrays)

Split arrays into train and validation subsets in deterministic order.

    The last part of data will become validation data.

Args:
        arrays: Tensors to split. Allowed inputs are arbitrarily nested structures
        of Tensors and NumPy arrays.
        validation_split: Float between 0 and 1. The proportion of the dataset to
        include in the validation split. The rest of the dataset will be included
        in the training split.

Returns:
        `(train_arrays, validation_arrays)`


def __init__(self, config: MaskFormerConfig):
    super().__init__()
    if getattr(config, 'backbone_config') is not None and config.backbone_config.model_type == 'swin':
        backbone_config = config.backbone_config
        backbone_config = MaskFormerSwinConfig.from_dict(backbone_config.to_dict())
        backbone_config.out_features = ['stage1', 'stage2', 'stage3', 'stage4']
        config.backbone_config = backbone_config
    self.encoder = load_backbone(config)
    feature_channels = self.encoder.channels
    self.decoder = MaskFormerPixelDecoder(in_features=feature_channels[-1], feature_size=config.fpn_feature_size, mask_feature_size=config.mask_feature_size, lateral_widths=feature_channels[:-1])

Pixel Level Module proposed in [Per-Pixel Classification is Not All You Need for Semantic
    Segmentation](https://huggingface.co/papers/2107.06278). It runs the input image through a backbone and a pixel
    decoder, generating an image feature map and pixel embeddings.

Args:
        config ([`MaskFormerConfig`]):
            The configuration used to instantiate this model.


def correct_tables(self, generation: str) -> str:
    for l in generation.split('\n'):
        if l.count('\\begin{tabular}') > 15 or l.count('\\multicolumn') > 60 or l.count('&') > 400:
            generation = generation.replace(l, '')
    generation = generation.replace('\\begin{table} \\begin{tabular}', '\\begin{table}\n\\begin{tabular}')
    generation = generation.replace('\\end{tabular} \\end{table}', '\\end{tabular}\n\\end{table}')
    generation = generation.replace('\\end{table} Tab', '\\end{table}\nTab')
    generation = re.sub('(^.+)\\\\begin{tab', '\\1\\n\\\\begin{tab', generation, flags=re.M)
    generation = generation.replace('\\begin{tabular}{l l}  & \\\\ \\end{tabular}', '')
    generation = generation.replace('\\begin{tabular}{}\n\n\\end{tabular}', '')
    return generation

Takes a generated string and fixes tables/tabulars to make them match the markdown format needed.

Args:
        generation (str): The generated text to be postprocessed.

Returns:
        str: The postprocessed text.

Example:
    ```python
    correct_tables("\begin{table} \begin{tabular}{l l} & \ \end{tabular} \end{table}")
    "\begin{table}
    \begin{tabular}{l l} & \ \end{tabular}
    \end{table}"
    ```


def __call__(self, raw_speech: Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]], padding: Union[bool, str, PaddingStrategy]=False, max_length: Optional[int]=None, pad_to_multiple_of: Optional[int]=None, padding_side: Optional[str]=None, return_tensors: Optional[Union[str, TensorType]]=None, verbose: bool=True, **kwargs) -> BatchEncoding:
    is_batched_numpy = isinstance(raw_speech, np.ndarray) and len(raw_speech.shape) > 1
    if is_batched_numpy and len(raw_speech.shape) > 2:
        raise ValueError(f'Only mono-channel audio is supported for input to {self}')
    is_batched = is_batched_numpy or (isinstance(raw_speech, (list, tuple)) and isinstance(raw_speech[0], (np.ndarray, tuple, list)))
    if is_batched and (not isinstance(raw_speech[0], np.ndarray)):
        raw_speech = [np.asarray(speech) for speech in raw_speech]
    elif not is_batched and (not isinstance(raw_speech, np.ndarray)):
        raw_speech = np.asarray(raw_speech)
    if not is_batched:
        raw_speech = [raw_speech]
    if self.do_normalize:
        raw_speech = [(x - np.mean(x)) / np.sqrt(np.var(x) + 1e-05) for x in raw_speech]
    encoded_inputs = BatchEncoding({'input_values': raw_speech})
    padded_inputs = self.pad(encoded_inputs, padding=padding, max_length=max_length, pad_to_multiple_of=pad_to_multiple_of, padding_side=padding_side, return_attention_mask=self.return_attention_mask, return_tensors=return_tensors, verbose=verbose)
    return padded_inputs

Main method to tokenize and prepare for the model one or several sequence(s) or one or several pair(s) of
    sequences.

Args:
        raw_speech (`np.ndarray`, `List[float]`, `List[np.ndarray]`, `List[List[float]]`):
            The sequence or batch of sequences to be padded. Each sequence can be a numpy array, a list of float
            values, a list of numpy array or a list of list of float values. Must be mono channel audio, not
            stereo, i.e. single float per timestep.

        padding_side (`str`, *optional*):
            The side on which the model should have padding applied. Should be selected between ['right', 'left'].
            Default value is picked from the class attribute of the same name.


def _parse_args():
    parser = argparse.ArgumentParser(description='preprocess_coco_minival: Preprocess COCO minival dataset')
    parser.add_argument('--images_folder', type=str, help='Full path of the validation images folder.', required=True)
    parser.add_argument('--instances_file', type=str, help='Full path of the input JSON file, like instances_val20xx.json.', required=True)
    parser.add_argument('--allowlist_file', type=str, help='File with COCO image ids to preprocess, one on each line.', required=False)
    parser.add_argument('--num_images', type=int, help='Number of allowlisted images to preprocess into the output folder.', required=False)
    parser.add_argument('--output_folder', type=str, help='Full path to output images & text proto files into.', required=True)
    return parser.parse_known_args(args=sys.argv[1:])[0]

Creates a parser that parse the command line arguments.

Returns:
        A namespace parsed from command line arguments.


def _section_from_possible_title(possible_title):
    for section in SECTION_TITLES:
        if _matches_section(possible_title, section):
            return section
    return None

Returns a section matched by the possible title, or None if none match.

Args:
        possible_title: A string that may be the title of a new section.

Returns:
        A Section type if one matches, or None if no section type matches.


def __init__(self, *, separator_stride_cls: Type[message.Message], code_cls: Type[message.Message], default_timezone: str) -> None:
    self.separator_stride_cls = separator_stride_cls
    self.code_cls = code_cls
    self.default_timezone = default_timezone

Creates a new instance of primitive_wrappers.Context.

Args:
        separator_stride_cls: The Base64BinarySeparatorStride type to use when
        parsing/printing Base64Binary FHIR primitives.
        code_cls: The Code type to use when parsing/printing profiled-Code
        primitives.
        default_timezone: The default timezone to use for date/time-like primitive
        parsing/printing.


def diff_contains_doc_examples(repo: Repo, branching_point: str, filename: str) -> bool:
    folder = Path(repo.working_dir)
    with checkout_commit(repo, branching_point):
        with open(folder / filename, 'r', encoding='utf-8') as f:
            old_content = f.read()
    with open(folder / filename, 'r', encoding='utf-8') as f:
        new_content = f.read()
    old_content_clean = keep_doc_examples_only(old_content)
    new_content_clean = keep_doc_examples_only(new_content)
    return old_content_clean != new_content_clean

Check if the diff is only in code examples of the doc in a filename.

Args:
        repo (`git.Repo`): A git repository (for instance the Transformers repo).
        branching_point (`str`): The commit reference of where to compare for the diff.
        filename (`str`): The filename where we want to know if the diff is only in codes examples.

Returns:
        `bool`: Whether the diff is only in code examples of the doc or not.


def write_object_proto_for_resource_variable(resource_variable, proto, options, enforce_naming=True):
    proto.variable.SetInParent()
    if enforce_naming and (not resource_variable.name.endswith(':0')):
        raise ValueError(f"Cowardly refusing to save variable {resource_variable.name} because of unexpected suffix in the name (expected ':0')which won't be restored.")
    proto.variable.name = tensor_module.get_op_name(resource_variable.name)
    proto.variable.trainable = resource_variable.trainable
    proto.variable.dtype = resource_variable.dtype.as_datatype_enum
    proto.variable.synchronization = resource_variable.synchronization.value
    proto.variable.aggregation = resource_variable.aggregation.value
    proto.variable.shape.CopyFrom(resource_variable.shape.as_proto())
    if options.experimental_variable_policy._save_variable_devices():
        if hasattr(resource_variable, 'device'):
            proto.variable.device = resource_variable.device

Writes additional information of the variable into the SavedObject proto.

    This allows users to define a `hook` to provide extra information of the
    variable to the SavedObject.

    For example, DistributedVariable class would fill in components in the
    distributed context.

Args:
        resource_variable: A `ResourceVariable` or `DistributedValue` that has the
        information to be saved into the proto.
        proto: `SavedObject` proto to update.
        options: A `SaveOption` instance that configures save behavior.
        enforce_naming: A bool determining whether to check that names end in the
        expected string ':0'


def get(logdir):
    with FileWriterCache._lock:
        if logdir not in FileWriterCache._cache:
            FileWriterCache._cache[logdir] = FileWriter(logdir, graph=ops.get_default_graph())
        return FileWriterCache._cache[logdir]

Returns the FileWriter for the specified directory.

Args:
        logdir: str, name of the directory.

Returns:
        A `FileWriter`.


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

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

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

Returns:
        decomposed_rel_pos (`torch.Tensor`):
            decomposed relative position embeddings.


def _call_wrapped_cell(self, inputs, state, cell_call_fn, **kwargs):

    def _should_dropout(p):
        return not isinstance(p, float) or p < 1
    if _should_dropout(self._input_keep_prob):
        inputs = self._dropout(inputs, 'input', self._recurrent_input_noise, self._input_keep_prob)
    output, new_state = cell_call_fn(inputs, state, **kwargs)
    if _should_dropout(self._state_keep_prob):
        shallow_filtered_substructure = nest.get_traverse_shallow_structure(self._dropout_state_filter, new_state)
        new_state = self._dropout(new_state, 'state', self._recurrent_state_noise, self._state_keep_prob, shallow_filtered_substructure)
    if _should_dropout(self._output_keep_prob):
        output = self._dropout(output, 'output', self._recurrent_output_noise, self._output_keep_prob)
    return (output, new_state)

Runs the wrapped cell and applies dropout.

Args:
        inputs: A tensor with wrapped cell's input.
        state: A tensor or tuple of tensors with wrapped cell's state.
        cell_call_fn: Wrapped cell's method to use for step computation (cell's
        `__call__` or 'call' method).
        **kwargs: Additional arguments.

Returns:
        A pair containing:

        - Output: A tensor with cell's output.
        - New state: A tensor or tuple of tensors with new wrapped cell's state.


def decode(self, music_tokens, start_level=0, end_level=None, bs_chunks=1) -> torch.Tensor:
    token_chunks = [torch.chunk(token, bs_chunks, dim=0) for token in music_tokens]
    dequantised_states = []
    for i in range(bs_chunks):
        music_tokens_i = [chunks[i] for chunks in token_chunks]
        dequantised_state = self._decode(music_tokens_i, start_level=start_level, end_level=end_level)
        dequantised_states.append(dequantised_state)
    return torch.cat(dequantised_states, dim=0)

Transforms the input `music_tokens` to their `raw_audio` representation.

Args:
        music_tokens (`torch.LongTensor`):
            Tensor of music tokens which will be decoded to raw audio by using the codebook. Each music token
            should be an index to a corresponding `code` vector in the codebook.
        start_level (`int`, *optional*):
            Level at which the decoding process will start. Default to 0.
        end_level (`int`, *optional*):
            Level at which the decoding process will start. Default to None.
        bs_chunks (int, *optional*):
            Number of chunks to process at the same time.


def _print_tensor(tensor_name, num_elements, tensor, output_tensor):
    if self._parameters.is_brief_mode():
        if tensor_name not in tensor_trace_order.tensorname_to_cache_idx:
            raise ValueError('Tensor %s with name %s is not in the tensorname_to_cache_idx' % (tensor, tensor_name))
        msg = '%d' % tensor_trace_order.tensorname_to_cache_idx[tensor_name]
    else:
        msg = '"%s"' % tensor_name
    if self._parameters.trace_dir:
        output_path = os.path.join(self._parameters.trace_dir, _TRACE_FILE_NAME + self._get_outfile_suffix())
        output_stream = _OUTPUT_STREAM_ESCAPE + output_path
    else:
        output_stream = sys.stderr
    return logging_ops.print_v2(msg, array_ops.shape(output_tensor), '@', self._replica_id, '\n', output_tensor, '\n', summarize=num_elements, output_stream=output_stream)

Prints a tensor value to a file.

Args:
        tensor_name: name of the tensor being traced.
        num_elements: number of elements to print (-1 means print all).
        tensor: the tensor needs to be returned.
        output_tensor: the tensor needs to be printed.

Returns:
        The same tensor passed via the "tensor" argument.

Raises:
        ValueError: If tensor_name is not already in
                tensor_trace_order.tensorname_to_cache_idx.


def user_set_string(namespace, name, metric, ptransform=None):
    labels = create_labels(ptransform=ptransform, namespace=namespace, name=name)
    if isinstance(metric, StringSetData):
        metric = metric.string_set
    if isinstance(metric, set):
        metric = list(metric)
    if isinstance(metric, list):
        metric = coders.IterableCoder(coders.StrUtf8Coder()).encode(metric)
    return create_monitoring_info(USER_STRING_SET_URN, STRING_SET_TYPE, metric, labels)

Return the string set monitoring info for the URN, metric and labels.

Args:
        namespace: User-defined namespace of StringSet.
        name: Name of StringSet.
        metric: The StringSetData representing the metrics.
        ptransform: The ptransform id used as a label.


def weighted_moving_average(value, decay, weight, truediv=True, collections=None, name=None):
    if collections is None:
        collections = [ops.GraphKeys.GLOBAL_VARIABLES]
    with variable_scope.variable_scope(name, 'WeightedMovingAvg', [value, weight, decay]) as scope:
        value_x_weight_var = variable_scope.get_variable('value_x_weight', shape=value.get_shape(), dtype=value.dtype, initializer=init_ops.zeros_initializer(), trainable=False, collections=collections)
        weight_var = variable_scope.get_variable('weight', shape=weight.get_shape(), dtype=weight.dtype, initializer=init_ops.zeros_initializer(), trainable=False, collections=collections)
        numerator = assign_moving_average(value_x_weight_var, value * weight, decay, zero_debias=False)
        denominator = assign_moving_average(weight_var, weight, decay, zero_debias=False)
        if truediv:
            return math_ops.truediv(numerator, denominator, name=scope.name)
        else:
            return math_ops.divide(numerator, denominator, name=scope.name)

Compute the weighted moving average of `value`.

    Conceptually, the weighted moving average is:
        `moving_average(value * weight) / moving_average(weight)`,
    where a moving average updates by the rule
        `new_value = decay * old_value + (1 - decay) * update`
    Internally, this Op keeps moving average variables of both `value * weight`
    and `weight`.

Args:
        value: A numeric `Tensor`.
        decay: A float `Tensor` or float value. The moving average decay.
        weight:  `Tensor` that keeps the current value of a weight. Shape should be
        able to multiply `value`.
        truediv:  Boolean, if `True`, dividing by `moving_average(weight)` is
        floating point division.  If `False`, use division implied by dtypes.
        collections:  List of graph collections keys to add the internal variables
        `value * weight` and `weight` to. Defaults to
        `[GraphKeys.GLOBAL_VARIABLES]`.
        name: Optional name of the returned operation. Defaults to
        "WeightedMovingAvg".

Returns:
        An Operation that updates and returns the weighted moving average.


def _get_numeric_relations(self, question, column_ids, row_ids, table):
    numeric_relations = [0] * len(column_ids)
    cell_indices_to_relations = collections.defaultdict(set)
    if question is not None and table is not None:
        for numeric_value_span in question.numeric_spans:
            for value in numeric_value_span.values:
                for column_index in range(len(table.columns)):
                    table_numeric_values = self._get_column_values(table, column_index)
                    sort_key_fn = self._get_numeric_sort_key_fn(table_numeric_values, value)
                    if sort_key_fn is None:
                        continue
                    for row_index, cell_value in table_numeric_values.items():
                        relation = get_numeric_relation(value, cell_value, sort_key_fn)
                        if relation is not None:
                            cell_indices_to_relations[column_index, row_index].add(relation)
    for (column_index, row_index), relations in cell_indices_to_relations.items():
        relation_set_index = 0
        for relation in relations:
            assert relation.value >= Relation.EQ.value
            relation_set_index += 2 ** (relation.value - Relation.EQ.value)
        for cell_token_index in self._get_cell_token_indexes(column_ids, row_ids, column_index, row_index):
            numeric_relations[cell_token_index] = relation_set_index
    return numeric_relations

Returns numeric relations embeddings

Args:
        question: Question object.
        column_ids: Maps word piece position to column id.
        row_ids: Maps word piece position to row id.
        table: The table containing the numeric cell values.


def cosine_similarity(y_true, y_pred, axis=-1):
    y_pred = ops.convert_to_tensor(y_pred)
    y_true = ops.convert_to_tensor(y_true, dtype=y_pred.dtype)
    y_true, y_pred = squeeze_or_expand_to_same_rank(y_true, y_pred)
    y_pred = normalize(y_pred, axis=axis)
    y_true = normalize(y_true, axis=axis)
    return -ops.sum(y_true * y_pred, axis=axis)

Computes the cosine similarity between labels and predictions.

    Formula:
    ```python
    loss = -sum(l2_norm(y_true) * l2_norm(y_pred))
    ```

    Note that it is a number between -1 and 1. When it is a negative number
    between -1 and 0, 0 indicates orthogonality and values closer to -1
    indicate greater similarity. This makes it usable as a loss function in a
    setting where you try to maximize the proximity between predictions and
    targets. If either `y_true` or `y_pred` is a zero vector, cosine
    similarity will be 0 regardless of the proximity between predictions
    and targets.

Args:
        y_true: Tensor of true targets.
        y_pred: Tensor of predicted targets.
        axis: Axis along which to determine similarity. Defaults to `-1`.

Returns:
        Cosine similarity tensor.

Example:
    >>> y_true = [[0., 1.], [1., 1.], [1., 1.]]
    >>> y_pred = [[1., 0.], [1., 1.], [-1., -1.]]
    >>> loss = keras.losses.cosine_similarity(y_true, y_pred, axis=-1)
    [-0., -0.99999994, 0.99999994]


def convert_to_date_tensor(date_inputs):
    if isinstance(date_inputs, DateTensor):
        return date_inputs
    if hasattr(date_inputs, 'year'):
        return from_datetimes(date_inputs)
    if isinstance(date_inputs, np.ndarray):
        date_inputs = date_inputs.astype('datetime64[D]')
        return from_np_datetimes(date_inputs)
    if tf.is_tensor(date_inputs):
        return from_ordinals(date_inputs)
    if isinstance(date_inputs, collections.abc.Sequence):
        if not date_inputs:
            return from_ordinals([])
        test_element = date_inputs[0]
        if hasattr(test_element, 'year'):
            return from_datetimes(date_inputs)
        if isinstance(test_element, collections.abc.Sequence):
            return from_tuples(date_inputs)
        if len(date_inputs) == 3:
            return from_year_month_day(date_inputs[0], date_inputs[1], date_inputs[2])
    try:
        as_ordinals = tf.convert_to_tensor(date_inputs, dtype=tf.int32)
        return from_ordinals(as_ordinals)
    except ValueError as e:
        raise ValueError('Failed to convert inputs to DateTensor. Unrecognized format. Error: ' + e)

Converts supplied data to a `DateTensor` if possible.

Args:
        date_inputs: One of the supported types that can be converted to a
        DateTensor. The following input formats are supported. 1. Sequence of
        `datetime.datetime`, `datetime.date`, or any other structure with data
        attributes called 'year', 'month' and 'day'. 2. A numpy array of
        `datetime64` type. 3. Sequence of (year, month, day) Tuples. Months are
        1-based (with January as 1) and constants.Months enum may be used instead
        of ints. Days are also 1-based. 4. A tuple of three int32 `Tensor`s
        containing year, month and date as positive integers in that order. 5. A
        single int32 `Tensor` containing ordinals (i.e. number of days since 31
        Dec 0 with 1 being 1 Jan 1.)

Returns:
        A `DateTensor` object representing the supplied dates.

Raises:
        ValueError: If conversion fails for any reason.


def query_origin_stack(self):
    ret = []
    for stack, id_to_string in zip(self._origin_stacks, self._origin_id_to_strings):
        ret.append(self._code_def_to_traceback(stack, id_to_string))
    return ret

Query the stack of the origin of the execution call.

Returns:
        A `list` of all tracebacks. Each item corresponds to an execution call,
        i.e., a `SendTracebacks` request. Each item is a `list` of 3-tuples:
        (filename, lineno, function_name).


def window_partition(self, hidden_states: torch.Tensor, window_size: int) -> Tuple[torch.Tensor, Tuple[int, int]]:
    batch_size, height, width, channel = hidden_states.shape
    pad_h = (window_size - height % window_size) % window_size
    pad_w = (window_size - width % window_size) % window_size
    hidden_states = F.pad(hidden_states, (0, 0, 0, pad_w, 0, pad_h))
    pad_height, pad_width = (height + pad_h, width + pad_w)
    hidden_states = hidden_states.reshape(batch_size, pad_height // window_size, window_size, pad_width // window_size, window_size, channel)
    windows = hidden_states.permute(0, 1, 3, 2, 4, 5).contiguous().reshape(-1, window_size, window_size, channel)
    return (windows, (pad_height, pad_width))

Args:
    Partition into non-overlapping windows with padding if needed.
        hidden_states (tensor): input tokens with [batch_size, height, width, channel]. window_size (int): window
        size.

Returns:
        windows: windows after partition with [batch_size * num_windows, window_size, window_size, channel].
        (pad_height, pad_width): padded height and width before partition


def export_run_metadata(self):
    if not self._context_handle:
        return None
    with c_api_util.tf_buffer() as buffer_:
        pywrap_tfe.TFE_ContextExportRunMetadata(self._context_handle, buffer_)
        proto_data = pywrap_tf_session.TF_GetBuffer(buffer_)
    run_metadata = config_pb2.RunMetadata()
    run_metadata.ParseFromString(compat.as_bytes(proto_data))
    return run_metadata

Returns a RunMetadata proto with accumulated information.

    The returned protocol buffer contains information since the most recent call
    to either enable_run_metadata or export_run_metadata.

Returns:
        A RunMetadata protocol buffer. Or None if not enabled.


def add_loss(loss, loss_collection=ops.GraphKeys.LOSSES):
    if loss_collection and (not context.executing_eagerly()):
        ops.add_to_collection(loss_collection, loss)

Adds a externally defined loss to the collection of losses.

Args:
        loss: A loss `Tensor`.
        loss_collection: Optional collection to add the loss to.


def last_updated(self, path):
    try:
        return self._blobstorageIO().last_updated(path)
    except Exception as e:
        raise BeamIOError('Last updated operation failed', {path: e})

Get UNIX Epoch time in seconds on the FileSystem.

Args:
        path: string path of file.

    Returns: float UNIX Epoch time

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


def sequence_categorical_column_with_hash_bucket(key, hash_bucket_size, dtype=dtypes.string):
    return fc.SequenceCategoricalColumn(fc.categorical_column_with_hash_bucket(key=key, hash_bucket_size=hash_bucket_size, dtype=dtype))

A sequence of categorical terms where ids are set by hashing.

    Pass this to `embedding_column` or `indicator_column` to convert sequence
    categorical data into dense representation for input to sequence NN, such as
    RNN.

Example:
    ```python
    tokens = sequence_categorical_column_with_hash_bucket(
        'tokens', hash_bucket_size=1000)
    tokens_embedding = embedding_column(tokens, dimension=10)
    columns = [tokens_embedding]

    features = tf.io.parse_example(..., features=make_parse_example_spec(columns))
    sequence_feature_layer = SequenceFeatures(columns)
    sequence_input, sequence_length = sequence_feature_layer(features)
    sequence_length_mask = tf.sequence_mask(sequence_length)

    rnn_cell = tf.keras.layers.SimpleRNNCell(hidden_size)
    rnn_layer = tf.keras.layers.RNN(rnn_cell)
    outputs, state = rnn_layer(sequence_input, mask=sequence_length_mask)
    ```

Args:
        key: A unique string identifying the input feature.
        hash_bucket_size: An int > 1. The number of buckets.
        dtype: The type of features. Only string and integer types are supported.

Returns:
        A `SequenceCategoricalColumn`.

Raises:
        ValueError: `hash_bucket_size` is not greater than 1.
        ValueError: `dtype` is neither string nor integer.


def apply_grad(self, grad, local_step=0, name=None):
    grad = ops.convert_to_tensor(grad, self._dtype)
    grad.get_shape().assert_is_compatible_with(self._shape)
    local_step = math_ops.cast(ops.convert_to_tensor(local_step), _dtypes.int64)
    return gen_data_flow_ops.resource_accumulator_apply_gradient(self._accumulator_ref, local_step=local_step, gradient=grad, name=name)

Attempts to apply a gradient to the accumulator.

    The attempt is silently dropped if the gradient is stale, i.e., local_step
    is less than the accumulator's global time step.

Args:
        grad: The gradient tensor to be applied.
        local_step: Time step at which the gradient was computed.
        name: Optional name for the operation.

Returns:
        The operation that (conditionally) applies a gradient to the accumulator.

Raises:
        ValueError: If grad is of the wrong shape


def _expand_dims_impl(st, axis, name=None):
    axis = array_ops.get_positive_axis(axis, st.rank + 1, axis_name='axis', ndims_name='rank(st)')
    with ops.name_scope(name, 'ExpandDims', [st, axis]):
        new_fields = {k: array_ops.expand_dims(v, axis) for k, v in st._fields.items()}
        new_shape = st.shape[:axis] + (1,) + st.shape[axis:]
        new_row_partitions = _expand_st_row_partitions(st, axis)
        new_nrows = st.nrows() if axis > 0 else 1
        return StructuredTensor.from_fields(new_fields, shape=new_shape, row_partitions=new_row_partitions, nrows=new_nrows)

Creates a StructuredTensor with a length 1 axis inserted at index `axis`.

    This is an implementation of tf.expand_dims for StructuredTensor. Note
    that the `axis` must be less than or equal to rank.

    >>> st = StructuredTensor.from_pyval([[{"x": 1}, {"x": 2}], [{"x": 3}]])
    >>> tf.expand_dims(st, 0).to_pyval()
    [[[{'x': 1}, {'x': 2}], [{'x': 3}]]]
    >>> tf.expand_dims(st, 1).to_pyval()
    [[[{'x': 1}, {'x': 2}]], [[{'x': 3}]]]
    >>> tf.expand_dims(st, 2).to_pyval()
    [[[{'x': 1}], [{'x': 2}]], [[{'x': 3}]]]
    >>> tf.expand_dims(st, -1).to_pyval()  # -1 is the same as 2
    [[[{'x': 1}], [{'x': 2}]], [[{'x': 3}]]]

Args:
        st: the original StructuredTensor.
        axis: the axis to insert the dimension: `-(rank + 1) <= axis <= rank`
        name: the name of the op.

Returns:
        a new structured tensor with larger rank.

Raises:
        an error if `axis < -(rank + 1)` or `rank < axis`.


def create_attention_mask_from_sequences(self, query_ids: List[int], table_values: List[TableValue]) -> List[int]:
    return [1] * (1 + len(query_ids) + 1 + len(table_values))

Creates the attention mask according to the query token IDs and a list of table values.

Args:
        query_ids (`List[int]`): list of token IDs corresponding to the ID.
        table_values (`List[TableValue]`): lift of table values, which are named tuples containing the
            token value, the column ID and the row ID of said token.

Returns:
        `List[int]`: List of ints containing the attention mask values.


def check_attribute_being_used(config_class, attributes, default_value, source_strings):
    attribute_used = False
    for attribute in attributes:
        for modeling_source in source_strings:
            if f'config.{attribute}' in modeling_source or f'getattr(config, "{attribute}"' in modeling_source or f'getattr(self.config, "{attribute}"' in modeling_source or ('TextConfig' in config_class.__name__ and f'config.get_text_config().{attribute}' in modeling_source):
                attribute_used = True
            elif re.search(f'getattr[ \\t\\v\\n\\r\\f]*\\([ \\t\\v\\n\\r\\f]*(self\\.)?config,[ \\t\\v\\n\\r\\f]*"{attribute}"', modeling_source) is not None:
                attribute_used = True
            if attribute_used:
                break
        if attribute_used:
            break
    attributes_to_allow = ['initializer_range', 'bos_index', 'eos_index', 'pad_index', 'unk_index', 'mask_index', 'image_token_id', 'video_token_id', 'image_seq_length', 'video_seq_length', 'image_size', 'text_config', 'use_cache', 'out_features', 'out_indices', 'sampling_rate', 'use_pretrained_backbone', 'backbone', 'backbone_config', 'use_timm_backbone', 'backbone_kwargs', 'rope_theta', 'partial_rotary_factor', 'pretraining_tp', 'boi_token_id', 'eoi_token_id']
    attributes_used_in_generation = ['encoder_no_repeat_ngram_size']
    case_allowed = True
    if not attribute_used:
        case_allowed = False
        for attribute in attributes:
            if attribute in ['is_encoder_decoder'] and default_value is True:
                case_allowed = True
            elif attribute in ['tie_word_embeddings'] and default_value is False:
                case_allowed = True
            elif attribute in attributes_to_allow + attributes_used_in_generation:
                case_allowed = True
            elif attribute.endswith('_token_id'):
                case_allowed = True
            if not case_allowed:
                allowed_cases = SPECIAL_CASES_TO_ALLOW.get(config_class.__name__, [])
                case_allowed = allowed_cases is True or attribute in allowed_cases
    return attribute_used or case_allowed

Check if any name in `attributes` is used in one of the strings in `source_strings`

Args:
        config_class (`type`):
            The configuration class for which the arguments in its `__init__` will be checked.
        attributes (`List[str]`):
            The name of an argument (or attribute) and its variant names if any.
        default_value (`Any`):
            A default value for the attribute in `attributes` assigned in the `__init__` of `config_class`.
        source_strings (`List[str]`):
            The python source code strings in the same modeling directory where `config_class` is defined. The file
            containing the definition of `config_class` should be excluded.


def numpy(self) -> npt.ArrayLike:
    maybe_arr = self._numpy()
    return maybe_arr.copy() if isinstance(maybe_arr, np.ndarray) else maybe_arr

Copy of the contents of this Tensor into a NumPy array or scalar.

    Unlike NumPy arrays, Tensors are immutable, so this method has to copy
    the contents to ensure safety. Use `memoryview` to get a readonly
    view of the contents without doing a copy:

    >>> t = tf.constant([42])
    >>> np.asarray(memoryview(t))
    array([42], dtype=int32)

    Note that `memoryview` is only zero-copy for Tensors on CPU. If a Tensor
    is on GPU, it will have to be transferred to CPU first in order for
    `memoryview` to work.

Returns:
        A NumPy array of the same shape and dtype or a NumPy scalar, if this
        Tensor has rank 0.

Raises:
        ValueError: If the dtype of this Tensor does not have a compatible
        NumPy dtype.


def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str]=None) -> Tuple[str]:
    if not os.path.isdir(save_directory):
        logger.error(f'Vocabulary path ({save_directory}) should be a directory')
        return
    vocab_file = os.path.join(save_directory, (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file'])
    with open(vocab_file, 'w', encoding='utf-8') as f:
        f.write(json.dumps(self.get_vocab(), ensure_ascii=False))
    return (vocab_file,)

Save the vocabulary and special tokens file to a directory.

Args:
        save_directory (`str`):
            The directory in which to save the vocabulary.

Returns:
        `Tuple(str)`: Paths to the files saved.


def _is_sequence_right_padded(mask):
    max_seq_length = mask.shape[1]
    count_of_true = torch.sum(mask, dim=1)
    batch_size = mask.shape[0]
    indices = torch.arange(max_seq_length, device=mask.device).repeat(batch_size, 1)
    right_padded_mask = indices < count_of_true.unsqueeze(1)
    return torch.all(mask == right_padded_mask)

Check the mask tensor and see if it right padded.

    cuDNN uses the sequence length param to skip the tailing
    timestep. If the data is left padded, or not a strict right padding (has
    masked value in the middle of the sequence), then cuDNN won't work
    properly in those cases.

    Left padded data: [[False, False, True, True, True]].
    Right padded data: [[True, True, True, False, False]].
    Mixture of mask/unmasked data: [[True, False, True, False, False]].

    Note that for the mixed data example above, the actually data RNN should see
    are those 2 Trues (index 0 and 2), the index 1 False should be ignored and
    not pollute the internal states.

Args:
        mask: the Boolean tensor with shape [batch, timestep]

Returns:
        boolean scalar tensor, whether the mask is strictly right padded.


def with_forward_compatibility_horizons(*horizons: Optional[tuple[int, int, int]]) -> Callable[[Callable[..., Any]], Callable[..., None]]:
    if not horizons:
        raise ValueError('Expected at least one horizon.')
    for horizon in horizons:
        if not (horizon is None or (len(horizon) == 3 and all((isinstance(x, int) for x in horizon)))):
            raise ValueError('Bad horizon value: %r' % horizon)

    def decorator(f: Callable[..., Any]) -> Callable[..., None]:
        if tf_inspect.isclass(f):
            raise ValueError('`with_forward_compatibility_horizons` only supports test methods.')

        def decorated(*args, **kwargs):
            for horizon in horizons:
                if horizon is None:
                    f(*args, **kwargs)
                else:
                    year, month, day = horizon
                    with forward_compatibility_horizon(year, month, day):
                        f(*args, **kwargs)
        return tf_decorator.make_decorator(f, decorated)
    return decorator

Executes the decorated test with the specified forward-compat horizons.

Args:
        *horizons: A list of (year, month, day) tuples.  If the list includes
        `None`, then the test will also be run with no forward-compatibility
        horizon set.

Returns:
        A decorator that will execute the test with the specified horizons.


def _supervised_signature_def(method_name, inputs, loss=None, predictions=None, metrics=None):
    if inputs is None or not inputs:
        raise ValueError('{} inputs cannot be None or empty.'.format(method_name))
    signature_inputs = {key: utils.build_tensor_info(tensor) for key, tensor in inputs.items()}
    signature_outputs = {}
    for output_set in (loss, predictions, metrics):
        if output_set is not None:
            sig_out = {key: utils.build_tensor_info(tensor) for key, tensor in output_set.items()}
            signature_outputs.update(sig_out)
    signature_def = signature_def_utils.build_signature_def(signature_inputs, signature_outputs, method_name)
    return signature_def

Creates a signature for training and eval data.

    This function produces signatures that describe the inputs and outputs
    of a supervised process, such as training or evaluation, that
    results in loss, metrics, and the like. Note that this function only requires
    inputs to be not None.

Args:
        method_name: Method name of the SignatureDef as a string.
        inputs: dict of string to `Tensor`.
        loss: dict of string to `Tensor` representing computed loss.
        predictions: dict of string to `Tensor` representing the output predictions.
        metrics: dict of string to `Tensor` representing metric ops.

Returns:
        A train- or eval-flavored signature_def.

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


def get_output_shape_at(self, node_index):
    return self._get_node_attribute_at_index(node_index, 'output_shapes', 'output shape')

Retrieves the output shape(s) of a layer at a given node.

Args:
        node_index: Integer, index of the node
            from which to retrieve the attribute.
            E.g. `node_index=0` will correspond to the
            first time the layer was called.

Returns:
        A shape tuple
        (or list of shape tuples if the layer has multiple outputs).

Raises:
        RuntimeError: If called in Eager mode.


def douglas_adi_step(theta=0.5):

    def _step_fn(time, next_time, coord_grid, value_grid, boundary_conditions, second_order_coeff_fn, first_order_coeff_fn, zeroth_order_coeff_fn, inner_second_order_coeff_fn, inner_first_order_coeff_fn, num_steps_performed, dtype=None, name=None):
        """Performs the step."""
        del num_steps_performed
        name = name or 'douglas_adi_step'
        return multidim_parabolic_equation_step(time, next_time, coord_grid, value_grid, boundary_conditions, douglas_adi_scheme(theta), second_order_coeff_fn, first_order_coeff_fn, zeroth_order_coeff_fn, inner_second_order_coeff_fn, inner_first_order_coeff_fn, dtype=dtype, name=name)
    return _step_fn

Creates a stepper function with Crank-Nicolson time marching scheme.

    Douglas ADI scheme is the simplest time marching scheme for solving parabolic
    PDEs with multiple spatial dimensions. The time step consists of several
    substeps: the first one is fully explicit, and the following `N` steps are
    implicit with respect to contributions of one of the `N` axes (hence "ADI" -
    alternating direction implicit). See `douglas_adi_scheme` below for more
    details.

Args:
        theta: positive Number. `theta = 0` corresponds to fully explicit scheme.
        The larger `theta` the stronger are the corrections by the implicit
        substeps. The recommended value is `theta = 0.5`, because the scheme is
        second order accurate in that case, unless mixed second derivative terms are
        present in the PDE.

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


def __init__(self, columns: list[str], split_string_by_delimiter: Optional[str]=None, *, ngram_range: tuple[int, int]=(1, 1), ngrams_separator: Optional[str]=None, compute_word_count: bool=False, key_vocab_filename: Optional[str]=None, name: Optional[str]=None):
    self.columns = columns
    self.ngram_range = ngram_range
    self.ngrams_separator = ngrams_separator
    self.name = name
    self.split_string_by_delimiter = split_string_by_delimiter
    self.key_vocab_filename = key_vocab_filename
    if compute_word_count:
        self.compute_word_count_fn = count_unique_words
    else:
        self.compute_word_count_fn = lambda *args, **kwargs: None
    if ngram_range != (1, 1) and (not ngrams_separator):
        raise ValueError('ngrams_separator must be specified when ngram_range is not (1, 1)')

Bag of words contains the unique words present in the input text.
    This operation applies a bag of words transformation to specified
    columns of incoming data. Also, the transformation accepts a Tuple of
    integers specifying the range of n-gram sizes. The transformation
    splits the input data into a set of consecutive n-grams if ngram_range
    is specified. The n-grams are then converted to a bag of words.
    Also, you can specify a seperator string that will be inserted between
    each ngram.

Args:
        columns: A list of column names to apply the transformation on.
        split_string_by_delimiter: (Optional) A string that specifies the
        delimiter to split the input strings before computing ngrams.
        ngram_range: A tuple of integers(inclusive) specifying the range of
        n-gram sizes.
        seperator: A string that will be inserted between each ngram.
        compute_word_count: A boolean that specifies whether to compute
        the unique word count over the entire dataset. Defaults to False.
        key_vocab_filename: The file name for the key vocabulary file when
        compute_word_count is True. If empty, a file name
        will be chosen based on the current scope. If provided, the vocab
        file will be suffixed with the column name.
        name: A name for the operation (optional).

    Note that original order of the input may not be preserved.


def forward(self, hidden_state):
    residual = hidden_state
    hidden_state = self.norm(hidden_state)
    if self.self_attn:
        batch_size, n_vars, num_patches, d_model = hidden_state.shape
        hidden_state_reshaped = hidden_state.reshape(batch_size * n_vars, num_patches, d_model)
        x_attn, _, _ = self.self_attn_layer(hidden_state_reshaped, output_attentions=False)
        x_attn = x_attn.reshape(batch_size, n_vars, num_patches, d_model)
    hidden_state = hidden_state.transpose(2, 3)
    hidden_state = self.mlp(hidden_state)
    if self.gated_attn:
        hidden_state = self.gating_block(hidden_state)
    hidden_state = hidden_state.transpose(2, 3)
    if self.self_attn:
        hidden_state = self.norm_attn(hidden_state + x_attn)
    out = hidden_state + residual
    return out

Args:
    hidden_state (`torch.Tensor`): Input tensor.

Returns:
    `torch.Tensor`: Transformed tensor.


def f(options, expected_tf_failures=0):
    test_parameters = [{'ksize': [[1, 1, 1, 1, 1], [1, 2, 2, 2, 1], [1, 2, 3, 4, 1]], 'strides': [[1, 1, 1, 1, 1], [1, 2, 1, 2, 1], [1, 2, 2, 4, 1]], 'input_shape': [[1, 1, 1, 1, 1], [1, 16, 15, 14, 1], [3, 16, 15, 14, 3]], 'padding': ['SAME', 'VALID'], 'data_format': ['NDHWC']}]

    def build_graph(parameters):
        input_tensor = tf.compat.v1.placeholder(dtype=tf.float32, name='input', shape=parameters['input_shape'])
        out = pool_op(input_tensor, ksize=parameters['ksize'], strides=parameters['strides'], data_format=parameters['data_format'], padding=parameters['padding'])
        return ([input_tensor], [out])

    def build_inputs(parameters, sess, inputs, outputs):
        input_values = create_tensor_data(tf.float32, parameters['input_shape'])
        return ([input_values], sess.run(outputs, feed_dict=dict(zip(inputs, [input_values]))))
    extra_convert_options = ExtraConvertOptions()
    extra_convert_options.allow_custom_ops = True
    make_zip_of_tests(options, test_parameters, build_graph, build_inputs, extra_convert_options, expected_tf_failures=expected_tf_failures)

Actual function that generates examples.

Args:
        options: An Options instance.
        expected_tf_failures: number of expected tensorflow failures.


def leaky_relu(features, alpha=0.2, name=None):
    with ops.name_scope(name, 'LeakyRelu', [features, alpha]) as name:
        features = ops.convert_to_tensor(features, name='features')
        if features.dtype.is_integer:
            features = math_ops.cast(features, dtypes.float32)
        if isinstance(alpha, np.ndarray):
            alpha = alpha.item()
        return gen_nn_ops.leaky_relu(features, alpha=alpha, name=name)

Compute the Leaky ReLU activation function.

    Source: [Rectifier Nonlinearities Improve Neural Network Acoustic Models.
    AL Maas, AY Hannun, AY Ng - Proc. ICML, 2013]
    (https://ai.stanford.edu/~amaas/papers/relu_hybrid_icml2013_final.pdf).

Args:
        features: A `Tensor` representing preactivation values. Must be one of
        the following types: `float16`, `float32`, `float64`, `int32`, `int64`.
        alpha: Slope of the activation function at x < 0.
        name: A name for the operation (optional).

Returns:
        The activation value.

    References:
        Rectifier Nonlinearities Improve Neural Network Acoustic Models:
        [Maas et al., 2013]
        (http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.693.1422)
        ([pdf]
        (http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.693.1422&rep=rep1&type=pdf))


def infer_steps_for_dataset(model, dataset, steps, epochs=1, steps_name='steps'):
    assert isinstance(dataset, data_types.DatasetV2)
    if model._in_multi_worker_mode() and dataset.options().experimental_distribute.auto_shard_policy != options_lib.AutoShardPolicy.OFF:
        return None
    size = backend.get_value(cardinality.cardinality(dataset))
    if size == cardinality.INFINITE and steps is None:
        raise ValueError('When passing an infinitely repeating dataset, you must specify the `%s` argument.' % (steps_name,))
    if size >= 0:
        if steps is not None and steps * epochs > size:
            if epochs > 1:
                raise ValueError('The dataset you passed contains %s batches, but you passed `epochs=%s` and `%s=%s`, which is a total of %s steps. We cannot draw that many steps from this dataset. We suggest to set `%s=%s`.' % (size, epochs, steps_name, steps, steps * epochs, steps_name, size // epochs))
            else:
                raise ValueError('The dataset you passed contains %s batches, but you passed `%s=%s`. We cannot draw that many steps from this dataset. We suggest to set `%s=%s`.' % (size, steps_name, steps, steps_name, size))
    if steps is None:
        if size >= 0:
            return size
        return None
    return steps

Infers steps_per_epoch needed to loop through a dataset.

Args:
        model: Keras model instance.
        dataset: Input data of type tf.data.Dataset.
        steps: Number of steps to draw from the dataset (may be None if unknown).
        epochs: Number of times to iterate over the dataset.
        steps_name: The string name of the steps argument, either `steps`,
        `validation_steps`, or `steps_per_epoch`. Only used for error message
        formatting.

Returns:
        Integer or `None`. Inferred number of steps to loop through the dataset.
        `None` is returned if 1) the size of the dataset is unknown and `steps` was
        not specified, or 2) this is multi-worker training and auto sharding is
        enabled.

Raises:
        ValueError: In case of invalid argument values.


def get_cuda_compute_capability(source_from_url=False):
    if not GPU_TYPE:
        if FLAGS.debug:
            print('Warning: GPU_TYPE is empty. Make sure to call `get_gpu_type()` first.')
    elif GPU_TYPE == 'unknown':
        if FLAGS.debug:
            print('Warning: Unknown GPU is detected. Skipping CUDA compute capability retrieval.')
    else:
        if source_from_url:
            cuda_compute_capa = cuda_compute_capability.retrieve_from_web()
        else:
            cuda_compute_capa = cuda_compute_capability.retrieve_from_golden()
        return cuda_compute_capa[GPU_TYPE]
    return

Retrieves CUDA compute capability based on the detected GPU type.

    This function uses the `cuda_compute_capability` module to retrieve the
    corresponding CUDA compute capability for the given GPU type.

Args:
        source_from_url: Boolean deciding whether to source compute capability
                    from NVIDIA website or from a local golden file.

Returns:
        List of all supported CUDA compute capabilities for the given GPU type.
        e.g. ['3.5', '3.7']


def __enter__(self):
    self._old = self._var_scope_store.current_scope
    if isinstance(self._name_or_scope, VariableScope):
        self._var_scope_store.open_variable_scope(self._new_name)
        self._old_subscopes = copy.copy(self._var_scope_store.variable_scopes_count)
        variable_scope_object = self._cached_variable_scope_object
    else:
        self._new_name = self._old.name + '/' + self._name_or_scope if self._old.name else self._name_or_scope
        self._reuse = self._reuse or self._old.reuse
        if self._old_name_scope is None:
            name_scope = self._name_or_scope
        else:
            name_scope = self._old_name_scope
        variable_scope_object = VariableScope(self._reuse, name=self._new_name, initializer=self._old.initializer, regularizer=self._old.regularizer, caching_device=self._old.caching_device, partitioner=self._old.partitioner, dtype=self._old.dtype, use_resource=self._old.use_resource, custom_getter=self._old.custom_getter, name_scope=name_scope, constraint=self._constraint)
        if self._initializer is not None:
            variable_scope_object.set_initializer(self._initializer)
        if self._regularizer is not None:
            variable_scope_object.set_regularizer(self._regularizer)
        if self._caching_device is not None:
            variable_scope_object.set_caching_device(self._caching_device)
        if self._partitioner is not None:
            variable_scope_object.set_partitioner(self._partitioner)
        if self._custom_getter is not None:
            variable_scope_object.set_custom_getter(_maybe_wrap_custom_getter(self._custom_getter, self._old.custom_getter))
        if self._dtype is not None:
            variable_scope_object.set_dtype(self._dtype)
        if self._use_resource is not None:
            variable_scope_object.set_use_resource(self._use_resource)
        self._var_scope_store.open_variable_scope(self._new_name)
    self._var_scope_store.current_scope = variable_scope_object
    self._last_variable_scope_object = variable_scope_object
    return variable_scope_object

Begins the scope block.

Returns:
        A VariableScope.

Raises:
        ValueError: when trying to reuse within a create scope, or create within
        a reuse scope, or if reuse is not `None` or `True`.
        TypeError: when the types of some arguments are not appropriate.


def unwrap_outputs(distribution_strategy, grouped_outputs, with_loss_tensor=False):
    if not with_loss_tensor:
        return flatten_per_replica_values(distribution_strategy, grouped_outputs)
    if not isinstance(grouped_outputs, list):
        grouped_outputs = [grouped_outputs]
    loss = distribution_strategy.reduce(reduce_util.ReduceOp.SUM, grouped_outputs[0], axis=None)
    all_outputs = flatten_per_replica_values(distribution_strategy, grouped_outputs[1:])
    if backend.is_tpu_strategy(distribution_strategy) and ops.executing_eagerly_outside_functions():
        all_outputs = all_outputs[::distribution_strategy.num_replicas_in_sync]
    return [loss] + all_outputs

Unwrap the list of outputs contained in the PerReplica parameters.

    This function calls `flatten_per_replica_values` to parse each of the input
    parameters into a list of outputs on the different devices. If we set
    `with_loss_tensor` to be True, we also call `reduce` on the list of losses on
    the different devices to give us one loss tensor.

Args:
        distribution_strategy: DistributionStrategy used to distribute training and
        validation.
        grouped_outputs: PerReplica outputs returned from the train or test function
        that we ran on each device.
        with_loss_tensor: Boolean that indicates if we need to add the reduced loss
        tensor as one of the outputs.

Returns:
        Values of each of the PerReplica outputs.


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, **kwargs):
    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)
        if sequence_length != 1:
            causal_mask = torch.triu(causal_mask, diagonal=1)
        causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1)
        causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
        if attention_mask is not None:
            causal_mask = causal_mask.clone()
            mask_length = attention_mask.shape[-1]
            padding_attention_mask = (attention_mask[:, None, None, :] == attention_mask[:, None, :, None])[:, :, -sequence_length:, :].to(dtype)
            padding_mask = causal_mask[:, :, :, :mask_length] + padding_attention_mask
            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.


def is_deterministic(self):
    return False

Whether this coder is guaranteed to encode values deterministically.

    A deterministic coder is required for key coders in GroupByKey operations
    to produce consistent results.

    For example, note that the default coder, the PickleCoder, is not
    deterministic: the ordering of picked entries in maps may vary across
    executions since there is no defined order, and such a coder is not in
    general suitable for usage as a key coder in GroupByKey operations, since
    each instance of the same key may be encoded differently.

Returns:
        Whether coder is deterministic.


def determine_framework(model: str, framework: Optional[str]=None) -> str:
    if framework is not None:
        return framework
    framework_map = {'pt': 'PyTorch', 'tf': 'TensorFlow'}
    exporter_map = {'pt': 'torch', 'tf': 'tf2onnx'}
    if os.path.isdir(model):
        if os.path.isfile(os.path.join(model, WEIGHTS_NAME)):
            framework = 'pt'
        elif os.path.isfile(os.path.join(model, TF2_WEIGHTS_NAME)):
            framework = 'tf'
        else:
            raise FileNotFoundError(f'Cannot determine framework from given checkpoint location. There should be a {WEIGHTS_NAME} for PyTorch or {TF2_WEIGHTS_NAME} for TensorFlow.')
        logger.info(f'Local {framework_map[framework]} model found.')
    elif is_torch_available():
        framework = 'pt'
    elif is_tf_available():
        framework = 'tf'
    else:
        raise OSError('Neither PyTorch nor TensorFlow found in environment. Cannot export to ONNX.')
    logger.info(f'Framework not requested. Using {exporter_map[framework]} to export to ONNX.')
    return framework

Determines the framework to use for the export.

    The priority is in the following order:
        1. User input via `framework`.
        2. If local checkpoint is provided, use the same framework as the checkpoint.
        3. Available framework in environment, with priority given to PyTorch

Args:
        model (`str`):
            The name of the model to export.
        framework (`str`, *optional*, defaults to `None`):
            The framework to use for the export. See above for priority if none provided.

Returns:
        The framework to use for the export.


def __init__(self, filename, separator=',', append=False):
    self.sep = separator
    self.filename = path_to_string(filename)
    self.append = append
    self.writer = None
    self.keys = None
    self.append_header = True
    super(CSVLogger, self).__init__()

Callback that streams epoch results to a CSV file.

    Supports all values that can be represented as a string,
    including 1D iterables such as `np.ndarray`.

Example:
    ```python
    csv_logger = CSVLogger('training.log')
    model.fit(X_train, Y_train, callbacks=[csv_logger])
    ```

Args:
        filename: Filename of the CSV file, e.g. `'run/log.csv'`.
        separator: String used to separate elements in the CSV file.
        append: Boolean. True: append if file exists (useful for continuing
            training). False: overwrite existing file.


def resize_positional_embeddings(positional_embeddings: torch.Tensor, spatial_shapes: torch.LongTensor, max_length: int) -> torch.Tensor:
    batch_size = spatial_shapes.shape[0]
    embed_dim = positional_embeddings.shape[-1]
    source_dtype = positional_embeddings.dtype
    resulted_positional_embeddings = torch.empty((batch_size, max_length, embed_dim), device=positional_embeddings.device, dtype=source_dtype)
    positional_embeddings = positional_embeddings.permute(2, 0, 1).unsqueeze(0)
    if positional_embeddings.device.type == 'cpu':
        positional_embeddings = positional_embeddings.to(torch.float32)
    for i in range(batch_size):
        height, width = spatial_shapes[i]
        resized_embeddings = F.interpolate(positional_embeddings, size=(height, width), mode='bilinear', align_corners=False, antialias=True)
        resized_embeddings = resized_embeddings.reshape(embed_dim, height * width).transpose(0, 1)
        resized_embeddings = resized_embeddings.to(source_dtype)
        resulted_positional_embeddings[i, :height * width] = resized_embeddings
        resulted_positional_embeddings[i, height * width:] = resized_embeddings[0]
    return resulted_positional_embeddings

Resize positional embeddings to image-specific size and pad to a fixed size.

Args:
        positional_embeddings (`torch.Tensor`):
            Position embeddings of shape (height, width, embed_dim)
        spatial_shapes (`torch.LongTensor`):
            Spatial shapes of shape (batch_size, 2) to resize the positional embeddings to
        max_length (`int`):
            Maximum length of the positional embeddings to pad resized positional embeddings to

Returns:
        `torch.Tensor`: Embeddings of shape (batch_size, max_length, embed_dim)


def resize(self, image: np.ndarray, size: Dict[str, int], resample: PILImageResampling=PILImageResampling.BILINEAR, data_format: Optional[Union[str, ChannelDimension]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None, **kwargs) -> np.ndarray:
    default_to_square = True
    if 'shortest_edge' in size:
        size = size['shortest_edge']
        default_to_square = False
    elif 'height' in size and 'width' in size:
        size = (size['height'], size['width'])
    else:
        raise ValueError("Size must contain either 'shortest_edge' or 'height' and 'width'.")
    output_size = get_resize_output_image_size(image, size=size, default_to_square=default_to_square, input_data_format=input_data_format)
    return resize(image, size=output_size, resample=resample, data_format=data_format, input_data_format=input_data_format, **kwargs)

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

Args:
        image (`np.ndarray`):
            Image to resize.
        size (`Dict[str, int]`):
            Size of the output image.
        resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BILINEAR`):
            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 (`ChannelDimension` or `str`, *optional*):
            The channel dimension format of the input image. If not provided, it will be inferred.


def build_inputs_with_special_tokens(self, token_ids_0: List[int], token_ids_1: Optional[List[int]]=None) -> List[int]:
    if token_ids_1 is None:
        return self.prefix_tokens + token_ids_0 + self.suffix_tokens
    return self.prefix_tokens + token_ids_0 + token_ids_1 + self.suffix_tokens

Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
    adding special tokens. An NLLB sequence has the following format, where `X` represents the sequence:

    - `input_ids` (for encoder) `X [eos, src_lang_code]`
    - `decoder_input_ids`: (for decoder) `X [eos, tgt_lang_code]`

    BOS is never used. Pairs of sequences are not the expected use case, but they will be handled without a
    separator.

Args:
        token_ids_0 (`List[int]`):
            List of IDs to which the special tokens will be added.
        token_ids_1 (`List[int]`, *optional*):
            Optional second list of IDs for sequence pairs.

Returns:
        `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.


def __init__(self, config: PretrainedConfig, generation_config: GenerationConfig, device: torch.device, dtype: torch.dtype=torch.float16, layer_device_map: Optional[Dict[int, Union[str, torch.device, int]]]=None, initial_prompt_shapes: Optional[List[List[int]]]=None) -> None:
    self.num_key_value_heads = config.num_attention_heads if getattr(config, 'num_key_value_heads', None) is None else config.num_key_value_heads
    self.head_dim = config.head_dim if hasattr(config, 'head_dim') else config.hidden_size // config.num_attention_heads
    self.num_hidden_layers = config.num_hidden_layers
    num_blocks = getattr(generation_config, 'num_blocks', None)
    block_size = getattr(generation_config, 'block_size', None)
    if num_blocks is None or block_size is None:
        logger.info('Calculating optimal block size and number...')
        num_blocks, block_size = compute_optimal_blocks(device, config, generation_config, initial_prompt_shapes or [], dtype, median_prefill_length=200)
        logger.info(f'Using calculated num_blocks={num_blocks}, block_size={block_size}')
    self.block_size = block_size
    self.num_blocks = num_blocks
    self.cache_shape = (self.num_key_value_heads, num_blocks, self.block_size, self.head_dim)
    self.dtype = dtype
    self.device = device
    self.key_cache: List[torch.Tensor] = []
    self.value_cache: List[torch.Tensor] = []
    for idx in range(config.num_hidden_layers):
        layer_device = layer_device_map[idx] if layer_device_map is not None else device
        new_layer_key_cache = torch.zeros(self.cache_shape, dtype=self.dtype, device=layer_device)
        new_layer_value_cache = torch.zeros(self.cache_shape, dtype=self.dtype, device=layer_device)
        torch._dynamo.mark_static_address(new_layer_key_cache)
        torch._dynamo.mark_static_address(new_layer_value_cache)
        self.key_cache.append(new_layer_key_cache)
        self.value_cache.append(new_layer_value_cache)
    self._free_blocks = deque(range(num_blocks))
    self._block_tables: Dict[str, List[int]] = {}

Initialize a paged attention cache for efficient memory usage.

Args:
        config: Model configuration
        generation_config: Generation configuration containing cache parameters
        device: Device for the cache tensors
        dtype: Data type for the cache tensors
        layer_device_map: Optional mapping of layer indices to devices
        initial_prompt_shapes: Optional sample prompts to help calculate optimal cache size


def get_all_plugin_assets(graph=None):
    if graph is None:
        graph = ops.get_default_graph()
    out = []
    for name in graph.get_collection(_PLUGIN_ASSET_PREFIX):
        collection = graph.get_collection(_PLUGIN_ASSET_PREFIX + name)
        if len(collection) != 1:
            raise ValueError('Collection for %s had %d items, expected 1' % (name, len(collection)))
        out.append(collection[0])
    return out

Retrieve all PluginAssets stored in the graph collection.

Args:
        graph: Optionally, the graph to get assets from. If unspecified, the default
        graph is used.

Returns:
        A list with all PluginAsset instances in the graph.

Raises:
        ValueError: if we unexpectedly find a collection with the wrong number of
        PluginAssets.


def _find_suite_class():
    test_suites = _find_suite_classes_in_module(sys.modules['__main__'])
    if len(test_suites) == 0:
        logging.debug('No suite class found in the __main__ module, trying to find it in the module of the caller of suite_runner.run_suite_class method.')
        stacks = inspect.stack()
        if len(stacks) < 2:
            logging.debug('Failed to get the caller stack of run_suite_class. Got stacks: %s', stacks)
        else:
            run_suite_class_caller_frame_info = inspect.stack()[2]
            caller_frame = run_suite_class_caller_frame_info.frame
            module = inspect.getmodule(caller_frame)
            if module is None:
                logging.debug('Failed to find module for frame %s', caller_frame)
            else:
                test_suites = _find_suite_classes_in_module(module)
    if len(test_suites) != 1:
        logging.error('Expected 1 test class per file, found %s.', [t.__name__ for t in test_suites])
        sys.exit(1)
    return test_suites[0]

Finds the test suite class.

    First search for test suite classes in the __main__ module. If no test suite
    class is found, search in the module that is calling
    `suite_runner.run_suite_class`.

    Walk through module members and find the subclass of BaseSuite. Only
    one subclass is allowed.

Returns:
        The test suite class in the test module.


def _convert_concrete_functions_to_saved_model(self, output_dir):
    if len(self._funcs) == 0:
        raise ValueError('No ConcreteFunction is specified.')
    if not self.experimental_lower_to_saved_model:
        return (None, None, None)
    if not self._trackable_obj or isinstance(self._trackable_obj, (_function.ConcreteFunction, _def_function.Function)):
        return (None, None, None)
    signatures = {}
    signature_keys = []
    try:
        if len(self._funcs) == 1:
            signatures[_signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = self._funcs[0]
            signature_keys = [_signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
        else:
            for func in self._funcs:
                signatures[func.graph.name] = func
                signature_keys.append(func.graph.name)
        _save.save(self._trackable_obj, output_dir, signatures=signatures, options=_save_options.SaveOptions(save_debug_info=True))
    except Exception:
        return (None, None, None)
    self.saved_model_dir = output_dir
    self._saved_model_tags = set([_tag_constants.SERVING])
    self._saved_model_exported_names = signature_keys
    self._parse_saved_model_args(always_enable_saved_model_import=True)
    if self.saved_model_dir:
        graph_def, input_tensors, output_tensors = self._load_saved_model(self.saved_model_dir, self._saved_model_tags)
        self._trackable_obj = _load(self.saved_model_dir, self._saved_model_tags)
        return (graph_def, input_tensors, output_tensors)
    return (None, None, None)

Save concrete functions to the SavedModel format.

Args:
        output_dir: The output directory to save the SavedModel.

Returns:
        graph_def: The frozen GraphDef.
        input_tensors: List of input tensors.
        output_tensors: List of output tensors.


def _finalize_func(string_handle):
    iterator_resource = gen_dataset_ops.iterator_from_string_handle_v2(string_handle, **self._input_dataset._flat_structure)
    with ops.control_dependencies([resource_variable_ops.destroy_resource_op(iterator_resource, ignore_lookup_error=True)]):
        return array_ops.constant(0, dtypes.int64)

Destroys the iterator resource created.

Args:
        string_handle: An iterator string handle created by _init_func

Returns:
        Tensor constant 0


def count(self) -> 'Builder':
    return self._to_builder(_evaluation.CountFunction(self.node.context, self.node, []))

The FHIRPath count() function.

Returns:
        An expression that evaluates to the count of items in the parent.


def with_outputs(self, *tags, main=None, allow_unknown_tags=None):
    if main in tags:
        raise ValueError('Main output tag %r must be different from side output tags %r.' % (main, tags))
    return _MultiParDo(self, tags, main, allow_unknown_tags)

Returns a tagged tuple allowing access to the outputs of a
    :class:`ParDo`.

    The resulting object supports access to the
    :class:`~apache_beam.pvalue.PCollection` associated with a tag
    (e.g. ``o.tag``, ``o[tag]``) and iterating over the available tags
    (e.g. ``for tag in o: ...``).

Args:
        *tags: if non-empty, list of valid tags. If a list of valid tags is given,
        it will be an error to use an undeclared tag later in the pipeline.
        **main_kw: dictionary empty or with one key ``'main'`` defining the tag to
        be used for the main output (which will not have a tag associated with
        it).

Returns:
        ~apache_beam.pvalue.DoOutputsTuple: An object of type
        :class:`~apache_beam.pvalue.DoOutputsTuple` that bundles together all
        the outputs of a :class:`ParDo` transform and allows accessing the
        individual :class:`~apache_beam.pvalue.PCollection` s for each output
        using an ``object.tag`` syntax.

Raises:
        TypeError: if the **self** object is not a
        :class:`~apache_beam.pvalue.PCollection` that is the result of a
        :class:`ParDo` transform.
        ValueError: if **main_kw** contains any key other than
        ``'main'``.


def _get_compatible_spec(value_or_spec1, value_or_spec2):
    spec1 = _get_spec_for(value_or_spec1)
    spec2 = _get_spec_for(value_or_spec2)
    common = spec1._without_tensor_names().most_specific_common_supertype([spec2._without_tensor_names()])
    if common is None:
        raise TypeError(f'No common supertype of {spec1} and {spec2}.')
    return common

Returns the most specific compatible spec.

Args:
        value_or_spec1: A TypeSpecs or a value that has a defined TypeSpec.
        value_or_spec2: A TypeSpecs or a value that has a defined TypeSpec.

Returns:
        The most specific compatible TypeSpecs of the input.

Raises:
        ValueError: If value_or_spec1 is not compatible with value_or_spec2.


def fields_to_string(fields, values={}):
    items = [repr(field['name']) + ':' + repr(values.get(field['name'], field.get('default', ''))) + ',' + ('  # %s' % field['description'] if 'description' in field else '') for field in fields]
    return '{\n  %s\n}' % '\n  '.join(items) if items else '{}'

Converts fields to a dictionary of parameters as a string.

    Used to generate input blocks in generated code.

    For example:
        {
        'auth':'user', # authentication to use for call.
        'name':'value', # parameter to pass to function.
        }

Args:
    - fields: (list) Contains {"field":{...}} objects to be rendered as imputs.
    - values: (dict) Default values to use for each field.

Returns:
        String representing python code to be written to a generated file.