def set_clang_compiler_path_win(environ_cp):
    default_clang_path = 'C:/Program Files/LLVM/bin/clang.exe'
    if not os.path.exists(default_clang_path):
        default_clang_path = shutil.which('clang') or ''
    clang_compiler_path = prompt_loop_or_load_from_env(environ_cp, var_name='CLANG_COMPILER_PATH', var_default=default_clang_path, ask_for_var='Please specify the path to clang executable.', check_success=os.path.exists, resolve_symlinks=True, error_msg='Invalid clang path. %s cannot be found. Note that Clang is nowpreferred compiler. You may use MSVC by removing --config=win_clang')
    write_action_env_to_bazelrc('CLANG_COMPILER_PATH', clang_compiler_path)
    write_to_bazelrc(f'build --repo_env=CC="{clang_compiler_path}"')
    write_to_bazelrc(f'build --repo_env=BAZEL_COMPILER="{clang_compiler_path}"')
    return clang_compiler_path

Set CLANG_COMPILER_PATH and environment variables.

    Loop over user prompts for clang path until receiving a valid response.
    Default is used if no input is given. Set CLANG_COMPILER_PATH and write
    environment variables CC and BAZEL_COMPILER to .bazelrc.

Args:
        environ_cp: (Dict) copy of the os.environ.

Returns:
        string value for clang_compiler_path.


def sample_paths(self, times: types.RealTensor, num_samples: Optional[int]=1, initial_state: Optional[types.RealTensor]=None, random_type: Optional[random.RandomType]=None, seed: Optional[types.IntTensor]=None, swap_memory: Optional[bool]=True, time_step: Optional[types.RealTensor]=None, num_time_steps: Optional[types.IntTensor]=None, skip: Optional[types.IntTensor]=0, precompute_normal_draws: Optional[bool]=True, times_grid: Optional[types.RealTensor]=None, normal_draws: Optional[types.RealTensor]=None, watch_params: Optional[List[types.RealTensor]]=None, validate_args: Optional[bool]=False, name: Optional[str]=None) -> types.RealTensor:
    name = name or self._name + '_log_sample_path'
    with tf.name_scope(name):
        if initial_state is not None:
            initial_state = tf.math.log(tf.convert_to_tensor(initial_state, dtype_hint=tf.float64))
        if self.precompute_iv():
            if time_step is not None or num_time_steps is not None or times_grid is not None:
                raise ValueError('`time_step`, `num_time_steps`, or `times_grid` cannot be usedwith the interpolated LVM')
            times_grid = self._times_grid
        return tf.math.exp(super(LocalVolatilityModel, self).sample_paths(times=times, num_samples=num_samples, initial_state=initial_state, random_type=random_type, seed=seed, swap_memory=swap_memory, name=name, time_step=time_step, num_time_steps=num_time_steps, skip=skip, precompute_normal_draws=precompute_normal_draws, times_grid=times_grid, normal_draws=normal_draws, watch_params=watch_params, validate_args=validate_args))

Returns samples from the LV process.

    See GenericItoProcess.sample_paths. If `times_grid` is supplied to
    `__init__`, then `times_grid` cannot be supplied here.

Raises:
        ValueError: If `precompute_iv` is True, but `time_step`, `num_time_steps`
        or `times_grid` are given.


def enable_debug_mode():
    if context.executing_eagerly():
        toggle_debug_mode(True)
    else:
        raise ValueError('`enable_debug_mode() is only supported in eager mode.')

Enables debug mode for tf.data.

    Example usage with pdb module:
    ```
    import tensorflow as tf
    import pdb

    tf.data.experimental.enable_debug_mode()

    def func(x):
        # Python 3.7 and older requires `pdb.Pdb(nosigint=True).set_trace()`
        pdb.set_trace()
        x = x + 1
        return x

    dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])
    dataset = dataset.map(func)

    for item in dataset:
        print(item)
    ```

    The effect of debug mode is two-fold:

    1) Any transformations that would introduce asynchrony, parallelism, or
    non-determinism to the input pipeline execution will be forced to execute
    synchronously, sequentially, and deterministically.

    2) Any user-defined functions passed into tf.data transformations such as
    `map` will be wrapped in `tf.py_function` so that their body is executed
    "eagerly" as a Python function as opposed to a traced TensorFlow graph, which
    is the default behavior. Note that even when debug mode is enabled, the
    user-defined function is still traced  to infer the shape and type of its
    outputs; as a consequence, any `print` statements or breakpoints will be
    triggered once during the tracing before the actual execution of the input
    pipeline.

    NOTE: As the debug mode setting affects the construction of the tf.data input
    pipeline, it should be enabled before any tf.data definitions.

Raises:
        ValueError: When invoked from graph mode.


def _from_spec(cls, spec: Union['DynamicRaggedShape.Spec', ragged_tensor.RaggedTensorSpec, tensor_lib.TensorSpec], dtype: dtypes.DType=dtypes.int64) -> 'DynamicRaggedShape.Spec':
    if isinstance(spec, DynamicRaggedShape.Spec):
        return spec
    elif isinstance(spec, ragged_tensor.RaggedTensorSpec):
        return cls._from_tensor_shape(spec.shape, spec.ragged_rank, spec.row_splits_dtype)
    elif isinstance(spec, tensor_lib.TensorSpec):
        return cls._from_tensor_shape(shape=spec.shape, num_row_partitions=0, dtype=dtype)

Create a TypeSpec for the shape of an object with a given TypeSpec.

    I.e., if `x_spec = tf.type_spec_from_value(x)`, then
    `DynamicRaggedShape.from_spec(x_spec)` returns a TypeSpec compatible with
    `tf.type_spec_from_value(tf.shape(x))`.

    >>> rt = tf.ragged.constant([[1, 2], [3], [4, 5, 6]])
    >>> rt_spec = tf.type_spec_from_value(rt)
    >>> rt_shape = DynamicRaggedShape.from_tensor(rt)

    >>> shape_spec_1 = tf.type_spec_from_value(rt_shape)
    >>> shape_spec_2 = DynamicRaggedShape.Spec._from_spec(rt_spec)
    >>> assert shape_spec_1.is_compatible_with(shape_spec_2)

Args:
        spec: a Spec of a Tensor or RaggedTensor.
        dtype: the default dtype (if necessary).

Returns:
        A Spec of the shape of a Tensor or RaggedTensor.


def _old_init(cls, fields, shape, nrows, row_partitions, internal=False):
    assert isinstance(fields, dict), fields
    assert isinstance(shape, tensor_shape.TensorShape), shape
    assert nrows is None or isinstance(nrows, tensor.Tensor), nrows
    assert row_partitions is None or isinstance(row_partitions, tuple), row_partitions
    return StructuredTensor(fields=fields, ragged_shape=_dynamic_ragged_shape_init(fields, shape, nrows, row_partitions))

Private constructor -- use factory methods to create StructuredTensors.

    This constructor builds a `StructuredTensor` from the given attributes,
    performing minimal validation.

Args:
        fields: A dictionary mapping from string to `Tensor`, `RaggedTensor`, or
        `StructuredTensor`.  (This dict is not copied, so the caller must ensure
        that it does not get mutated via leaked references.)
        shape: `tf.TensorShape` with statically known rank.
        nrows: scalar integer `tf.Tensor`, or `None` if `shape.rank==0`.
        row_partitions: tuple of `RowPartition`s, with length `shape.rank-1`.
        internal: ignored argument.

Returns:
        a StructuredTensor.


def forward(self, input_ids: Optional[torch.LongTensor]=None, pixel_values: Optional[torch.FloatTensor]=None, attention_mask: Optional[torch.Tensor]=None, position_ids: Optional[torch.LongTensor]=None, return_loss: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, interpolate_pos_encoding: bool=True, return_dict: Optional[bool]=None) -> Union[Tuple, CLIPSegOutput]:
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    vision_outputs = self.vision_model(pixel_values=pixel_values, output_attentions=output_attentions, output_hidden_states=output_hidden_states, interpolate_pos_encoding=interpolate_pos_encoding, return_dict=return_dict)
    text_outputs = self.text_model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    image_embeds = vision_outputs[1]
    image_embeds = self.visual_projection(image_embeds)
    text_embeds = text_outputs[1]
    text_embeds = self.text_projection(text_embeds)
    image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True)
    text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True)
    logit_scale = self.logit_scale.exp()
    logits_per_text = torch.matmul(text_embeds, image_embeds.t()) * logit_scale
    logits_per_image = logits_per_text.t()
    loss = None
    if return_loss:
        loss = clipseg_loss(logits_per_text)
    if not return_dict:
        output = (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)
        return (loss,) + output if loss is not None else output
    return CLIPSegOutput(loss=loss, logits_per_image=logits_per_image, logits_per_text=logits_per_text, text_embeds=text_embeds, image_embeds=image_embeds, text_model_output=text_outputs, vision_model_output=vision_outputs)

return_loss (`bool`, *optional*):
    Whether or not to return the contrastive loss.

Example:
    ```python
    >>> from PIL import Image
    >>> import requests
    >>> from transformers import AutoProcessor, CLIPSegModel

    >>> processor = AutoProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
    >>> model = CLIPSegModel.from_pretrained("CIDAS/clipseg-rd64-refined")

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

    >>> inputs = processor(
    ...     text=["a photo of a cat", "a photo of a dog"], images=image, return_tensors="pt", padding=True
    ... )

    >>> outputs = model(**inputs)
    >>> logits_per_image = outputs.logits_per_image  # this is the image-text similarity score
    >>> probs = logits_per_image.softmax(dim=1)  # we can take the softmax to get the label probabilities
    ```


def ExpandSuperClasses(self, t):
    superclasses = set()
    self._CollectSuperclasses(t, superclasses)
    return superclasses

Generate a list of all (known) superclasses for a type.

Args:
        t: A type name. E.g. "int".

Returns:
        A set of types. This set includes t as well as all its superclasses. For
        example, this will return "bool", "int" and "object" for "bool".


def __floordiv__(self, other):
    try:
        other = as_dimension(other)
    except (TypeError, ValueError):
        return NotImplemented
    if self._value is None or other.value is None:
        return Dimension(None)
    else:
        return Dimension(self._value // other.value)

Returns the quotient of `self` and `other` rounded down.

    Dimensions are divided as follows:

    ```python
    tf.compat.v1.Dimension(m)    // tf.compat.v1.Dimension(n)     ==
    tf.compat.v1.Dimension(m // n)
    tf.compat.v1.Dimension(m)    // tf.compat.v1.Dimension(None)  # equiv. to
    tf.compat.v1.Dimension(None)
    tf.compat.v1.Dimension(None) // tf.compat.v1.Dimension(n)     # equiv. to
    tf.compat.v1.Dimension(None)
    tf.compat.v1.Dimension(None) // tf.compat.v1.Dimension(None)  # equiv. to
    tf.compat.v1.Dimension(None)
    ```

Args:
        other: Another Dimension, or a value accepted by `as_dimension`.

Returns:
        A `Dimension` whose value is the integer quotient of `self` and `other`.


def clip(x, min_value, max_value):
    if isinstance(min_value, (int, float)) and isinstance(max_value, (int, float)):
        if max_value < min_value:
            max_value = min_value
    if min_value is None:
        min_value = -np.inf
    if max_value is None:
        max_value = np.inf
    return clip_ops.clip_by_value(x, min_value, max_value)

Element-wise value clipping.

Args:
        x: Tensor or variable.
        min_value: Python float, integer, or tensor.
        max_value: Python float, integer, or tensor.

Returns:
        A tensor.


def true_positives(labels, predictions, weights=None, metrics_collections=None, updates_collections=None, name=None):
    if context.executing_eagerly():
        raise RuntimeError('tf.metrics.true_positives is not supported when eager execution is enabled.')
    with variable_scope.variable_scope(name, 'true_positives', (predictions, labels, weights)):
        predictions, labels, weights = _remove_squeezable_dimensions(predictions=math_ops.cast(predictions, dtype=dtypes.bool), labels=math_ops.cast(labels, dtype=dtypes.bool), weights=weights)
        is_true_positive = math_ops.logical_and(math_ops.equal(labels, True), math_ops.equal(predictions, True))
        return _count_condition(is_true_positive, weights, metrics_collections, updates_collections)

Sum the weights of true_positives.

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

Args:
        labels: The ground truth values, a `Tensor` whose dimensions must match
        `predictions`. Will be cast to `bool`.
        predictions: The predicted values, a `Tensor` of arbitrary dimensions. Will
        be cast to `bool`.
        weights: Optional `Tensor` whose rank is either 0, or the same rank as
        `labels`, and must be broadcastable to `labels` (i.e., all dimensions must
        be either `1`, or the same as the corresponding `labels` dimension).
        metrics_collections: An optional list of collections that the metric
        value variable should be added to.
        updates_collections: An optional list of collections that the metric update
        ops should be added to.
        name: An optional variable_scope name.

Returns:
        value_tensor: A `Tensor` representing the current value of the metric.
        update_op: An operation that accumulates the error from a batch of data.

Raises:
        ValueError: If `predictions` and `labels` have mismatched shapes, or if
        `weights` is not `None` and its shape doesn't match `predictions`, or if
        either `metrics_collections` or `updates_collections` are not a list or
        tuple.
        RuntimeError: If eager execution is enabled.


def convert_to_rgb(image: np.ndarray, palette: Optional[PIL.ImagePalette.ImagePalette]=None, data_format: Optional[Union[str, ChannelDimension]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None) -> ImageInput:
    if input_data_format is None:
        input_data_format = infer_channel_dimension_format(image, num_channels=(1, 3, 4))
    data_format = input_data_format if data_format is None else data_format
    mode = 'P' if palette is not None else None
    image = to_pil_image(image, image_mode=mode, input_data_format=input_data_format)
    if image.mode == 'P' and palette is not None:
        image.putpalette(palette)
    image_rgba = image.convert('RGBA')
    background = Image.new('RGBA', image_rgba.size, (255, 255, 255))
    alpha_composite = Image.alpha_composite(background, image_rgba)
    alpha_composite = alpha_composite.convert('RGB')
    output_array = np.array(alpha_composite)
    output_array = to_channel_dimension_format(output_array, data_format, input_channel_dim=ChannelDimension.LAST)
    return output_array

Converts an image to RGB format.

Args:
        image (`np.ndarray`):
            The image to convert.
        palette (List[int], *optional*):
            The palette to use if given.
        data_format (ChannelDimension or str, *optional*):
            The channel dimension format for the output 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.


def aggregate_events(events: List[str]) -> List[Tuple[str, int]]:
    frequency_list = collections.Counter(events)
    return sorted(frequency_list.items(), key=lambda x: (-x[1], x[0]))

Returns a list of tuples: (event string, count of times they appear).

    The list is sorted descending by count and then ascending by event string.

Args:
        events: A list of strings defining an event (either an error or a warning).

Returns:
        List of tuples : (event string, number of times they appear).


def analyze(model_path=None, model_content=None, gpu_compatibility=False, **kwargs):
    if not model_path and (not model_content):
        raise ValueError('neither `model_path` nor `model_content` is provided')
    if model_path:
        print(f'=== {model_path} ===\n')
        tflite_model = model_path
        input_is_filepath = True
    else:
        print('=== TFLite ModelAnalyzer ===\n')
        tflite_model = model_content
        input_is_filepath = False
    if kwargs.get('experimental_use_mlir', False):
        print(wrap_converter.wrapped_flat_buffer_file_to_mlir(tflite_model, input_is_filepath))
    else:
        print(_analyzer_wrapper.ModelAnalyzer(tflite_model, input_is_filepath, gpu_compatibility))

Analyzes the given tflite_model with dumping model structure.

    This tool provides a way to understand users' TFLite flatbuffer model by
    dumping internal graph structure. It also provides additional features
    like checking GPU delegate compatibility.

    WARNING: Experimental interface, subject to change.
            The output format is not guaranteed to stay stable, so don't
            write scripts to this.

Args:
        model_path: TFLite flatbuffer model path.
        model_content: TFLite flatbuffer model object.
        gpu_compatibility: Whether to check GPU delegate compatibility.
        **kwargs: Experimental keyword arguments to analyze API.

Returns:
        Print analyzed report via console output.


def execute_with_cancellation(op_name, num_outputs, inputs, attrs, ctx, cancellation_manager, name=None):
    device_name = ctx.device_name
    try:
        ctx.ensure_initialized()
        tensors = pywrap_tfe.TFE_Py_ExecuteCancelable(ctx._handle, device_name, op_name, inputs, attrs, cancellation_manager._impl, num_outputs)
    except core._NotOkStatusException as e:
        if name is not None:
            e.message += ' name: ' + name
        raise core._status_to_exception(e) from None
    except TypeError as e:
        keras_symbolic_tensors = [x for x in inputs if _is_keras_symbolic_tensor(x)]
        if keras_symbolic_tensors:
            raise core._SymbolicException('Inputs to eager execution function cannot be Keras symbolic tensors, but found {}'.format(keras_symbolic_tensors))
        raise e
    return tensors

Execute a TensorFlow operation.

Args:
        op_name: Name of the TensorFlow operation (see REGISTER_OP in C++ code) to
        execute.
        num_outputs: The number of outputs of the operation to fetch. (Explicitly
        provided instead of being inferred for performance reasons).
        inputs: A list of inputs to the operation. Each entry should be a Tensor, or
        a value which can be passed to the Tensor constructor to create one.
        attrs: A tuple with alternating string attr names and attr values for this
        operation.
        ctx: The value of context.context().
        cancellation_manager: a `CancellationManager` object that can be used to
        cancel the operation.
        name: Customized name for the operation.

Returns:
        List of output Tensor objects. The list is empty if there are no outputs

Raises:
        An exception on error.


def call(self, input_ids: TFModelInputType | None=None, attention_mask: np.ndarray | tf.Tensor | None=None, token_type_ids: np.ndarray | tf.Tensor | None=None, position_ids: np.ndarray | tf.Tensor | None=None, head_mask: np.ndarray | tf.Tensor | None=None, inputs_embeds: np.ndarray | tf.Tensor | None=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, labels: np.ndarray | tf.Tensor | None=None, training: Optional[bool]=False) -> Union[TFSequenceClassifierOutput, Tuple[tf.Tensor]]:
    outputs = self.tapas(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, training=training)
    pooled_output = outputs[1]
    pooled_output = self.dropout(inputs=pooled_output, training=training)
    logits = self.classifier(inputs=pooled_output)
    loss = None if labels is None else self.hf_compute_loss(labels=labels, logits=logits)
    if not return_dict:
        output = (logits,) + outputs[2:]
        return (loss,) + output if loss is not None else output
    return TFSequenceClassifierOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions)

labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
    Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
    config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
    `config.num_labels > 1` a classification loss is computed (Cross-Entropy). Note: this is called
    "classification_class_index" in the original implementation.

Returns:

Example:
    ```python
    >>> from transformers import AutoTokenizer, TapasForSequenceClassification
    >>> import tensorflow as tf
    >>> import pandas as pd

    >>> tokenizer = AutoTokenizer.from_pretrained("google/tapas-base-finetuned-tabfact")
    >>> model = TapasForSequenceClassification.from_pretrained("google/tapas-base-finetuned-tabfact")

    >>> data = {
    ...     "Actors": ["Brad Pitt", "Leonardo Di Caprio", "George Clooney"],
    ...     "Age": ["56", "45", "59"],
    ...     "Number of movies": ["87", "53", "69"],
    ... }
    >>> table = pd.DataFrame.from_dict(data)
    >>> queries = [
    ...     "There is only one actor who is 45 years old",
    ...     "There are 3 actors which played in more than 60 movies",
    ... ]

    >>> inputs = tokenizer(table=table, queries=queries, padding="max_length", return_tensors="tf")
    >>> labels = tf.convert_to_tensor([1, 0])  # 1 means entailed, 0 means refuted

    >>> outputs = model(**inputs, labels=labels)
    >>> loss = outputs.loss
    >>> logits = outputs.logits
    ```


def coordination_leader(cluster_spec):
    cluster_spec = normalize_cluster_spec(cluster_spec)
    if not cluster_spec.as_dict():
        return ''
    if 'ps' in cluster_spec.jobs:
        return '/job:ps/replica:0/task:0'
    if 'chief' in cluster_spec.jobs:
        return '/job:chief/replica:0/task:0'
    assert 'worker' in cluster_spec.jobs
    return '/job:worker/replica:0/task:0'

Return the task name of the coordination service leader.

Args:
        cluster_spec: a dict, `ClusterDef` or `ClusterSpec` object sxpecifying the
        cluster configurations.

Returns:
        a string indicating the task name of the coordination service leader.


Prepares two tables, both with three distinct rows.
    The first table has two columns:
    1.0, 2.0 | 3.0
    2.0, 0.0 | 1.0
    1.0, 3.0 | 4.0
    The second table has three columns:
    1.0 | 2.0 | 3.0
    2.0 | 0.0 | 1.0
    1.0 | 3.0 | 4.0

Returns:
    SegmentedTensors with the tables.


def convert_typing_to_builtin(typ):
    origin = getattr(typ, '__origin__', None)
    args = getattr(typ, '__args__', None)
    if origin not in _BUILTINS:
        return typ
    if not args:
        return origin
    if origin is list:
        return list[convert_typing_to_builtin(args[0])]
    elif origin is dict:
        return dict[convert_typing_to_builtin(args[0]), convert_typing_to_builtin(args[1])]
    elif origin is tuple:
        return tuple[tuple(convert_typing_to_builtin(args))]
    elif origin is set:
        return set[convert_typing_to_builtin(args)]
    elif origin is frozenset:
        return frozenset[convert_typing_to_builtin(args)]

Converts a given typing collections type to its builtin counterpart.

Args:
        typ: A typing type (e.g., typing.List[int]).

Returns:
        type: The corresponding builtin type (e.g., list[int]).


def _alt_inner_shape(self, new_inner_rank):
    if new_inner_rank == 0:
        raise ValueError('new_inner_rank cannot be zero')
    elif self.inner_rank == 0:
        raise ValueError('old inner_rank cannot be zero')
    elif new_inner_rank == self.inner_rank:
        return self.inner_shape
    elif new_inner_rank < self.inner_rank:
        if self._static_inner_shape.is_fully_defined():
            return _alt_inner_shape_from_tensor_shape(self._static_inner_shape, self.dtype, new_inner_rank)
        first_dimension = self._num_slices_in_dimension(-new_inner_rank)
        if new_inner_rank == 1:
            return array_ops.expand_dims(first_dimension, 0)
        remaining_dimensions = self.inner_shape[1 - new_inner_rank:]
        return array_ops.concat([array_ops.expand_dims(first_dimension, 0), remaining_dimensions], axis=0)
    else:
        assert new_inner_rank > self.inner_rank
        new_dimensions = new_inner_rank - self.inner_rank
        if any([not x.is_uniform() for x in self.row_partitions[-new_dimensions:]]):
            raise ValueError('Cannot get an inner shape over a ragged dimension')
        first_dimension = self._num_slices_in_dimension(-new_inner_rank)
        new_dimensions = new_inner_rank - self.inner_rank
        new_dims = [first_dimension] + [x.uniform_row_length() for x in self.row_partitions[-new_dimensions:]]
        return array_ops.concat([array_ops_stack.stack(new_dims), self.inner_shape[1:]], axis=0)

Get an alternative inner shape with higher or lower rank.

    For the rank of the inner shape to be be higher, the last few ragged
    dimensions must have uniform_row_length.

Args:
        new_inner_rank: the new rank of the inner_shape

Returns:
        A new inner_shape of rank new_inner_rank.


def from_sub_models_config(cls, text_encoder_config: PretrainedConfig, audio_encoder_config: PretrainedConfig, decoder_config: MusicgenMelodyDecoderConfig, **kwargs):
    return cls(text_encoder=text_encoder_config.to_dict(), audio_encoder=audio_encoder_config.to_dict(), decoder=decoder_config.to_dict(), **kwargs)

Instantiate a [`MusicgenMelodyConfig`] (or a derived class) from text encoder, audio encoder and decoder
    configurations.

Returns:
        [`MusicgenMelodyConfig`]: An instance of a configuration object


def assert_no_legacy_layers(layers):
    legacy_layers = [l for l in layers if getattr(l, '_is_legacy_layer', None)]
    if legacy_layers:
        layer_str = '\n'.join(('  ' + str(l) for l in legacy_layers))
        raise TypeError('The following are legacy tf.layers.Layers:\n{}\nTo use keras as a framework (for instance using the Network, Model, or Sequential classes), please use the tf.keras.layers implementation instead. (Or, if writing custom layers, subclass from tf.keras.layers rather than tf.layers)'.format(layer_str))

Prevent tf.layers.Layers from being used with Keras.

    Certain legacy layers inherit from their keras analogs; however they are
    not supported with keras and can lead to subtle and hard to diagnose bugs.

Args:
        layers: A list of layers to check

Raises:
        TypeError: If any elements of layers are tf.layers.Layers


def get_full_path(path):
    if path_utils.isabs(path):
        return path
    else:
        return path_utils.join(_pytype_source_dir(), path)

Full path to a file or directory within the pytype source tree.

Args:
        path: An absolute or relative path.

Returns:
        path for absolute paths.
        full path resolved relative to pytype/ for relative paths.


def set_distribution(value):
    global_state.set_global_attribute(GLOBAL_ATTRIBUTE_NAME, value)

Set the distribution as the global distribution setting.

Args:
        value: a `Distribution` instance.


def get_prefix_bias_name(self) -> Union[None, str]:
    warnings.warn('The method get_prefix_bias_name is deprecated. Please use `get_bias` instead.', FutureWarning)
    return None

Get the concatenated _prefix name of the bias from the model name to the parent layer

Returns:
        `str`: The _prefix name of the bias.


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

    @parameterized.named_parameters(('_v1_graph', 'v1_graph'), ('_function', 'function'))
    @functools.wraps(func)
    def decorated(self: 'TensorFlowTestCase', run_mode: str, *args, **kwargs) -> None:
        if run_mode == 'v1_graph':
            with ops.Graph().as_default():
                func(self, *args, **kwargs)
        elif run_mode == 'function':

            @def_function.function
            def function_in_eager():
                func(self, *args, **kwargs)
            graph_for_eager_test = ops.Graph()
            with graph_for_eager_test.as_default(), context.eager_mode():
                function_in_eager()
            ops.dismantle_graph(graph_for_eager_test)
        else:
            raise ValueError('Unknown run mode %s' % run_mode)
    return decorated

Run a test case in v1 graph mode and inside tf.function in eager mode.

    WARNING: This decorator can only be used in test cases that statically checks
    generated graph. Attempting to evaluate graph or function results via.
    session.run() or self.evaluate() will fail.

    WARNING: This decorator can only be used for test cases that inherit from
    absl.testing.parameterized.TestCase.

Args:
        func: Test case function to be decorated.

Returns:
        Decorated test case function.


def has_leak(output: EventSetNodeCollection, input: Optional[EventSetNodeCollection]=None) -> bool:
    if input is None:
        normalized_input = None
    else:
        normalized_input = _normalize_query(input)
    normalized_output = _normalize_query(output)
    graph = infer_graph(inputs=normalized_input, outputs=normalized_output)
    leak_key = LeakOperator.operator_key()
    for operator in graph.operators:
        if operator.operator_key() == leak_key:
            return True
    return False

Tests if a node depends on a leak operator.

    Tests if a [`EventSetNode`][temporian.EventSetNode] or collection of nodes
    depends on the only operator that can introduce future leakage:
    [`EventSet.leak()`][temporian.EventSet.leak].

    Single input output example:
        ```python
        >>> a = tp.input_node([("f", tp.float32)])
        >>> b = a.moving_sum(5)
        >>> c = b.leak(6)
        >>> d = c.prefix("my_prefix_")
        >>> e = d.moving_sum(7)
        >>> # The computation of "e" contains a leak.
        >>> tp.has_leak(e)
        True
        >>> # The computation of "e" given "d" does not contain a leak.
        >>> tp.has_leak(e, d)
        False

        ```

Args:
        output: Nodes to compute. Supports Node, dict of Nodes and list of
            Nodes.
        input: Optional input nodes. Supports Node, dict of Nodes and list of
            Nodes. If not specified, assumes for the input nodes to be the the
            raw data inputs e.g. [`tp.input_node()`][temporian.input_node] and
            [`tp.event_set()`][temporian.event_set].

Returns:
        True if and only if the computation of `output` from `inputs` depends
        on a [`EventSet.leak()`][temporian.EventSet.leak] operator.


def get_feature_extractor_dict(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> tuple[dict[str, Any], dict[str, Any]]:
    cache_dir = kwargs.pop('cache_dir', None)
    force_download = kwargs.pop('force_download', False)
    resume_download = kwargs.pop('resume_download', None)
    proxies = kwargs.pop('proxies', None)
    subfolder = kwargs.pop('subfolder', None)
    token = kwargs.pop('token', None)
    use_auth_token = kwargs.pop('use_auth_token', None)
    local_files_only = kwargs.pop('local_files_only', False)
    revision = kwargs.pop('revision', None)
    if use_auth_token is not None:
        warnings.warn('The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.', FutureWarning)
        if token is not None:
            raise ValueError('`token` and `use_auth_token` are both specified. Please set only the argument `token`.')
        token = use_auth_token
    from_pipeline = kwargs.pop('_from_pipeline', None)
    from_auto_class = kwargs.pop('_from_auto', False)
    user_agent = {'file_type': 'feature extractor', 'from_auto_class': from_auto_class}
    if from_pipeline is not None:
        user_agent['using_pipeline'] = from_pipeline
    if is_offline_mode() and (not local_files_only):
        logger.info('Offline mode: forcing local_files_only=True')
        local_files_only = True
    pretrained_model_name_or_path = str(pretrained_model_name_or_path)
    is_local = os.path.isdir(pretrained_model_name_or_path)
    if os.path.isdir(pretrained_model_name_or_path):
        feature_extractor_file = os.path.join(pretrained_model_name_or_path, FEATURE_EXTRACTOR_NAME)
    if os.path.isfile(pretrained_model_name_or_path):
        resolved_feature_extractor_file = pretrained_model_name_or_path
        is_local = True
    elif is_remote_url(pretrained_model_name_or_path):
        feature_extractor_file = pretrained_model_name_or_path
        resolved_feature_extractor_file = download_url(pretrained_model_name_or_path)
    else:
        feature_extractor_file = FEATURE_EXTRACTOR_NAME
        try:
            resolved_feature_extractor_file = cached_file(pretrained_model_name_or_path, feature_extractor_file, cache_dir=cache_dir, force_download=force_download, proxies=proxies, resume_download=resume_download, local_files_only=local_files_only, subfolder=subfolder, token=token, user_agent=user_agent, revision=revision)
        except OSError:
            raise
        except Exception:
            raise OSError(f"Can't load feature extractor for '{pretrained_model_name_or_path}'. If you were trying to load it from 'https://huggingface.co/models', make sure you don't have a local directory with the same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory containing a {FEATURE_EXTRACTOR_NAME} file")
    try:
        with open(resolved_feature_extractor_file, encoding='utf-8') as reader:
            text = reader.read()
        feature_extractor_dict = json.loads(text)
    except json.JSONDecodeError:
        raise OSError(f"It looks like the config file at '{resolved_feature_extractor_file}' is not a valid JSON file.")
    if is_local:
        logger.info(f'loading configuration file {resolved_feature_extractor_file}')
    else:
        logger.info(f'loading configuration file {feature_extractor_file} from cache at {resolved_feature_extractor_file}')
    return (feature_extractor_dict, kwargs)

From a `pretrained_model_name_or_path`, resolve to a dictionary of parameters, to be used for instantiating a
    feature extractor of type [`~feature_extraction_utils.FeatureExtractionMixin`] using `from_dict`.

    Parameters:
        pretrained_model_name_or_path (`str` or `os.PathLike`):
            The identifier of the pre-trained checkpoint from which we want the dictionary of parameters.

Returns:
        `Tuple[Dict, Dict]`: The dictionary(ies) that will be used to instantiate the feature extractor object.


def forward(self, x):
    if self.training and self.drop_prob > 0:
        return XDropout.apply(x, self.get_context())
    return x

Call the module

Args:
        x (`torch.tensor`): The input tensor to apply dropout


def _TestGetItem(self, struct, slice_spec, expected):
    tensor_slice_spec1 = _make_tensor_slice_spec(slice_spec, True)
    tensor_slice_spec2 = _make_tensor_slice_spec(slice_spec, False)
    value1 = struct.__getitem__(slice_spec)
    value2 = struct.__getitem__(tensor_slice_spec1)
    value3 = struct.__getitem__(tensor_slice_spec2)
    self.assertAllEqual(value1, expected, 'slice_spec=%s' % (slice_spec,))
    self.assertAllEqual(value2, expected, 'slice_spec=%s' % (slice_spec,))
    self.assertAllEqual(value3, expected, 'slice_spec=%s' % (slice_spec,))

Helper function for testing StructuredTensor.__getitem__.

    Checks that calling `struct.__getitem__(slice_spec) returns the expected
    value.  Checks three different configurations for each slice spec:

        * Call __getitem__ with the slice spec as-is (with int values)
        * Call __getitem__ with int values in the slice spec wrapped in
        `tf.constant()`.
        * Call __getitem__ with int values in the slice spec wrapped in
        `tf.compat.v1.placeholder()` (so value is not known at graph
        construction time).

Args:
        struct: The StructuredTensor to test.
        slice_spec: The slice spec.
        expected: The expected value of struct.__getitem__(slice_spec), as a
        python list.


def clear_signature_defs(tflite_model):
    model = tflite_model
    if not isinstance(tflite_model, bytearray):
        model = bytearray(tflite_model)
    return signature_def_util.ClearSignatureDefs(model)

Clears SignatureDefs from the Metadata of a TfLite flatbuffer buffer.

Args:
        tflite_model: TFLite model buffer to remove signature_defs.

Returns:
        buffer: A TFLite model binary identical to model buffer with
        no SignatureDef metadata.

Raises:
        ValueError:
        tflite_model buffer does not contain a valid TFLite model.


def DummyMethod(name, *params):

    def make_param(param):
        return pytd.Parameter(param, type=pytd.AnythingType(), kind=pytd.ParameterKind.REGULAR, optional=False, mutated_type=None)
    sig = pytd.Signature(tuple((make_param(param) for param in params)), starargs=None, starstarargs=None, return_type=pytd.AnythingType(), exceptions=(), template=())
    return pytd.Function(name=name, signatures=(sig,), kind=pytd.MethodKind.METHOD, flags=pytd.MethodFlag.NONE)

Create a simple method using only "Any"s as types.

Args:
        name: The name of the method
        *params: The parameter names.

Returns:
        A pytd.Function.


def __init__(self, query_engine: engine.Engine, value_set_codes_table: str) -> None:
    super().__init__()
    self._client = query_engine
    self._value_set_codes_table = value_set_codes_table

Initializes the SparkValueSetManager with user-provided sqlalchemy engine.

Args:
        query_engine: Sqlalchemy engine to communicate with Spark cluster
        value_set_codes_table: A table containing value set expansions. If
        `value_set_codes_table` is a string, it must included a project ID if
        not in the client's default project, dataset ID and table ID, each
        separated by a '.'. The table must match the schema described by
        https://github.com/FHIR/sql-on-fhir/blob/master/sql-on-fhir.md#valueset-support


def slice_inputs(self, indices_dataset, inputs):
    dataset = dataset_ops.DatasetV2.zip((indices_dataset, dataset_ops.DatasetV2.from_tensors(inputs).repeat()))

    def grab_batch(i, data):
        return nest.map_structure(lambda d: array_ops.gather(d, i, axis=0), data)
    dataset = dataset.map(grab_batch, num_parallel_calls=dataset_ops.AUTOTUNE)
    options = options_lib.Options()
    options.experimental_optimization.apply_default_optimizations = False
    if self._shuffle:
        options.experimental_external_state_policy = options_lib.ExternalStatePolicy.IGNORE
    dataset = dataset.with_options(options)
    return dataset

Slice inputs into a Dataset of batches.

    Given a Dataset of batch indices and the unsliced inputs,
    this step slices the inputs in a parallelized fashion
    and produces a dataset of input batches.

Args:
        indices_dataset: A Dataset of batched indices
        inputs: A python data structure that contains the inputs, targets,
        and possibly sample weights.

Returns:
        A Dataset of input batches matching the batch indices.


def index_of(self, file_path, line_number, called_function_name, called_file_path, called_function_start_line):
    location_key = (file_path, called_function_name, line_number)
    if location_key in self._location_key_to_location:
        location = self._location_key_to_location[location_key]
        return location.id
    else:
        location_index = len(self._location_key_to_location) + 1
        location = profile_pb2.Location()
        location.id = location_index
        self._location_key_to_location[location_key] = location
        line = location.line.add()
        line.function_id = self._functions.index_of(called_file_path, called_function_name, called_function_start_line)
        line.line = line_number
        return location_index

Returns index of the location, adding the location if needed.

Args:
        file_path: (string) Path to file that makes the call.
        line_number: (integer) Call line number.
        called_function_name: (string) Function name of the function called at
        `file_path` and `line_number`.
        called_file_path: (string) Path to file where the called function is
        defined.
        called_function_start_line: (integer) Start line number of called
        function definition in `called_file_path` file.

Returns:
        Index of location.


def pack(value):
    if is_packed(value):
        return value
    spec = value._type_spec._tf_extension_type_with_packed(True)
    try:
        variant = composite_tensor_ops.composite_tensor_to_variants(value)
    except nested_structure_coder.NotEncodableError as e:
        raise ValueError('ExtensionTypes must have a __name__ field in order to be packed.') from e
    return _create_object_from_type_and_dict(type(value), {'_tf_extension_type_cached_type_spec': spec, '_tf_extension_type_packed_variant': variant})

Returns a copy of `value` with fields packed in a single Variant.

Args:
        value: An `ExtensionType` object.

Returns:
        An `ExtensionType` object.


def set_size(a, validate_indices=True):
    a = sparse_tensor.convert_to_tensor_or_sparse_tensor(a, name='a')
    if not isinstance(a, sparse_tensor.SparseTensor):
        raise TypeError('Expected `SparseTensor`, got %s.' % a)
    if a.values.dtype.base_dtype not in _VALID_DTYPES:
        raise TypeError(f'Invalid dtype `{a.values.dtype}` not in supported dtypes: `{_VALID_DTYPES}`.')
    return gen_set_ops.set_size(a.indices, a.values, a.dense_shape, validate_indices)

Compute number of unique elements along last dimension of `a`.

Args:
        a: `SparseTensor`, with indices sorted in row-major order.
        validate_indices: Whether to validate the order and range of sparse indices
        in `a`. Note that setting this to `false` allows for undefined behavior
        when calling this function with invalid indices.

Returns:
        `int32` `Tensor` of set sizes. For `a` ranked `n`, this is a `Tensor` with
        rank `n-1`, and the same 1st `n-1` dimensions as `a`. Each value is the
        number of unique elements in the corresponding `[0...n-1]` dimension of `a`.

Raises:
        TypeError: If `a` is an invalid types.


def get(self):
    while self.is_running():
        try:
            inputs = self.ready_queue.get(block=False)
            yield inputs
            continue
        except queue.Empty:
            pass
        try:
            value = self.future_queue.get(block=True, timeout=5)
            self.future_queue.task_done()
            if isinstance(value, Exception):
                raise value
            inputs = value.get()
            if inputs is not None:
                yield inputs
        except queue.Empty:
            pass
        except Exception as e:
            self.stop(drain_queue_and_join=True)
            raise e
    raise ValueError('Iterator called after `on_epoch_end` or before `on_epoch_begin`.')

Creates a generator to extract data from the queue.

    Skip the data if it is `None`.

    This method is called from the main thread.

Yields:
        The next element in the queue, i.e. a tuple
        `(inputs, targets)` or
        `(inputs, targets, sample_weights)`.


Create an accumulator based on 'output'.

    This method creates a new accumulator with identical internal state to the
    one used to create the data in 'output'. This means that if you do

    output_data = combiner.extract(accumulator_1)
    accumulator_2 = combiner.restore(output_data)

    then accumulator_1 and accumulator_2 will have identical internal state, and
    computations using either of them will be equivalent.

Args:
        output: The data output from a previous computation. Should be in the same
        form as provided by 'extract_output'.

Returns:
        A new accumulator.


def __init__(self, *content: WritableTypes, style_files: Optional[Iterable[str]]=None, styles: Optional[Iterable[str]]=None, script_files: Optional[Iterable[str]]=None, scripts: Optional[Iterable[str]]=None) -> None:
    super().__init__(*content, style_files=Html.StyleFiles(*(style_files or [])), styles=Html.Styles(*(styles or [])), script_files=Html.ScriptFiles(*(script_files or [])), scripts=Html.Scripts(*(scripts or [])))

Constructor.

Args:
        *content: One or multiple body part (str, Html, lambda, None) of the HTML.
        style_files: URLs for external styles to include.
        styles: CSS styles to include.
        script_files: URLs for external scripts to include.
        scripts: JavaScript scripts to include.


def sparse_top_k_categorical_accuracy(y_true, y_pred, k=5):
    y_pred_rank = tensor_conversion.convert_to_tensor_v2_with_dispatch(y_pred).shape.ndims
    y_true_rank = tensor_conversion.convert_to_tensor_v2_with_dispatch(y_true).shape.ndims
    if y_true_rank is not None and y_pred_rank is not None:
        if y_pred_rank > 2:
            y_pred = array_ops.reshape(y_pred, [-1, y_pred.shape[-1]])
        if y_true_rank > 1:
            y_true = array_ops.reshape(y_true, [-1])
    return math_ops.cast(nn.in_top_k(y_pred, math_ops.cast(y_true, 'int32'), k), backend.floatx())

Computes how often integer targets are in the top `K` predictions.

    Standalone usage:
    >>> y_true = [2, 1]
    >>> y_pred = [[0.1, 0.9, 0.8], [0.05, 0.95, 0]]
    >>> m = tf.keras.metrics.sparse_top_k_categorical_accuracy(
    ...     y_true, y_pred, k=3)
    >>> assert m.shape == (2,)
    >>> m.numpy()
    array([1., 1.], dtype=float32)

Args:
        y_true: tensor of true targets.
        y_pred: tensor of predicted targets.
        k: (Optional) Number of top elements to look at for computing accuracy.
        Defaults to 5.

Returns:
        Sparse top K categorical accuracy value.


def ensure_value_to_cell(value):

    def dummy_fn():
        value
    cell_value = dummy_fn.__closure__[0]
    if not isinstance(value, type(cell_value)):
        return cell_value
    return value

Ensures that a value is converted to a python cell object.

Args:
        value: Any value that needs to be casted to the cell type

Returns:
        A value wrapped as a cell object (see function "func_load")


def __init__(self, pipeline: Union[beam_runner_api_pb2.Pipeline, beam.Pipeline], default_vertex_attrs={'shape': 'box'}, default_edge_attrs=None, render_option=None):
    self._lock = threading.Lock()
    self._graph: pydot.Dot = None
    self._pipeline_instrument = None
    if isinstance(pipeline, beam.Pipeline):
        self._pipeline_instrument = inst.PipelineInstrument(pipeline, pipeline._options)
        self._pipeline_instrument.preprocess()
    if isinstance(pipeline, beam_runner_api_pb2.Pipeline):
        self._pipeline_proto = pipeline
    elif isinstance(pipeline, beam.Pipeline):
        self._pipeline_proto = pipeline.to_runner_api()
    else:
        raise TypeError('pipeline should either be a %s or %s, while %s is given' % (beam_runner_api_pb2.Pipeline, beam.Pipeline, type(pipeline)))
    self._consumers: DefaultDict[str, List[str]] = collections.defaultdict(list)
    self._producers: Dict[str, str] = {}
    for transform_id, transform_proto in self._top_level_transforms():
        for pcoll_id in transform_proto.inputs.values():
            self._consumers[pcoll_id].append(transform_id)
        for pcoll_id in transform_proto.outputs.values():
            self._producers[pcoll_id] = transform_id
    default_vertex_attrs = default_vertex_attrs or {'shape': 'box'}
    if 'color' not in default_vertex_attrs:
        default_vertex_attrs['color'] = 'blue'
    if 'fontcolor' not in default_vertex_attrs:
        default_vertex_attrs['fontcolor'] = 'blue'
    vertex_dict, edge_dict = self._generate_graph_dicts()
    self._construct_graph(vertex_dict, edge_dict, default_vertex_attrs, default_edge_attrs)
    self._renderer = pipeline_graph_renderer.get_renderer(render_option)

Constructor of PipelineGraph.

Example:
        graph = pipeline_graph.PipelineGraph(pipeline_proto)
        graph.get_dot()

        or

        graph = pipeline_graph.PipelineGraph(pipeline)
        graph.get_dot()

Args:
        pipeline: (Pipeline proto) or (Pipeline) pipeline to be rendered.
        default_vertex_attrs: (Dict[str, str]) a dict of default vertex attributes
        default_edge_attrs: (Dict[str, str]) a dict of default edge attributes
        render_option: (str) this parameter decides how the pipeline graph is
        rendered. See display.pipeline_graph_renderer for available options.


def get_maybe_abstract_instance(self, data):
    if data.is_concrete:
        data_type = type(data.pyval)
        if data_type in self.primitive_instances:
            return self.primitive_instances[data_type]
    return data

Get an instance of the same type as the given data, abstract if possible.

    Get an abstract instance of primitive data stored as a
    ConcreteValue. Return any other data as-is. This is used by
    constant_to_var to discard concrete values that have been kept
    around for InterpreterFunction.

    This method intentionally does not descend into containers, as doing so
    causes new timeouts. If you need to discard concrete values inside
    containers, use abstract_utils.abstractify_variable instead.

Args:
        data: The data.

Returns:
        An instance of the same type as the data, abstract if possible.


def adb_call(args=None, shell=False, timeout=None, stderr=None) -> bytes:
    return self._exec_adb_cmd(name, args, shell=shell, timeout=timeout, stderr=stderr)

Wrapper for an ADB command.

Args:
        args: string or list of strings, arguments to the adb command.
        See subprocess.Proc() documentation.
        shell: bool, True to run this command through the system shell,
        False to invoke it directly. See subprocess.Proc() docs.
        timeout: float, the number of seconds to wait before timing out.
        If not specified, no timeout takes effect.
        stderr: a Byte stream, like io.BytesIO, stderr of the command
        will be written to this object if provided.

Returns:
        The output of the adb command run if exit code is 0.


def dates(self):
    return _gen_business_days(self._start_date, self._end_date, self._holiday_calendar, self._backward)

Returns the dates as computed from the schedule as a DateTensor.

    Constructs the date schedule from the supplied data. For more details see
    the initializer docstring.

Returns:
        `DateTensor` of rank one more than `start_date` or `end_date`
        (depending on `backwards`), representing schedules for each element
        of the input.


def barrier(mesh: layout.Mesh, barrier_name: Optional[str]=None, timeout_in_ms: Optional[int]=None):
    if barrier_name is None:
        barrier_name = '(barrier)'
    logging.info('entering barrier before op: %s', barrier_name)
    context.async_wait()
    component = array_ops.reshape(1.0, [1] * len(mesh.shape()))
    ones = api.pack([component] * mesh.num_local_devices(), layout.Layout(mesh.dim_names, mesh))
    mesh_size = math_ops.reduce_sum(ones)
    if mesh_size != mesh.size:
        raise ValueError('Global barrier produced wrong mesh size : {0} while mesh has actualsize : {1}'.format(mesh_size, mesh.size))
    context.async_wait()
    if context.context().coordination_service:
        if timeout_in_ms is None:
            timeout_in_ms = 24 * 60 * 60 * 1000
        num_calls = _BARRIER_DICT.setdefault(barrier_name, 0)
        _BARRIER_DICT[barrier_name] = num_calls + 1
        barrier_id = f'{barrier_name}:{num_calls}'
        context.context().wait_at_barrier(barrier_id, timeout_in_ms)
    logging.info('finished running barrier across all clients after op: %s', barrier_name)

Runs a barrier on the mesh.

    Upon returning from the barrier, all operations run before the barrier
    would have completed across all clients. Currently we allocate a fully
    sharded tensor with mesh shape and run an all_reduce on it.

Example:
    A barrier can be used before application exit to ensure completion of pending
    ops.

    ```python

    x = [1, 2, 3]
    x = dtensor.relayout(x, dtensor.Layout.batch_sharded(mesh, 'batch', 1))
    dtensor.barrier(mesh)

    # At this point all devices on all clients in the mesh have completed
    # operations before the barrier. Therefore it is OK to tear down the clients.
    sys.exit()
    ```

Args:
        mesh: The mesh to run the barrier on.
        barrier_name: The name of the barrier. Mainly used for logging purpose.
        timeout_in_ms: The timeout of the barrier in ms. If omitted, blocks
        indefinitely till the barrier is reached from all clients.


def _get_object_type(filename, filepath):
    filename_no_ext = os.path.splitext(filename)[0].lower()
    if filename_no_ext.endswith(PrecompiledExampleType.test_ends):
        object_type = PRECOMPILED_OBJECT_TYPE_UNIT_TEST
    elif PrecompiledExampleType.katas in filepath.split(os.sep):
        object_type = PRECOMPILED_OBJECT_TYPE_KATA
    elif PrecompiledExampleType.examples in filepath.split(os.sep):
        object_type = PRECOMPILED_OBJECT_TYPE_EXAMPLE
    else:
        object_type = PRECOMPILED_OBJECT_TYPE_UNSPECIFIED
    return object_type

Get type of an object based on it filename/filepath

Args:
        filename: object's filename
        filepath: object's filepath

    Returns: type of the object (example, kata, unit-test)


def min(x, axis=None, keepdims=False):
    return math_ops.reduce_min(x, axis, keepdims)

Minimum value in a tensor.

Args:
        x: A tensor or variable.
        axis: An integer, the axis to find minimum values.
        keepdims: A boolean, whether to keep the dimensions or not.
            If `keepdims` is `False`, the rank of the tensor is reduced
            by 1. If `keepdims` is `True`,
            the reduced dimension is retained with length 1.

Returns:
        A tensor with minimum values of `x`.


def get_extended_attention_mask(self, attention_mask: Tensor, input_shape: Tuple[int], device: device, is_decoder: bool) -> Tensor:
    if attention_mask.dim() == 3:
        extended_attention_mask = attention_mask[:, None, :, :]
    elif attention_mask.dim() == 2:
        if is_decoder:
            batch_size, seq_length = input_shape
            seq_ids = torch.arange(seq_length, device=device)
            causal_mask = seq_ids[None, None, :].repeat(batch_size, seq_length, 1) <= seq_ids[None, :, None]
            causal_mask = causal_mask.to(attention_mask.dtype)
            if causal_mask.shape[1] < attention_mask.shape[1]:
                prefix_seq_len = attention_mask.shape[1] - causal_mask.shape[1]
                causal_mask = torch.cat([torch.ones((batch_size, seq_length, prefix_seq_len), device=device, dtype=causal_mask.dtype), causal_mask], axis=-1)
            extended_attention_mask = causal_mask[:, None, :, :] * attention_mask[:, None, None, :]
        else:
            extended_attention_mask = attention_mask[:, None, None, :]
    else:
        raise ValueError('Wrong shape for input_ids (shape {}) or attention_mask (shape {})'.format(input_shape, attention_mask.shape))
    extended_attention_mask = extended_attention_mask.to(dtype=self.dtype)
    extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
    return extended_attention_mask

Makes broadcastable attention and causal masks so that future and masked tokens are ignored.

Args:
        attention_mask (`torch.Tensor`):
            Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
        input_shape (`Tuple[int]`):
            The shape of the input to the model.
        device (`torch.device`):
            The device of the input to the model.

Returns:
        `torch.Tensor` The extended attention mask, with a the same dtype as `attention_mask.dtype`.


def peek(self, index, name=None):
    if name is None:
        name = '%s_peek' % self._name
    fn = lambda: gen_data_flow_ops.stage_peek(index, dtypes=self._dtypes, shared_name=self._name, name=name, capacity=self._capacity, memory_limit=self._memory_limit)
    return self.__internal_get(fn, name)

Peeks at an element in the staging area.

    If the staging area is too small to contain the element at
    the specified index, it will block until enough elements
    are inserted to complete the operation.

    The placement of the returned tensor will be determined by
    the current device scope when this function is called.

Args:
        index: The index of the tensor within the staging area
            to look up.
        name: A name for the operation (optional).

Returns:
        The tuple of tensors that was gotten.


def add_callback(self, callback):
    self.callback_handler.add_callback(callback)

Add a callback to the current list of [`~transformers.TrainerCallback`].

Args:
        callback (`type` or [`~transformers.TrainerCallback]`):
        A [`~transformers.TrainerCallback`] class or an instance of a [`~transformers.TrainerCallback`]. In the
        first case, will instantiate a member of that class.


def get_identity_broadcaster(cls, nvals, dtype=None):
    return _GatherLayerBroadcaster(math_ops.range(nvals, dtype=dtype))

Create an identity broadcaster.

    TODO(martinz): an identity broadcaster can be far more efficient than a
    generic broadcaster. Add an optimized implementation.

Args:
        nvals: the number of values for the broadcaster.
        dtype: the dtype of the broadcaster, or None to use the dtype of nvals.

Returns:
        an identity broadcaster from [0....nvals-1] to [0...nvals-1]


def get_group_key(self, devices):
    with self._lock:
        devices_key = ','.join(devices)
        if devices_key not in self._known_groups:
            self._known_groups[devices_key] = self._get_new_group_key(devices)
        return self._known_groups[devices_key]

Returns a group key for the list of local devices.

    The same group key is returned if the list of local devices is the same.

Args:
        devices: a list of local canonical device strings in a collective group.

Returns:
        a group key.


def set_image_data_format(data_format):
    global _IMAGE_DATA_FORMAT
    data_format = str(data_format).lower()
    if data_format not in {'channels_first', 'channels_last'}:
        raise ValueError(f"The `data_format` argument must be one of {{'channels_first', 'channels_last'}}. Received: data_format={data_format}")
    _IMAGE_DATA_FORMAT = data_format

Set the value of the image data format convention.

Args:
        data_format: string. `'channels_first'` or `'channels_last'`.

Example:
    >>> keras.config.image_data_format()
    'channels_last'

    >>> keras.config.set_image_data_format('channels_first')
    >>> keras.config.image_data_format()
    'channels_first'

    >>> # Set it back to `'channels_last'`
    >>> keras.config.set_image_data_format('channels_last')


def matmul(a: ragged_tensor.RaggedOrDense, b: ragged_tensor.RaggedOrDense, transpose_a=False, transpose_b=False, adjoint_a=False, adjoint_b=False, a_is_sparse=False, b_is_sparse=False, output_type=None, grad_a=False, grad_b=False, name=None):
    del grad_a
    del grad_b
    if transpose_a and adjoint_a:
        raise ValueError('Only one of transpose_a and adjoint_a can be True.')
    if transpose_b and adjoint_b:
        raise ValueError('Only one of transpose_b and adjoint_b can be True.')
    kwargs = dict(transpose_a=transpose_a, transpose_b=transpose_b, adjoint_a=adjoint_a, adjoint_b=adjoint_b, a_is_sparse=a_is_sparse, b_is_sparse=b_is_sparse, output_type=output_type)
    with ops.name_scope(name, 'RaggedMatMul', [a, b]) as name:
        a = ragged_tensor.convert_to_tensor_or_ragged_tensor(a, name='a')
        b = ragged_tensor.convert_to_tensor_or_ragged_tensor(b, name='b')
        a_is_ragged = isinstance(a, ragged_tensor.RaggedTensor)
        b_is_ragged = isinstance(b, ragged_tensor.RaggedTensor)
        if not (a_is_ragged or b_is_ragged):
            return math_ops.matmul(a, b, **kwargs)
        if a.dtype != b.dtype:
            raise ValueError('`a` and `b` must have the same dtype.')
        if a.shape.rank is None:
            if b.shape.rank is None:
                raise ValueError('matmul requires at least one input to have known rank if either input is ragged.')
            rank = b.shape.rank
        else:
            if b.shape.rank is not None and a.shape.rank != b.shape.rank:
                raise ValueError('`a` and `b` must have the same rank.')
            rank = a.shape.rank
        if rank < 2:
            raise ValueError('`a` and `b` must have the same rank.')
        if rank > 3:
            shape_err = 'Batch dimensions of `a` and `b` do not have the same size.'
            if not a_is_ragged:
                a = ragged_tensor.RaggedTensor.from_tensor(a, ragged_rank=1)
            if not b_is_ragged:
                b = ragged_tensor.RaggedTensor.from_tensor(b, ragged_rank=1)
            with ops.control_dependencies([check_ops.assert_equal(a.row_splits, b.row_splits, message=shape_err)]):
                flat_result = matmul(a.values, b.values, **kwargs)
                return a.with_values(flat_result)
        if rank == 2:
            return _matmul_2d(a, b, **kwargs)
        assert rank == 3
        a_ragged_rank = a.ragged_rank if a_is_ragged else 0
        if a_ragged_rank == 1 and (not (b_is_ragged or transpose_a or adjoint_a)):
            return _matmul_3d_with_batch_dim_folding(a, b, **kwargs)
        else:
            return _matmul_3d_with_map_fn(a, b, **kwargs)

Multiplies matrix `a` by matrix `b`.

    If all transpose or adjoint attributes are `False` then:

    ```
    output[..., i, j] = sum_k (a[..., i, k] * b[..., k, j]), for all indices i, j.
    ```

    The inputs `a` and `b` must have `rank >= 2`, where the outermost `rank - 2`
    dimensions are batch dimensions.  The inputs must have the same dtype.  See
    `tf.matmul` for more information.

Args:
        a: `tf.Tensor` or `RaggedTensor` with `rank > 1`.
        b: `tf.Tensor` or `RaggedTensor` with same type and rank as `a`.
        transpose_a: If `True`, `a` is transposed before multiplication.
        transpose_b: If `True`, `b` is transposed before multiplication.
        adjoint_a: If `True`, `a` is conjugated & transposed before multiplication.
        adjoint_b: If `True`, `b` is conjugated & transposed before multiplication.
        a_is_sparse: If `True`, optimize assuming `a` is mostly zero.
        b_is_sparse: If `True`, optimize assuming `b` is mostly zero.
        output_type: The output datatype (optional).
        grad_a: Unused.
        grad_b: Unused.
        name: Name for the operation (optional).

Returns:
        A `Tensor` or `RaggedTensor` with the same rank and shape as `a`, where
        each inner-most matrix is the product of the corresponding matrices in `a`
        and `b`.


def resolve(self, reference_path_or_paths: Optional[Union[str, List[str]]]=None) -> Union[Tuple[symbolic.Symbolic, utils.KeyPath], List[Tuple[symbolic.Symbolic, utils.KeyPath]]]:
    single_input = False
    if reference_path_or_paths is None:
        reference_paths = self.reference_paths
    elif isinstance(reference_path_or_paths, str):
        reference_paths = [utils.KeyPath.parse(reference_path_or_paths)]
        single_input = True
    elif isinstance(reference_path_or_paths, utils.KeyPath):
        reference_paths = [reference_path_or_paths]
        single_input = True
    elif isinstance(reference_path_or_paths, list):
        paths = []
        for path in reference_path_or_paths:
            if isinstance(path, str):
                path = utils.KeyPath.parse(path)
            elif not isinstance(path, utils.KeyPath):
                raise ValueError("Argument 'reference_path_or_paths' must be None, a string, KeyPath object, a list of strings, or a list of KeyPath objects.")
            paths.append(path)
        reference_paths = paths
    else:
        raise ValueError("Argument 'reference_path_or_paths' must be None, a string, KeyPath object, a list of strings, or a list of KeyPath objects.")
    resolved_paths = []
    for reference_path in reference_paths:
        parent = self.sym_parent
        while parent is not None and (not reference_path.exists(parent)):
            parent = getattr(parent, 'sym_parent', None)
        if parent is None:
            raise ValueError(f"Cannot resolve '{reference_path}': parent not found.")
        resolved_paths.append((parent, parent.sym_path + reference_path))
    return resolved_paths if not single_input else resolved_paths[0]

Resolve reference paths based on the location of this node.

Args:
        reference_path_or_paths: (Optional) a string or KeyPath as a reference
        path or a list of strings or KeyPath objects as a list of
        reference paths.
        If this argument is not provided, prebound reference paths of this
        object will be used.

Returns:
        A tuple (or list of tuple) of (resolved parent, resolved full path)


def __init__(self, graph, control_inputs) -> None:
    self._graph = graph
    if control_inputs is None:
        self._control_inputs_val = []
        self._new_stack = True
    else:
        self._control_inputs_val = control_inputs
        self._new_stack = False
    self._seen_nodes = set()
    self._old_stack = None
    self._old_control_flow_context = None

Create a new `_ControlDependenciesController`.

    A `_ControlDependenciesController` is the context manager for
    `with tf.control_dependencies()` blocks.  These normally nest,
    as described in the documentation for `control_dependencies()`.

    The `control_inputs` argument list control dependencies that must be
    added to the current set of control dependencies.  Because of
    uniquification the set can be empty even if the caller passed a list of
    ops.  The special value `None` indicates that we want to start a new
    empty set of control dependencies instead of extending the current set.

    In that case we also clear the current control flow context, which is an
    additional mechanism to add control dependencies.

Args:
        graph: The graph that this controller is managing.
        control_inputs: List of ops to use as control inputs in addition to the
        current control dependencies.  None to indicate that the dependencies
        should be cleared.


def row_limits(self):
    return self._row_splits[1:]

Returns the limit indices for rows in this row partition.

    These indices specify where the values for each row end.
    `partition.row_limits()` is equal to `partition.row_splits()[:-1]`.

Returns:
        A 1-D integer Tensor with shape `[self.nrows]`.
        The returned tensor is nonnegative, and is sorted in ascending order.
        `self.row_limits()[-1] == self.nvals()`.


def create_constructor_args(cls, proto_list: List[ir_swap.InterestRateSwap], config: InterestRateSwapConfig=None) -> Dict[str, Any]:
    swap_data = proto_utils.from_protos_v2(proto_list, config)
    res = {}
    for key in swap_data:
        tensor_repr = proto_utils.tensor_repr(swap_data[key])
        res[key] = tensor_repr
    return res

Creates a dictionary to initialize InterestRateSwap.

    The output dictionary is such that the instruments can be initialized
    as follows:
    ```
    initializer = create_constructor_args(proto_list, config)
    swaps = [InterestRateSwap(**data) for data in initializer.values()]
    ```

    The keys of the output dictionary are unique identifiers of the batched
    instruments. This is useful for identifying an existing graph that could be
    reused for the instruments without the need of rebuilding the graph.

Args:
        proto_list: A list of protos for which the initialization arguments are
        constructed.
        config: An instance of `InterestRateSwapConfig`.

Returns:
        A possibly nested dictionary such that each value provides initialization
        arguments for the InterestRateSwap.


def __init__(self, input_energy: energy.BernoulliEnergy, num_expectation_samples: int, initial_seed: Union[None, tf.Tensor]=None, name: Union[None, str]=None):
    super().__init__(input_energy, num_expectation_samples, initial_seed, name)
    self._logits_variable = tf.Variable(input_energy.logits, trainable=False)
    self._distribution = tfd.Bernoulli(logits=self._logits_variable, dtype=tf.int8)

Initializes a BernoulliEnergyInference.

Args:
        input_energy: The parameterized energy function which defines this
        distribution via the equations of an energy based model.  This class
        assumes that all parameters of `energy` are `tf.Variable`s and that
        they are all returned by `energy.variables`.
        num_expectation_samples: Number of samples to draw and use for estimating
        the expectation value.
        initial_seed: PRNG seed; see tfp.random.sanitize_seed for details. This
        seed will be used in the `sample` method.  If None, the seed is updated
        after every inference call.  Otherwise, the seed is fixed.
        name: Optional name for the model.


def weighting_function(max_num_bins: int, up: torch.Tensor, reg_scale: int) -> torch.Tensor:
    upper_bound1 = abs(up[0]) * abs(reg_scale)
    upper_bound2 = abs(up[0]) * abs(reg_scale) * 2
    step = (upper_bound1 + 1) ** (2 / (max_num_bins - 2))
    left_values = [-step ** i + 1 for i in range(max_num_bins // 2 - 1, 0, -1)]
    right_values = [step ** i - 1 for i in range(1, max_num_bins // 2)]
    values = [-upper_bound2] + left_values + [torch.zeros_like(up[0][None])] + right_values + [upper_bound2]
    values = [v if v.dim() > 0 else v.unsqueeze(0) for v in values]
    values = torch.cat(values, 0)
    return values

Generates the non-uniform Weighting Function W(n) for bounding box regression.

Args:
        max_num_bins (int): Max number of the discrete bins.
        up (Tensor): Controls upper bounds of the sequence,
                    where maximum offset is ±up * H / W.
        reg_scale (float): Controls the curvature of the Weighting Function.
                        Larger values result in flatter weights near the central axis W(max_num_bins/2)=0
                        and steeper weights at both ends.

Returns:
        Tensor: Sequence of Weighting Function.


def model_to_dot(model, show_shapes=False, show_dtype=False, show_layer_names=True, rankdir='TB', expand_nested=False, dpi=200, subgraph=False, show_layer_activations=False, show_trainable=False, **kwargs):
    from keras.src.ops.function import make_node_key
    if not model.built:
        raise ValueError('This model has not yet been built. Build the model first by calling `build()` or by calling the model on a batch of data.')
    from keras.src.models import functional
    from keras.src.models import sequential
    if not check_pydot():
        raise ImportError('You must install pydot (`pip install pydot`) for model_to_dot to work.')
    if subgraph:
        dot = pydot.Cluster(style='dashed', graph_name=model.name)
        dot.set('label', model.name)
        dot.set('labeljust', 'l')
    else:
        dot = pydot.Dot()
        dot.set('rankdir', rankdir)
        dot.set('concentrate', True)
        dot.set('dpi', dpi)
        dot.set('splines', 'ortho')
        dot.set_node_defaults(shape='record')
    if kwargs.pop('layer_range', None) is not None:
        raise ValueError('Argument `layer_range` is no longer supported.')
    if kwargs:
        raise ValueError(f'Unrecognized keyword arguments: {kwargs}')
    kwargs = {'show_layer_names': show_layer_names, 'show_layer_activations': show_layer_activations, 'show_dtype': show_dtype, 'show_shapes': show_shapes, 'show_trainable': show_trainable}
    if isinstance(model, sequential.Sequential):
        layers = model.layers
    elif not isinstance(model, functional.Functional):
        node = make_node(model, **kwargs)
        dot.add_node(node)
        return dot
    else:
        layers = model._operations
    for i, layer in enumerate(layers):
        if expand_nested and isinstance(layer, (functional.Functional, sequential.Sequential)):
            submodel = model_to_dot(layer, show_shapes, show_dtype, show_layer_names, rankdir, expand_nested, subgraph=True, show_layer_activations=show_layer_activations, show_trainable=show_trainable)
            dot.add_subgraph(submodel)
        else:
            node = make_node(layer, **kwargs)
            dot.add_node(node)
    if isinstance(model, sequential.Sequential):
        if not expand_nested:
            for i in range(len(layers) - 1):
                add_edge(dot, layers[i], layers[i + 1])
            return dot
        else:
            layers = model.layers[1:]
    for layer in layers:
        for inbound_index, inbound_node in enumerate(layer._inbound_nodes):
            if isinstance(model, functional.Functional) and make_node_key(layer, inbound_index) not in model._nodes:
                continue
            for input_index, input_tensor in enumerate(inbound_node.input_tensors):
                input_history = input_tensor._keras_history
                if input_history.operation is None:
                    continue
                input_node = input_history.operation._inbound_nodes[input_history.node_index]
                output_index = input_history.tensor_index
                source = input_node.operation
                destination = layer
                if not expand_nested:
                    add_edge(dot, source, layer)
                    continue
                while isinstance(source, (functional.Functional, sequential.Sequential)):
                    source, _, output_index = source.outputs[output_index]._keras_history
                while isinstance(destination, (functional.Functional, sequential.Sequential)):
                    if isinstance(destination, functional.Functional):
                        destination = destination.inputs[input_index]._keras_history.operation
                    else:
                        destination = destination.layers[0]
                add_edge(dot, source, destination)
    return dot

Convert a Keras model to dot format.

Args:
        model: A Keras model instance.
        show_shapes: whether to display shape information.
        show_dtype: whether to display layer dtypes.
        show_layer_names: whether to display layer names.
        rankdir: `rankdir` argument passed to PyDot,
            a string specifying the format of the plot: `"TB"`
            creates a vertical plot; `"LR"` creates a horizontal plot.
        expand_nested: whether to expand nested Functional models
            into clusters.
        dpi: Image resolution in dots per inch.
        subgraph: whether to return a `pydot.Cluster` instance.
        show_layer_activations: Display layer activations (only for layers that
            have an `activation` property).
        show_trainable: whether to display if a layer is trainable.

Returns:
        A `pydot.Dot` instance representing the Keras model or
        a `pydot.Cluster` instance representing nested model if
        `subgraph=True`.


def is_placeholder(x):
    try:
        if ops.executing_eagerly_outside_functions():
            return hasattr(x, '_is_backend_placeholder')
        from tensorflow.python.keras.utils import tf_utils
        if tf_utils.is_extension_type(x):
            flat_components = nest.flatten(x, expand_composites=True)
            return py_any((is_placeholder(c) for c in flat_components))
        else:
            return x.op.type == 'Placeholder'
    except AttributeError:
        return False

Returns whether `x` is a placeholder.

Args:
        x: A candidate placeholder.

Returns:
        Boolean.


def _read_side_inputs(self, tags_and_types):
    assert self.side_input_maps is None
    for i, (side_tag, view_class, view_options) in enumerate(tags_and_types):
        sources = []
        for si in filter(lambda o: o.tag == side_tag, self.spec.side_inputs):
            if not isinstance(si, operation_specs.WorkerSideInputSource):
                raise NotImplementedError('Unknown side input type: %r' % si)
            sources.append(si.source)
        si_counter = opcounters.SideInputReadCounter(self.counter_factory, self.state_sampler, declaring_step=self.name_context.step_name, input_index=i + 1)
        element_counter = opcounters.OperationCounters(self.counter_factory, self.name_context.step_name, view_options['coder'], i, suffix='side-input')
        iterator_fn = sideinputs.get_iterator_fn_for_sources(sources, read_counter=si_counter, element_counter=element_counter)
        yield apache_sideinputs.SideInputMap(view_class, view_options, sideinputs.EmulatedIterable(iterator_fn))

Generator reading side inputs in the order prescribed by tags_and_types.

Args:
        tags_and_types: List of tuples (tag, type). Each side input has a string
        tag that is specified in the worker instruction. The type is actually
        a boolean which is True for singleton input (read just first value)
        and False for collection input (read all values).

Yields:
        With each iteration it yields the result of reading an entire side source
        either in singleton or collection mode according to the tags_and_types
        argument.


def _partitioner(shape, dtype):
    if axis >= len(shape):
        raise ValueError(f'Cannot partition variable along axis {axis} when shape is only {shape}')
    dtype = dtypes.as_dtype(dtype)
    if dtype.base_dtype == dtypes.string:
        bytes_per_element = bytes_per_string_element
    else:
        bytes_per_element = dtype.size
    total_size_bytes = shape.num_elements() * bytes_per_element
    partitions = total_size_bytes / min_slice_size
    partitions_list = [1] * len(shape)
    partitions_list[axis] = max(1, min(shape.dims[axis].value, max_partitions, int(math.ceil(partitions))))
    return partitions_list

Partitioner that partitions list for a variable of given shape and type.

    Ex: Consider partitioning a variable of type float32 with
        shape=[1024, 1024].
        If `max_partitions` >= 16, this function would return
        [(1024 * 1024 * 4) / (256 * 1024), 1] = [16, 1].
        If `max_partitions` < 16, this function would return
        [`max_partitions`, 1].

Args:
        shape: Shape of the variable.
        dtype: Type of the variable.

Returns:
        List of partitions for each axis (currently only one axis can be
        partitioned).

Raises:
        ValueError: If axis to partition along does not exist for the variable.


def to_json(self, **kwargs) -> JSONValueType:

Returns a plain Python value as a representation for this object.

    A plain Python value are basic python types that can be serialized into
    JSON, e.g: ``bool``, ``int``, ``float``, ``str``, ``dict`` (with string
    keys), ``list``, ``tuple`` where the container types should have plain
    Python values as their values.

Args:
        **kwargs: Keyword arguments as flags to control JSON conversion.

Returns:
        A plain Python value.


def __init__(self, window_coder):
    super().__init__()
    if window_coder is None:
        raise ValueError('window_coder should not be None')
    self._window_coder = window_coder

Create a new WindowFn with compatible coder.
    To be applied to PCollections with windows that are compatible with the
    given coder.

Args:
        window_coder: coders.Coder object to be used on windows.


def _make_1d_state(state_size, seed):
    if isinstance(seed, int):
        ls = []
        for _ in range(state_size):
            ls.append(seed & SEED_BIT_MASK)
            seed >>= SEED_TYPE_BITS
        seed = ls
    seed = nest.map_structure(_uint_to_int, seed)
    seed = math_ops.cast(seed, STATE_TYPE)
    seed = array_ops.reshape(seed, [-1])
    seed = seed[0:state_size]
    seed_size = seed.shape[0]
    if seed_size is None:
        seed_size = array_ops.shape(seed)[0]
    padding_size = math_ops.maximum(state_size - seed_size, 0)
    padding = array_ops.zeros([padding_size], seed.dtype)
    seed = array_ops.concat([padding, seed], axis=0)
    seed.set_shape([state_size])
    return seed

Makes a 1-D RNG state.

Args:
        state_size: an integer.
        seed: an integer or 1-D tensor.

Returns:
        a 1-D tensor of shape [state_size] and dtype STATE_TYPE.


def _reshape_tensors(tensors, shape):
    reshaped = []
    for t in tensors:
        with ops.colocate_with(t):
            reshaped.append(array_ops.reshape(t, shape))
    return reshaped

Reshape tensors flattened by _flatten_tensors.

Args:
        tensors: list of `tf.Tensor` of identical length 1D tensors.
        shape: list of integers describing the desired shape.  Product of
        the elements must equal the length of each tensor.

Returns:
        list of `tf.Tensor` which are the reshaped inputs.


A function that is executed upon a test being skipped.

    Implementation is optional.

Args:
        record: records.TestResultRecord, a copy of the test record for
        this test, containing all information of the test execution
        including exception objects.


def assert_is_fully_defined(self):
    if not self.is_fully_defined():
        raise ValueError('Shape %s is not fully defined' % self)

Raises an exception if `self` is not fully defined in every dimension.

Raises:
        ValueError: If `self` does not have a known value for every dimension.


def f():
    return constant_op.constant(1)

First sentence.

    Second sentence.

Returns:
        Something.


def get_dict_table_schema(schema):
    if isinstance(schema, (dict, value_provider.ValueProvider)) or callable(schema) or schema is None:
        return schema
    elif isinstance(schema, str):
        table_schema = get_table_schema_from_string(schema)
        return table_schema_to_dict(table_schema)
    elif isinstance(schema, bigquery.TableSchema):
        return table_schema_to_dict(schema)
    else:
        raise TypeError('Unexpected schema argument: %s.' % schema)

Transform the table schema into a dictionary instance.

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

Returns:
        Dict[str, Any]: The schema to be used if the BigQuery table to write has
        to be created but in the dictionary format.


def _validate_testbed_configs(testbed_configs):
    seen_names = set()
    for config in testbed_configs:
        name = config[keys.Config.key_testbed_name.value]
        _validate_testbed_name(name)
        if name in seen_names:
            raise MoblyConfigError('Duplicate testbed name %s found.' % name)
        seen_names.add(name)

Validates the testbed configurations.

Args:
        testbed_configs: A list of testbed configuration dicts.

Raises:
        MoblyConfigError: Some parts of the configuration is invalid.


def m_to_inches(value):
    if value is None:
        return None
    return value / 39.37

Converts distance in meters to inches

Args:
        value: floating point representing the distance in meters
    Returns: distance in inches


def forward(self, hidden_states: torch.Tensor):
    gate_score = self.act(self.w_0(hidden_states))
    hidden_states = self.w_1(hidden_states)
    hidden_states = gate_score * hidden_states
    return hidden_states

Transform an input tensor from one feature space to another via a nonlinear operation

Args:
        hidden_states (`torch.Tensor` of shape `(batch, seq_len, dim_in)`)


def score_retrieval(self, query_embeddings: Union['torch.Tensor', List['torch.Tensor']], passage_embeddings: Union['torch.Tensor', List['torch.Tensor']], batch_size: int=128, output_dtype: Optional['torch.dtype']=None, output_device: Union['torch.device', str]='cpu') -> 'torch.Tensor':
    if len(query_embeddings) == 0:
        raise ValueError('No queries provided')
    if len(passage_embeddings) == 0:
        raise ValueError('No passages provided')
    if query_embeddings[0].device != passage_embeddings[0].device:
        raise ValueError('Queries and passages must be on the same device')
    if query_embeddings[0].dtype != passage_embeddings[0].dtype:
        raise ValueError('Queries and passages must have the same dtype')
    if output_dtype is None:
        output_dtype = query_embeddings[0].dtype
    scores: List[torch.Tensor] = []
    for i in range(0, len(query_embeddings), batch_size):
        batch_scores: List[torch.Tensor] = []
        batch_queries = torch.nn.utils.rnn.pad_sequence(query_embeddings[i:i + batch_size], batch_first=True, padding_value=0)
        for j in range(0, len(passage_embeddings), batch_size):
            batch_passages = torch.nn.utils.rnn.pad_sequence(passage_embeddings[j:j + batch_size], batch_first=True, padding_value=0)
            batch_scores.append(torch.einsum('bnd,csd->bcns', batch_queries, batch_passages).max(dim=3)[0].sum(dim=2))
        scores.append(torch.cat(batch_scores, dim=1).to(output_dtype).to(output_device))
    return torch.cat(scores, dim=0)

Compute the late-interaction/MaxSim score (ColBERT-like) for the given multi-vector
    query embeddings (`qs`) and passage embeddings (`ps`). For ColQwen2, a passage is the
    image of a document page.

    Because the embedding tensors are multi-vector and can thus have different shapes, they
    should be fed as:
    (1) a list of tensors, where the i-th tensor is of shape (sequence_length_i, embedding_dim)
    (2) a single tensor of shape (n_passages, max_sequence_length, embedding_dim) -> usually
        obtained by padding the list of tensors.

Args:
        query_embeddings (`Union[torch.Tensor, List[torch.Tensor]`): Query embeddings.
        passage_embeddings (`Union[torch.Tensor, List[torch.Tensor]`): Passage embeddings.
        batch_size (`int`, *optional*, defaults to 128): Batch size for computing scores.
        output_dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): The dtype of the output tensor.
            If `None`, the dtype of the input embeddings is used.
        output_device (`torch.device` or `str`, *optional*, defaults to "cpu"): The device of the output tensor.

Returns:
        `torch.Tensor`: A tensor of shape `(n_queries, n_passages)` containing the scores. The score
        tensor is saved on the "cpu" device.


def call(self, input_ids: TFModelInputType | None=None, attention_mask: tf.Tensor | None=None, token_type_ids: tf.Tensor | None=None, inputs_embeds: tf.Tensor | None=None, output_attentions: bool | None=None, output_hidden_states: bool | None=None, return_dict: bool | None=None, training: bool=False) -> TFDPRQuestionEncoderOutput | Tuple[tf.Tensor, ...]:
    if input_ids is not None and inputs_embeds is not None:
        raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time')
    elif input_ids is not None:
        input_shape = shape_list(input_ids)
    elif inputs_embeds is not None:
        input_shape = shape_list(inputs_embeds)[:-1]
    else:
        raise ValueError('You have to specify either input_ids or inputs_embeds')
    if attention_mask is None:
        attention_mask = tf.ones(input_shape, dtype=tf.dtypes.int32) if input_ids is None else input_ids != self.config.pad_token_id
    if token_type_ids is None:
        token_type_ids = tf.zeros(input_shape, dtype=tf.dtypes.int32)
    outputs = self.question_encoder(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, inputs_embeds=inputs_embeds, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, training=training)
    if not return_dict:
        return outputs[1:]
    return TFDPRQuestionEncoderOutput(pooler_output=outputs.pooler_output, hidden_states=outputs.hidden_states, attentions=outputs.attentions)

Return:

Example:
    ```python
    >>> from transformers import TFDPRQuestionEncoder, DPRQuestionEncoderTokenizer

    >>> tokenizer = DPRQuestionEncoderTokenizer.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
    >>> model = TFDPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base", from_pt=True)
    >>> input_ids = tokenizer("Hello, is my dog cute ?", return_tensors="tf")["input_ids"]
    >>> embeddings = model(input_ids).pooler_output
    ```


def scale_translation(self, trans_scale_factor: float) -> Rigid:
    return self.apply_trans_fn(lambda t: t * trans_scale_factor)

Scales the translation by a constant factor.

Args:
        trans_scale_factor:
            The constant factor

Returns:
        A transformation object with a scaled translation.


def _tokenize(self, text, **kwargs):
    return self.sp_model.encode(text, out_type=str)

Args:
    text: TextInput
    Returns a tokenized string. The Gemma tokenizer never adds a prefix space.


def bessel_i1(x, name=None):
    with ops.name_scope(name, 'bessel_i1', [x]):
        return gen_special_math_ops.bessel_i1(x)

Computes the Bessel i1 function of `x` element-wise.

    Modified Bessel function of order 1.

    It is preferable to use the numerically stabler function `i1e(x)` instead.

    >>> tf.math.special.bessel_i1([-1., -0.5, 0.5, 1.]).numpy()
    array([-0.5651591 , -0.25789431,  0.25789431,  0.5651591 ], dtype=float32)

Args:
        x: A `Tensor` or `SparseTensor`. Must be one of the following types: `half`,
        `float32`, `float64`.
        name: A name for the operation (optional).

Returns:
        A `Tensor` or `SparseTensor`, respectively. Has the same type as `x`.

    @compatibility(scipy)
    Equivalent to scipy.special.i1
    @end_compatibility


def take_bug_report(self, test_name=None, begin_time=None, timeout=300, destination=None):
    prefix = DEFAULT_BUG_REPORT_NAME
    if test_name:
        prefix = '%s,%s' % (DEFAULT_BUG_REPORT_NAME, test_name)
    if begin_time is None:
        begin_time = mobly_logger.get_log_file_timestamp()
    new_br = True
    try:
        stdout = self.adb.shell('bugreportz -v').decode('utf-8')
        if 'not found' in stdout:
            new_br = False
    except adb.AdbError:
        new_br = False
    if destination is None:
        destination = os.path.join(self.log_path, 'BugReports')
    br_path = utils.abs_path(destination)
    utils.create_dir(br_path)
    filename = self.generate_filename(prefix, str(begin_time), 'txt')
    if new_br:
        filename = filename.replace('.txt', '.zip')
    full_out_path = os.path.join(br_path, filename)
    self.wait_for_boot_completion()
    self.log.debug('Start taking bugreport.')
    if new_br:
        out = self.adb.shell('bugreportz', timeout=timeout).decode('utf-8')
        if not out.startswith('OK'):
            raise DeviceError(self, 'Failed to take bugreport: %s' % out)
        br_out_path = out.split(':')[1].strip()
        self.adb.pull([br_out_path, full_out_path])
        self.adb.shell(['rm', br_out_path])
    else:
        self.adb.bugreport(' > "%s"' % full_out_path, shell=True, timeout=timeout)
    self.log.debug('Bugreport taken at %s.', full_out_path)
    return full_out_path

Takes a bug report on the device and stores it in a file.

Args:
        test_name: Name of the test method that triggered this bug report.
        begin_time: Timestamp of when the test started. If not set, then
        this will default to the current time.
        timeout: float, the number of seconds to wait for bugreport to
        complete, default is 5min.
        destination: string, path to the directory where the bugreport
        should be saved.

Returns:
        A string that is the absolute path to the bug report on the host.


def manual_variable_initialization(value):
    global _MANUAL_VAR_INIT
    _MANUAL_VAR_INIT = value

Sets the manual variable initialization flag.

    This boolean flag determines whether
    variables should be initialized
    as they are instantiated (default), or if
    the user should handle the initialization
    (e.g. via `tf.compat.v1.initialize_all_variables()`).

Args:
        value: Python boolean.


def input_node_from_schema(schema: Schema, same_sampling_as: Optional[EventSetNode]=None, name: Optional[str]=None) -> EventSetNode:
    return input_node(features=schema.features, indexes=schema.indexes, is_unix_timestamp=schema.is_unix_timestamp, same_sampling_as=same_sampling_as, name=name)

Creates an input [`EventSetNode`][temporian.EventSetNode] from a schema.

    Usage example:

        ```python
        >>> # Create two nodes with the same schema.
        >>> a = tp.input_node(features=[("f1", tp.float64), ("f2", tp.str_)])
        >>> b = tp.input_node_from_schema(a.schema)

        ```

Args:
        schema: Schema of the node.
        same_sampling_as: If set, the created EventSetNode is guaranteed to have the
            same sampling as same_sampling_as`. In this case, `indexes` and
            `is_unix_timestamp` should not be provided. Some operators require
            for input EventSetNodes to have the same sampling.
        name: Name for the EventSetNode.

Returns:
        EventSetNode with the given specifications.


def row_splits(self):
    return self._row_splits

Returns the row-split indices for this row partition.

    `row_splits` specifies where the values for each row begin and end.
    In particular, the values for row `i` are stored in the slice
    `values[row_splits[i]:row_splits[i+1]]`.

Returns:
        A 1-D integer `Tensor` with shape `[self.nrows+1]`.
        The returned tensor is non-empty, and is sorted in ascending order.
        `self.row_splits()[0] == 0`.
        `self.row_splits()[-1] == self.nvals()`.


def _has_attr(self, node, obj, attr):
    if isinstance(obj, abstract.AMBIGUOUS_OR_EMPTY):
        return (node, None)
    if not isinstance(attr, abstract.PythonConstant) or not isinstance(attr.pyval, str):
        return (node, None)
    node, ret = self.ctx.attribute_handler.get_attribute(node, obj, attr.pyval)
    return (node, ret is not None)

Check if the object has attribute attr.

Args:
        node: The given node.
        obj: A BaseValue, generally the left hand side of a hasattr() call.
        attr: A BaseValue, generally the right hand side of a hasattr() call.

Returns:
        (node, result) where result = True if the object has attribute attr, False
        if it does not, and None if it is ambiguous.