def options_price_from_samples(strikes: types.RealTensor, expiries: types.RealTensor, maturities: types.RealTensor, is_call_options: types.BoolTensor, sample_discount_curve_paths_fn: Callable[..., Tuple[types.RealTensor, types.RealTensor]], num_samples: types.IntTensor, time_step: types.RealTensor, dtype: tf.DType=None, name: str=None) -> types.RealTensor:
    name = name or 'options_price_from_samples'
    with tf.name_scope(name):
        sim_times, _ = tf.unique(tf.reshape(expiries, shape=[-1]))
        longest_expiry = tf.reduce_max(sim_times)
        sim_times, _ = tf.unique(tf.concat([sim_times, tf.range(time_step, longest_expiry, time_step)], axis=0))
        sim_times = tf.sort(sim_times, name='sort_sim_times')
        tau = maturities - expiries
        curve_times_builder, _ = tf.unique(tf.reshape(tau, shape=[-1]))
        curve_times = tf.sort(curve_times_builder, name='sort_curve_times')
        p_t_tau, r_t = sample_discount_curve_paths_fn(times=sim_times, curve_times=curve_times, num_samples=num_samples)
        dim = p_t_tau.shape[-1]
        dt_builder = tf.concat(axis=0, values=[tf.convert_to_tensor([0.0], dtype=dtype), sim_times[1:] - sim_times[:-1]])
        dt = tf.expand_dims(tf.expand_dims(dt_builder, axis=-1), axis=0)
        discount_factors_builder = tf.math.exp(-r_t * dt)
        discount_factors_builder = tf.transpose(utils.cumprod_using_matvec(tf.transpose(discount_factors_builder, [0, 2, 1])), [0, 2, 1])
        discount_factors_builder = tf.expand_dims(discount_factors_builder, axis=1)
        discount_factors_simulated = tf.repeat(discount_factors_builder, p_t_tau.shape.as_list()[1], axis=1)
        sim_time_index = tf.searchsorted(sim_times, tf.reshape(expiries, [-1]))
        curve_time_index = tf.searchsorted(curve_times, tf.reshape(tau, [-1]))
        curve_time_index, sim_time_index = tff_utils.broadcast_tensors(curve_time_index, sim_time_index)
        gather_index = _prepare_indices(tf.range(0, num_samples), curve_time_index, sim_time_index, tf.range(0, dim))
        payoff_discount_factors_builder = tf.gather_nd(discount_factors_simulated, gather_index)
        payoff_discount_factors = tf.reshape(payoff_discount_factors_builder, [num_samples] + strikes.shape + [dim])
        payoff_bond_price_builder = tf.gather_nd(p_t_tau, gather_index)
        payoff_bond_price = tf.reshape(payoff_bond_price_builder, [num_samples] + strikes.shape + [dim])
        is_call_options = tf.reshape(tf.broadcast_to(is_call_options, strikes.shape), [1] + strikes.shape + [1])
        strikes = tf.reshape(strikes, [1] + strikes.shape + [1])
        payoff = tf.where(is_call_options, tf.math.maximum(payoff_bond_price - strikes, 0.0), tf.math.maximum(strikes - payoff_bond_price, 0.0))
        option_value = tf.math.reduce_mean(payoff_discount_factors * payoff, axis=0)
        return option_value

Computes the zero coupon bond options price from simulated discount curves.

Args:
  strikes: A real `Tensor` of any shape and dtype. The strike price of the
    options. The shape of this input determines the number (and shape) of the
    options to be priced and the output.
  expiries: A real `Tensor` of the same dtype and compatible shape as
    `strikes`.  The time to expiry of each bond option.
  maturities: A real `Tensor` of the same dtype and compatible shape as
    `strikes`.  The time to maturity of the underlying zero coupon bonds.
  is_call_options: A boolean `Tensor` of a shape compatible with `strikes`.
    Indicates whether the option is a call (if True) or a put (if False).
  sample_discount_curve_paths_fn: Callable which takes the following args:

    1) times: Rank 1 `Tensor` of positive real values, specifying the times at
      which the path points are to be evaluated.
    2) curve_times: Rank 1 `Tensor` of positive real values, specifying the
      maturities at which the discount curve is to be computed at each
      simulation time.
    3) num_samples: Positive scalar integer specifying the number of paths to
      draw.

    and returns two `Tensor`s, the first being a Rank-4 tensor of shape
    `[num_samples, m, k, dim]` containing the simulated zero coupon bond
    curves, and the second being a `Tensor` of shape `[num_samples, k, dim]`
    containing the simulated short rate paths. Here, `m` is the size of
    `curve_times`, `k` is the size of `times`, and `dim` is the dimensionality
    of the paths.

  num_samples: Positive scalar `int32` `Tensor`. The number of simulation
    paths during Monte-Carlo valuation.
  time_step: Scalar real `Tensor`. Maximal distance between time grid points
    in Euler scheme. Relevant when Euler scheme is used for simulation.
  dtype: The default dtype to use when converting values to `Tensor`s.
    Default value: `None` which means that default dtypes inferred by
      TensorFlow are used.
  name: Python string. The name to give to the ops created by this function.
    Default value: `None` which maps to the default name
    `options_price_from_samples`.

Returns:
  A `Tensor` of real dtype and shape `strikes.shape + [dim]` containing the
  computed option prices.


def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)
    q_embed = q * cos + rotate_half(q) * sin
    k_embed = k * cos + rotate_half(k) * sin
    return (q_embed, k_embed)

Applies Rotary Position Embedding to the query and key tensors.

Args:
    q (`torch.Tensor`): The query tensor.
    k (`torch.Tensor`): The key tensor.
    cos (`torch.Tensor`): The cosine part of the rotary embedding.
    sin (`torch.Tensor`): The sine part of the rotary embedding.
    position_ids (`torch.Tensor`, *optional*):
        Deprecated and unused.
    unsqueeze_dim (`int`, *optional*, defaults to 1):
        The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
        sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
        that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
        k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
        cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
        the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
Returns:
    `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.


class LayoutLMv2FastImageProcessorKwargs(DefaultFastImageProcessorKwargs):
    apply_ocr: Optional[bool]
    ocr_lang: Optional[str]
    tesseract_config: Optional[str]

Args:
    apply_ocr (`bool`, *optional*, defaults to `True`):
        Whether to apply the Tesseract OCR engine to get words + normalized bounding boxes. Can be overridden by
        the `apply_ocr` parameter in the `preprocess` method.
    ocr_lang (`str`, *optional*):
        The language, specified by its ISO code, to be used by the Tesseract OCR engine. By default, English is
        used. Can be overridden by the `ocr_lang` parameter in the `preprocess` method.
    tesseract_config (`str`, *optional*):
        Any additional custom configuration flags that are forwarded to the `config` parameter when calling
        Tesseract. For example: '--psm 6'. Can be overridden by the `tesseract_config` parameter in the
        `preprocess` method.


def count_function(function: _evaluation.CountFunction, operand_result: Optional[_sql_data_types.Select], params_result: Collection[_sql_data_types.StandardSqlExpression]) -> _sql_data_types.Select:
    del function, params_result
    if operand_result is None:
        raise ValueError('count() cannot be called without an operand.')
    if operand_result.from_part is None:
        return _sql_data_types.Select(select_part=_sql_data_types.CountCall((_sql_data_types.RawExpression(operand_result.sql_alias, _sql_data_type=operand_result.sql_data_type),)), from_part=str(operand_result.to_subquery()), where_part=operand_result.where_part, sql_dialect=_sql_data_types.SqlDialect.SPARK)
    else:
        return dataclasses.replace(operand_result, select_part=_sql_data_types.CountCall((operand_result.select_part,)))

Returns an integer representing the number of elements in a collection.

By default, `_CountFunction` will return 0.

Args:
  function: The FHIRPath AST `HasValueFunction` node
  operand_result: The expression which is being evaluated
  params_result: The parameter passed in to function

Returns:
  A compiled Spark SQL expression.

Raises:
  ValueError: When the function is called without an operand


def set_help_intro(self, help_intro):
    self._command_handler_registry.set_help_intro(help_intro=help_intro)

Set an introductory message to the help output of the command registry.

Args:
  help_intro: (RichTextLines) Rich text lines appended to the beginning of
    the output of the command "help", as introductory information.


def set_state_tree(self, state_tree):
    for k, v in state_tree.items():
        path_value_dict = self._flatten_nested_dict(v)
        if k == 'trainable_variables':
            self._assign_variable_values(self.trainable_variables, path_value_dict)
        elif k == 'non_trainable_variables':
            self._assign_variable_values(self.non_trainable_variables, path_value_dict)
        elif k == 'optimizer_variables':
            self._assign_variable_values(self.optimizer.variables, path_value_dict)
        elif k == 'metrics_variables':
            self._assign_variable_values(self.metrics_variables, path_value_dict)
        else:
            raise ValueError(f'Unknown variable name: {k}')

Assigns values to variables of the model.

This method takes a dictionary of nested variable values, which
represents the state tree of the model, and assigns them to the
corresponding variables of the model. The dictionary keys represent the
variable names (e.g., `'trainable_variables'`, `'optimizer_variables'`),
and the values are nested dictionaries containing the variable
paths and their corresponding values.

Args:
    state_tree: A dictionary representing the state tree of the model.
        The keys are the variable names, and the values are nested
        dictionaries representing the variable paths and their values.


def __init__(self, counter_factory, state_sampler, declaring_step, input_index):
    super().__init__(counter_factory, state_sampler)
    self.declaring_step = declaring_step
    self.input_index = input_index
    self.update_current_step()

Create a side input read counter.

Args:
  counter_factory: A counters.CounterFactory to create byte counters.
  state_sampler: A statesampler.StateSampler to transition into read states.
  declaring_step: A string with the step name of the step that directly
    receives the side input initially.
  input_index: The index of the side input in the list of inputs of the
    declaring step.

The side input is uniquely identified by (declaring_step, input_index);
where declaring_step is the step that receives the PCollectionView as a
side input, and input_index is the index of the PCollectionView within
the list of inputs.


class Qwen2AudioEncoderConfig(PretrainedConfig):
    model_type = 'qwen2_audio_encoder'

    def __init__(self, num_mel_bins=128, encoder_layers=32, encoder_attention_heads=20, encoder_ffn_dim=5120, encoder_layerdrop=0.0, d_model=1280, dropout=0.0, attention_dropout=0.0, activation_function='gelu', activation_dropout=0.0, scale_embedding=False, initializer_range=0.02, max_source_positions=1500, **kwargs):
        super().__init__(**kwargs)
        self.num_mel_bins = num_mel_bins
        self.d_model = d_model
        self.encoder_layers = encoder_layers
        self.encoder_attention_heads = encoder_attention_heads
        self.encoder_ffn_dim = encoder_ffn_dim
        self.dropout = dropout
        self.attention_dropout = attention_dropout
        self.activation_function = activation_function
        self.activation_dropout = activation_dropout
        self.encoder_layerdrop = encoder_layerdrop
        self.num_hidden_layers = encoder_layers
        self.initializer_range = initializer_range
        self.scale_embedding = scale_embedding
        self.max_source_positions = max_source_positions

This is the configuration class to store the configuration of a [`Qwen2AudioEncoder`]. It is used to instantiate a
Qwen2-Audio audio encoder according to the specified arguments, defining the model architecture. Instantiating a
configuration with the defaults will yield a similar configuration to that of the audio encoder of the Qwen2-Audio
architecture.

e.g. [Qwen/Qwen2-Audio-7B](https://huggingface.co/Qwen/Qwen2-Audio-7B)

Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
documentation from [`PretrainedConfig`] for more information.

Args:
    num_mel_bins (`int`, *optional*, defaults to 128):
        Number of mel features used per input features. Should correspond to the value used in the
        `Qwen2AudioProcessor` class.
    encoder_layers (`int`, *optional*, defaults to 32):
        Number of encoder layers.
    encoder_attention_heads (`int`, *optional*, defaults to 20):
        Number of attention heads for each attention layer in the Transformer encoder.
    encoder_ffn_dim (`int`, *optional*, defaults to 5120):
        Dimensionality of the "intermediate" (often named feed-forward) layer in encoder.
    encoder_layerdrop (`float`, *optional*, defaults to 0.0):
        The LayerDrop probability for the encoder. See the [LayerDrop paper](see https://huggingface.co/papers/1909.11556)
        for more details.
    d_model (`int`, *optional*, defaults to 1280):
        Dimensionality of the layers.
    dropout (`float`, *optional*, defaults to 0.0):
        The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
    attention_dropout (`float`, *optional*, defaults to 0.0):
        The dropout ratio for the attention probabilities.
    activation_function (`str`, *optional*, defaults to `"gelu"`):
        The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
        `"relu"`, `"silu"` and `"gelu_new"` are supported.
    activation_dropout (`float`, *optional*, defaults to 0.0):
        The dropout ratio for activations inside the fully connected layer.
    scale_embedding (`bool`, *optional*, defaults to `False`):
        Scale embeddings by diving by sqrt(d_model).
    initializer_range (`float`, *optional*, defaults to 0.02):
        The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
    max_source_positions (`int`, *optional*, defaults to 1500):
        The maximum sequence length of log-mel filter-bank features that this model might ever be used with.

Example:

```python
>>> from transformers import Qwen2AudioEncoderConfig, Qwen2AudioEncoder

>>> # Initializing a Qwen2AudioEncoderConfig
>>> configuration = Qwen2AudioEncoderConfig()

>>> # Initializing a Qwen2AudioEncoder (with random weights)
>>> model = Qwen2AudioEncoder(configuration)

>>> # Accessing the model configuration
>>> configuration = model.config
```


def returns_scalar(return_type: Optional[FhirPathDataType]) -> bool:
    return not return_type or return_type.cardinality == Cardinality.SCALAR

Indicates if the return type evaluates to a scalar.

Args:
  return_type: The data type to describe.

Returns:
  True if `return_type` represents an element with cardinality less than or
  equal to one whose parents are all also scalars.
  False otherwise. For example, the path Patient.name.use does not return a
  scalar, despite 'use' being a scalar, because it is a child of the
  collection, 'name.'


def enable_mixed_precision_graph_rewrite_v1(opt, loss_scale='dynamic'):
    if mixed_precision_global_state.is_using_mixed_precision_policy():
        raise ValueError('The mixed precision graph rewrite cannot be enabled, because the global Keras dtype Policy has been set to a mixed precision policy. At most, one of the following can be called:\n\n  1. tf.keras.mixed_precision.set_global_policy() with a mixed precision policy (You called this first)\n\n  2. tf.train.experimental.enable_mixed_precision_graph_rewrite() (You called this second)\nYou called both functions, which is an error, because both functions enable you to use mixed precision. If in doubt which function to use, use the first, as it supports Eager execution and is more customizable.')
    if mixed_precision_global_state.non_mixed_precision_session_created():
        tf_logging.warn('You already have existing Sessions that do not use mixed precision. enable_mixed_precision_graph_rewrite() will not affect these Sessions.')
    opt = _wrap_optimizer(opt, loss_scale)
    config.set_optimizer_experimental_options({'auto_mixed_precision': True})
    mixed_precision_global_state.set_mixed_precision_graph_rewrite_enabled(True)
    return opt

Enable mixed precision via a graph rewrite.

Mixed precision is the use of both float32 and float16 data types when
training a model to improve performance. This is achieved via a graph rewrite
operation and a loss-scale optimizer.

Performing arithmetic operations in float16 takes advantage of specialized
processing units, such as NVIDIA Tensor Cores, for much higher arithmetic
throughput. However, due to the smaller representable range, performing the
entire training with float16 can result in gradient underflow, that is, small
gradient values becoming zeroes. Instead, performing only select arithmetic
operations in float16 results in higher throughput and decreased training
time when using compatible hardware accelerators while also reducing memory
usage, typically without sacrificing model accuracy.

Note: While the mixed precision rewrite changes the datatype of various
layers throughout the model, the same accuracy reached in float32 is
expected. If a `NaN` gradient occurs with dynamic loss scaling, the model
update for that batch is skipped. In this case, the global step count is not
incremented, and the `LossScaleOptimizer` attempts to decrease the loss
scaling value to avoid `NaN` values in subsequent iterations. This approach
has been shown to achieve the same accuracy as float32 and, in most cases,
better training throughput.

Example:

```python
model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(64, activation='softmax'),
])

opt = tf.keras.optimizers.SGD()
opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(opt)
model.compile(loss="mse", optimizer=opt)

x_train = np.random.random((1024, 64))
y_train = np.random.random((1024, 64))
model.fit(x_train, y_train)
```

Calling `enable_mixed_precision_graph_rewrite(opt)` enables the graph rewrite
operation before computing gradients. The function additionally returns an
`Optimizer` (`opt`) wrapped with a `LossScaleOptimizer`. This prevents
underflow in the float16 tensors during the backward pass. An optimizer of
type `tf.train.Optimizer` or `tf.keras.optimizers.Optimizer` must be passed
to this function, which will then be wrapped to use loss scaling.

The graph rewrite operation changes the `dtype` of certain operations in the
graph from float32 to float16. There are several categories of operations
that are either included or excluded by this rewrite operation. The following
categories of Ops are defined inside corresponding functions under the class
`AutoMixedPrecisionLists` in

auto_mixed_precision_lists.h:

* `ClearList`: Ops that do not have numerically significant adverse effects.
E.g. `ArgMax` and `Floor`.
* `AllowList`: Ops that are considered numerically safe for execution in
float16, and thus are always converted. E.g. `Conv2D`.
* `DenyList`: Ops that are numerically unsafe to execute in float16 and
can negatively affect downstream nodes. E.g. `Softmax`.
* `GrayList`: Ops that are considered numerically safe for execution in
float16 unless downstream from a DenyList Op. E.g. `Add` and `AvgPool`.

When this function is used, gradients should only be computed and applied
with the returned optimizer, either by calling `opt.minimize()` or
`opt.compute_gradients()` followed by `opt.apply_gradients()`.
Gradients should not be computed with `tf.gradients` or `tf.GradientTape`.
This is because the returned optimizer will apply loss scaling, and
`tf.gradients` or `tf.GradientTape` will not. If you do directly use
`tf.gradients` or `tf.GradientTape`, your model may not converge due to
float16 underflow problems.

When eager execution is enabled, the mixed precision graph rewrite is only
enabled within `tf.function`s, as outside `tf.function`s, there is no graph.

For NVIDIA GPUs with Tensor cores, as a general performance guide, dimensions
(such as batch size, input size, output size, and channel counts)
should be powers of two if under 256, or  otherwise divisible by 8 if above
256. For more information, check out the
[NVIDIA Deep Learning Performance Guide](
https://docs.nvidia.com/deeplearning/sdk/dl-performance-guide/index.html).

Currently, mixed precision is only enabled on NVIDIA Tensor Core GPUs with
Compute Capability 7.0 and above (Volta, Turing, or newer architectures). The
parts of the graph on CPUs and TPUs are untouched by the graph rewrite.

Raises:
  `ValueError`, if the `tf.keras.mixed_precision` API is also used by calling
  `tf.keras.mixed_precision.set_global_policy`. Only one mixed precision
  API can be used.

Args:
  opt: An instance of a `tf.keras.optimizers.Optimizer` or a
    `tf.train.Optimizer`.
  loss_scale: Either an int/float, the string `"dynamic"`, or an instance of
    a `tf.mixed_precision.experimental.LossScale`. The loss scale to use. It
    is recommended to keep this as its default value of `"dynamic"`, which
    will adjust the scaling automatically to prevent `Inf` or `NaN` values.

Returns:
  A version of `opt` that will use loss scaling to prevent underflow.


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

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

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

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

#### Examples

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

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


def reduce_to(self, reduce_op, value, destinations, options=None):
    with monitoring.MonitoredTimer(distributed_api_time_counter.get_cell(self.__class__.__name__, 'Reduce_to_eagerly')) if not ops.inside_function() else contextlib.nullcontext():
        if options is None:
            options = collective_util.Options()
        _require_cross_replica_or_default_context_extended(self)
        assert not isinstance(destinations, (list, tuple))
        assert not isinstance(reduce_op, variable_scope.VariableAggregation)
        if isinstance(reduce_op, six.string_types):
            reduce_op = reduce_util.ReduceOp(reduce_op.upper())
        assert reduce_op == reduce_util.ReduceOp.SUM or reduce_op == reduce_util.ReduceOp.MEAN
        return self._reduce_to(reduce_op, value, destinations, options)

Combine (via e.g. sum or mean) values across replicas.

`reduce_to` aggregates `tf.distribute.DistributedValues` and distributed
variables. It supports both dense values and `tf.IndexedSlices`.

This API currently can only be called in cross-replica context. Other
variants to reduce values across replicas are:
* `tf.distribute.StrategyExtended.batch_reduce_to`: the batch version of
  this API.
* `tf.distribute.ReplicaContext.all_reduce`: the counterpart of this API
  in replica context. It supports both batched and non-batched all-reduce.
* `tf.distribute.Strategy.reduce`: a more convenient method to reduce
  to the host in cross-replica context.

`destinations` specifies where to reduce the value to, e.g. "GPU:0". You can
also pass in a `Tensor`, and the destinations will be the device of that
tensor. For all-reduce, pass the same to `value` and `destinations`.

It can be used in `tf.distribute.ReplicaContext.merge_call` to write code
that works for all `tf.distribute.Strategy`.

@tf.function
def step_fn(var):

  def merge_fn(strategy, value, var):
    # All-reduce the value. Note that `value` here is a
    # `tf.distribute.DistributedValues`.
    reduced = strategy.extended.reduce_to(tf.distribute.ReduceOp.SUM,
        value, destinations=var)
    strategy.extended.update(var, lambda var, value: var.assign(value),
        args=(reduced,))

  value = tf.identity(1.)
  tf.distribute.get_replica_context().merge_call(merge_fn,
    args=(value, var))

def run(strategy):
  with strategy.scope():
    v = tf.Variable(0.)
    strategy.run(step_fn, args=(v,))
    return v

run(tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"]))
MirroredVariable:{
  0: <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=2.0>,
  1: <tf.Variable 'Variable/replica_1:0' shape=() dtype=float32, numpy=2.0>
}
run(tf.distribute.experimental.CentralStorageStrategy(
    compute_devices=["GPU:0", "GPU:1"], parameter_device="CPU:0"))
<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=2.0>
run(tf.distribute.OneDeviceStrategy("GPU:0"))
<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>

Args:
  reduce_op: a `tf.distribute.ReduceOp` value specifying how values should
    be combined. Allows using string representation of the enum such as
    "SUM", "MEAN".
  value: a `tf.distribute.DistributedValues`, or a `tf.Tensor` like object.
  destinations: a `tf.distribute.DistributedValues`, a `tf.Variable`, a
    `tf.Tensor` alike object, or a device string. It specifies the devices
    to reduce to. To perform an all-reduce, pass the same to `value` and
    `destinations`. Note that if it's a `tf.Variable`, the value is reduced
    to the devices of that variable, and this method doesn't update the
    variable.
  options: a `tf.distribute.experimental.CommunicationOptions`. Options to
    perform collective operations. This overrides the default options if the
    `tf.distribute.Strategy` takes one in the constructor. See
    `tf.distribute.experimental.CommunicationOptions` for details of the
    options.

Returns:
  A tensor or value reduced to `destinations`.


class AriaGroupedExpertsGemm(nn.Module):

    def __init__(self, in_features, out_features, groups):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.groups = groups
        self.weight = nn.Parameter(torch.empty(groups, in_features, out_features))

    def forward(self, input, tokens_per_expert):
        """
        Perform grouped matrix multiplication.

        Args:
            input (`torch.Tensor`):
                Input tensor of shape (num_tokens, in_features).
            tokens_per_expert (`torch.Tensor`):
                Number of tokens assigned to each expert.

        Returns:
            torch.Tensor: Output tensor of shape (num_tokens, out_features).
        """
        return sequential_experts_gemm(input, self.weight, tokens_per_expert.cpu())

Grouped GEMM (General Matrix Multiplication) module for efficient expert computation.
This module utilizes the grouped_gemm library (https://github.com/fanshiqing/grouped_gemm)
for optimized performance. If the grouped_gemm library is not installed, it gracefully
falls back to a sequential GEMM implementation, which may be slower but ensures
functionality.

Args:
    in_features (`int`):
        Number of input features.
    out_features (`int`):
        Number of output features.
    groups (`int`):
        Number of expert groups.


def ceil(x):
    if any_symbolic_tensors((x,)):
        return Ceil().symbolic_call(x)
    return backend.numpy.ceil(x)

Return the ceiling of the input, element-wise.

The ceil of the scalar `x` is the smallest integer `i`, such that
`i >= x`.

Args:
    x: Input tensor.

Returns:
    The ceiling of each element in `x`, with float dtype.


def trunc(x):
    if any_symbolic_tensors((x,)):
        return Trunc().symbolic_call(x)
    return backend.numpy.trunc(x)

Return the truncated value of the input, element-wise.

The truncated value of the scalar `x` is the nearest integer `i` which is
closer to zero than `x` is. In short, the fractional part of the signed
number `x` is discarded.

Args:
    x: Input tensor.

Returns:
    The truncated value of each element in `x`.

Example:
>>> x = ops.array([-1.7, -1.5, -0.2, 0.2, 1.5, 1.7, 2.0])
>>> ops.trunc(x)
array([-1.0, -1.0, -0.0, 0.0, 1.0, 1.0, 2.0])


def forward(self, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.LongTensor]=None, decoder_input_values: Optional[torch.FloatTensor]=None, decoder_attention_mask: Optional[torch.LongTensor]=None, head_mask: Optional[torch.FloatTensor]=None, decoder_head_mask: Optional[torch.FloatTensor]=None, cross_attn_head_mask: Optional[torch.Tensor]=None, encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]]=None, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, speaker_embeddings: Optional[torch.FloatTensor]=None, labels: Optional[torch.FloatTensor]=None, stop_labels: Optional[torch.Tensor]=None) -> Union[Tuple, Seq2SeqSpectrogramOutput]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    if labels is not None:
        if decoder_input_values is None:
            decoder_input_values, decoder_attention_mask = shift_spectrograms_right(labels, self.config.reduction_factor, decoder_attention_mask)
        if self.config.use_guided_attention_loss:
            output_attentions = True
    outputs = self.speecht5(input_values=input_ids, attention_mask=attention_mask, decoder_input_values=decoder_input_values, decoder_attention_mask=decoder_attention_mask, head_mask=head_mask, decoder_head_mask=decoder_head_mask, cross_attn_head_mask=cross_attn_head_mask, encoder_outputs=encoder_outputs, past_key_values=past_key_values, use_cache=use_cache, speaker_embeddings=speaker_embeddings, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=True)
    outputs_before_postnet, outputs_after_postnet, logits = self.speech_decoder_postnet(outputs[0])
    loss = None
    if labels is not None:
        criterion = SpeechT5SpectrogramLoss(self.config)
        loss = criterion(attention_mask, outputs_before_postnet, outputs_after_postnet, logits, labels, outputs.cross_attentions)
    if not return_dict:
        output = (outputs_after_postnet,) + outputs[1:]
        return (loss,) + output if loss is not None else output
    return Seq2SeqSpectrogramOutput(loss=loss, spectrogram=outputs_after_postnet, past_key_values=outputs.past_key_values, decoder_hidden_states=outputs.decoder_hidden_states, decoder_attentions=outputs.decoder_attentions, cross_attentions=outputs.cross_attentions, encoder_last_hidden_state=outputs.encoder_last_hidden_state, encoder_hidden_states=outputs.encoder_hidden_states, encoder_attentions=outputs.encoder_attentions)

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

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

    [What are input IDs?](../glossary#input-ids)
decoder_input_values (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_mel_bins)`):
    Float values of input mel spectrogram.

    SpeechT5 uses an all-zero spectrum as the starting token for `decoder_input_values` generation. If
    `past_key_values` is used, optionally only the last `decoder_input_values` have to be input (see
    `past_key_values`).
decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
    Default behavior: generate a tensor that ignores pad tokens in `decoder_input_values`. Causal mask will
    also be used by default.

    If you want to change padding behavior, you should read [`SpeechT5Decoder._prepare_decoder_attention_mask`]
    and modify to your needs. See diagram 1 in [the paper](https://huggingface.co/papers/1910.13461) for more
    information on the default strategy.
cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
    Mask to nullify selected heads of the cross-attention modules. Mask values selected in `[0, 1]`:

    - 1 indicates the head is **not masked**,
    - 0 indicates the head is **masked**.
speaker_embeddings (`torch.FloatTensor` of shape `(batch_size, config.speaker_embedding_dim)`, *optional*):
    Tensor containing the speaker embeddings.
labels (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_mel_bins)`, *optional*):
    Float values of target mel spectrogram. Timesteps set to `-100.0` are ignored (masked) for the loss
    computation. Spectrograms can be obtained using [`SpeechT5Processor`]. See [`SpeechT5Processor.__call__`]
    for details.
stop_labels (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
    Binary tensor indicating the position of the stop token in the sequence.

Example:

```python
>>> from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan, set_seed
>>> import torch

>>> processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
>>> model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
>>> vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")

>>> inputs = processor(text="Hello, my dog is cute", return_tensors="pt")
>>> speaker_embeddings = torch.zeros((1, 512))  # or load xvectors from a file

>>> set_seed(555)  # make deterministic

>>> # generate speech
>>> speech = model.generate(inputs["input_ids"], speaker_embeddings=speaker_embeddings, vocoder=vocoder)
>>> speech.shape
torch.Size([15872])
```


def FromId(os_id, error_on_unknown=True):
    if not os_id:
        return None
    for operating_system in OperatingSystem._ALL:
        if operating_system.id == os_id:
            return operating_system
    if error_on_unknown:
        raise InvalidEnumValue(os_id, 'Operating System', [value.id for value in OperatingSystem._ALL])
    return None

Gets the enum corresponding to the given operating system id.

Args:
  os_id: str, The operating system id to parse
  error_on_unknown: bool, True to raise an exception if the id is unknown,
    False to just return None.

Raises:
  InvalidEnumValue: If the given value cannot be parsed.

Returns:
  OperatingSystemTuple, One of the OperatingSystem constants or None if the
  input is None.


def apply_rotary_unpadded(qkv, cos, sin, cu_seqlens: Optional[torch.Tensor]=None, max_seqlen: Optional[int]=None):
    return ApplyRotaryEmbUnpad.apply(qkv, cos, sin, cu_seqlens, max_seqlen)

Arguments:
    qkv: (total_nnz, 3, nheads, headdim) - input tensor for packed QKV.
    cos, sin: (seqlen_rotary, rotary_dim / 2)
    interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead
        of 1st half and 2nd half (GPT-NeoX style).
    inplace: if True, apply rotary embedding in-place.
    seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount.
        Most commonly used in inference when we have KV cache.
    cu_seqlens: (batch + 1,) or None
    max_seqlen: int
Return:
    out: (total_nnz, dim)
rotary_dim must be <= headdim
Apply rotary embedding to the first rotary_dim of x.


def build_graph(device, n, m, k, transpose_a, transpose_b, dtype):
    with ops.device('%s' % device):
        if not transpose_a:
            x = variable_v1.VariableV1(random_ops.random_uniform([n, m], dtype=dtype), use_resource=False)
        else:
            x = variable_v1.VariableV1(random_ops.random_uniform([m, n], dtype=dtype), use_resource=False)
        if not transpose_b:
            y = variable_v1.VariableV1(random_ops.random_uniform([m, k], dtype=dtype), use_resource=False)
        else:
            y = variable_v1.VariableV1(random_ops.random_uniform([k, m], dtype=dtype), use_resource=False)
        z = math_ops.matmul(x, y, transpose_a=transpose_a, transpose_b=transpose_b)
        return control_flow_ops.group(z)

Build a graph containing a sequence of matmul operations.

Args:
  device: String, the device to run on.
  n: tensor A's first dimension size.
  m: tensor A's second dimension size.
  k: tensor B's second dimension size.
  transpose_a: boolean value to show if tensor A is transposed.
  transpose_b: boolean value to show if tensor B is transposed.
  dtype: numpy data type of the input tensor.

Returns:
  A matmul operation to run()


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

Returns the expected continuation value for each path.

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

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


def assert_that(actual, matcher, label='assert_that', reify_windows=False, use_global_window=True):
    assert isinstance(actual, pvalue.PCollection), '%s is not a supported type for Beam assert' % type(actual)
    pipeline = actual.pipeline
    if getattr(actual.pipeline, 'result', None):
        raise RuntimeError('assert_that must be used within a beam.Pipeline context. ' + 'Prior to Beam 2.60.0, asserts outside of the context of a pipeline ' + 'were silently ignored, starting with Beam 2.60.0 this is no longer ' + 'allowed. To fix, move your assert_that call into your pipeline ' + 'context so that it is added before the pipeline is run. For more ' + 'information, see https://github.com/apache/beam/pull/30771')
    if label in pipeline.applied_labels:
        label_idx = 2
        while f'{label}_{label_idx}' in pipeline.applied_labels:
            label_idx += 1
        label = f'{label}_{label_idx}'
    if isinstance(matcher, _EqualToPerWindowMatcher):
        reify_windows = True
        use_global_window = True

    class ReifyTimestampWindow(DoFn):

        def process(self, element, timestamp=DoFn.TimestampParam, window=DoFn.WindowParam, pane_info=DoFn.PaneInfoParam):
            return [TestWindowedValue(element, timestamp, [window], pane_info)]

    class AddWindow(DoFn):

        def process(self, element, window=DoFn.WindowParam):
            yield (element, window)

    class AssertThat(PTransform):

        def expand(self, pcoll):
            if reify_windows:
                pcoll = pcoll | ParDo(ReifyTimestampWindow())
            keyed_singleton = pcoll.pipeline | Create([(None, None)])
            keyed_singleton.is_bounded = True
            if use_global_window:
                pcoll = pcoll | WindowInto(window.GlobalWindows())
            keyed_actual = pcoll | 'ToVoidKey' >> Map(lambda v: (None, v))
            keyed_actual.is_bounded = True
            plain_actual = (keyed_singleton, keyed_actual) | 'Group' >> CoGroupByKey() | 'Unkey' >> Map(lambda k_values: list(k_values[1][1]))
            if not use_global_window:
                plain_actual = plain_actual | 'AddWindow' >> ParDo(AddWindow())
            return plain_actual | 'Match' >> Map(matcher)

        def default_label(self):
            return label
    return actual | AssertThat()

A PTransform that checks a PCollection has an expected value.

Note that assert_that should be used only for testing pipelines since the
check relies on materializing the entire PCollection being checked.

Args:
  actual: A PCollection.
  matcher: A matcher function taking as argument the actual value of a
    materialized PCollection. The matcher validates this actual value against
    expectations and raises BeamAssertException if they are not met.
  label: Optional string label. This is needed in case several assert_that
    transforms are introduced in the same pipeline.
  reify_windows: If True, matcher is passed a list of TestWindowedValue.
  use_global_window: If False, matcher is passed a dictionary of
    (k, v) = (window, elements in the window).

Returns:
  Ignored.


def auto_plot_array(*, video_min_num_frames: int=15, height: None | int | tuple[int, int]=(100, 250), show_images_kwargs: Optional[dict[str, Any]]=None, show_videos_kwargs: Optional[dict[str, Any]]=None) -> None:
    ipython = IPython.get_ipython()
    if ipython is None:
        return
    array_repr_html_fn = functools.partial(array_repr_html, video_min_num_frames=video_min_num_frames, height=height, show_images_kwargs=show_images_kwargs, show_videos_kwargs=show_videos_kwargs)
    print('Display big np/tf/jax arrays as image for nicer IPython display')
    formatter = ipython.display_formatter.formatters['text/html']
    try:
        jnp = enp.lazy.jnp
    except ImportError:
        pass
    else:
        jax_array_cls = type(jnp.zeros(shape=()))
        formatter.for_type(jax_array_cls, array_repr_html_fn)
    try:
        tf = enp.lazy.tf
    except ImportError:
        pass
    else:
        formatter.for_type(tf.Tensor, array_repr_html_fn)
    try:
        torch = enp.lazy.torch
    except ImportError:
        pass
    else:
        formatter.for_type(torch.Tensor, array_repr_html_fn)
    formatter.for_type(enp.lazy.np.ndarray, array_repr_html_fn)

If called, 2d/3d imgage arrays will be plotted as images in colab/jupyter.

Usage:

>>> ecolab.auto_plot_array()
>>> np.zeros((28, 28, 3))  # Displayed as image

Args:
  video_min_num_frames: Video `(num_frames, h, w, c)` with less than this
    number of frames will be displayed as individual images
  height: `(min, max)` image height in pixels. Images smaller/larger will be
    reshaped. `None` to disable. If a single number, assume `min == max`.
  show_images_kwargs: Kwargs forwarded to `mediapy.show_images`
  show_videos_kwargs: Kwargs forwarded to `mediapy.show_videos`


def _gen_save_and_restore_functions(checkpoint_factory_map: object_identity.ObjectIdentityDictionary) -> object_identity.ObjectIdentityDictionary:
    saveable_fn_map = object_identity.ObjectIdentityDictionary()
    for obj, factory_data_list in checkpoint_factory_map.items():
        if resource_variable_ops.is_resource_variable(obj) or not factory_data_list:
            continue
        if factory_data_list[0].name == trackable_utils.SERIALIZE_TO_TENSORS_NAME:
            assert len(factory_data_list) == 1
            saveable_fn_map[obj] = {trackable_utils.SERIALIZE_TO_TENSORS_NAME: tracing_utils.trace_save_and_restore(obj)}
        else:
            saveable_fn_map[obj] = trace_saveable_util.trace_save_restore_function_map(obj, factory_data_list)
    return saveable_fn_map

Generates global and individual save/restore concrete functions.

The global functions records the ops to save and restore the entire object to
a file prefix, while the individual functions save and restore value tensors
for resources.

This function is intended to run on the output of
`save_util_v1.get_checkpoint_factories_and_keys(object_names)`,
which returns the generated a map of `_CheckpointFactoryData`.

Args:
  checkpoint_factory_map: A dictionary mapping trackable objects to
    a list of `_CheckpointFactoryData`.

Returns:
  Tuple of (
    saveable_fn_map: Maps obj -> factory name -> (concrete save, restore)
    )


class TFGPT2DoubleHeadsModelOutput(ModelOutput):
    logits: Optional[tf.Tensor] = None
    mc_logits: Optional[tf.Tensor] = None
    past_key_values: List[tf.Tensor] | None = None
    hidden_states: Tuple[tf.Tensor] | None = None
    attentions: Tuple[tf.Tensor] | None = None

Base class for outputs of models predicting if two sentences are consecutive or not.

Args:
    logits (`tf.Tensor` of shape `(batch_size, num_choices, sequence_length, config.vocab_size)`):
        Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
    mc_logits (`tf.Tensor` of shape `(batch_size, num_choices)`):
        Prediction scores of the multiple choice classification head (scores for each choice before SoftMax).
    past_key_values (`List[tf.Tensor]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
        List of `tf.Tensor` of length `config.n_layers`, with each tensor of shape `(2, batch_size, num_heads,
        sequence_length, embed_size_per_head)`).

        Contains pre-computed hidden-states (key and values in the attention blocks) that can be used (see
        `past_key_values` input) to speed up sequential decoding.
    hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape
        `(batch_size, sequence_length, hidden_size)`.

        Hidden-states of the model at the output of each layer plus the initial embedding outputs.
    attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
        Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
        sequence_length)`.

        Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
        heads.


class IncMeanTracker(WindowedTracker, MeanTracker):

    def __init__(self, window_mode, **kwargs):
        super().__init__(window_mode=window_mode, **kwargs)
        self._mean = 0

    def push(self, x):
        """Pushes a new value and updates the incremental mean.

    Args:
      x: The new value to be pushed.
    """
        if not math.isnan(x):
            self._n += 1
            delta = x - self._mean
        else:
            delta = 0
        if self._window_mode == WindowMode.SLIDING:
            if len(self._queue) >= self._window_size and (not math.isnan((old_x := self.pop()))):
                self._n -= 1
                delta += self._mean - old_x
            super().push(x)
        if self._n > 0:
            self._mean += delta / self._n
        else:
            self._mean = 0

    def get(self):
        """Returns the current incremental mean.

    Returns:
      float: The current incremental mean value.
             Returns NaN if no valid (non-NaN) values have been pushed.
    """
        if self._n < 1:
            return float('nan')
        return self._mean

Base class for incremental mean trackers.

This class implements incremental calculation of the mean, which is more
efficient for streaming data as it updates the mean with each new data point
instead of recalculating from scratch.

Args:
  window_mode: A `WindowMode` enum specifying whether the window is `LANDMARK`
    or `SLIDING`.
  **kwargs: Keyword arguments passed to the parent class constructor.


def freeze_graph_with_def_protos(input_graph_def: Optional[graph_pb2.GraphDef], input_saver_def: Optional[saver_pb2.SaverDef], input_checkpoint: Optional[str], output_node_names: str, restore_op_name: Optional[str], filename_tensor_name: Optional[str], output_graph: str, clear_devices: bool, initializer_nodes: str, variable_names_whitelist: str='', variable_names_denylist: str='', input_meta_graph_def: Optional[meta_graph_pb2.MetaGraphDef]=None, input_saved_model_dir: Optional[str]=None, saved_model_tags: Optional[List[str]]=None, checkpoint_version: int=saver_pb2.SaverDef.V2) -> graph_pb2.GraphDef:
    del restore_op_name, filename_tensor_name
    if not input_saved_model_dir and (not checkpoint_management.checkpoint_exists(input_checkpoint)):
        raise ValueError("Input checkpoint '" + input_checkpoint + "' doesn't exist!")
    if not output_node_names:
        raise ValueError('You need to supply the name of a node to --output_node_names.')
    if clear_devices:
        if input_meta_graph_def:
            for node in input_meta_graph_def.graph_def.node:
                node.device = ''
        elif input_graph_def:
            for node in input_graph_def.node:
                node.device = ''
    if input_graph_def:
        _ = importer.import_graph_def(input_graph_def, name='')
    with session.Session() as sess:
        if input_saver_def:
            saver = saver_lib.Saver(saver_def=input_saver_def, write_version=checkpoint_version)
            saver.restore(sess, input_checkpoint)
        elif input_meta_graph_def:
            restorer = saver_lib.import_meta_graph(input_meta_graph_def, clear_devices=True)
            restorer.restore(sess, input_checkpoint)
            if initializer_nodes:
                sess.run(initializer_nodes.replace(' ', '').split(','))
        elif input_saved_model_dir:
            if saved_model_tags is None:
                saved_model_tags = []
            loader.load(sess, saved_model_tags, input_saved_model_dir)
        else:
            var_list = {}
            reader = py_checkpoint_reader.NewCheckpointReader(input_checkpoint)
            var_to_shape_map = reader.get_variable_to_shape_map()
            all_partition_variable_names = [tensor.name.split(':')[0] for op in sess.graph.get_operations() for tensor in op.values() if re.search('/part_\\d+/', tensor.name)]
            has_partition_var = False
            for key in var_to_shape_map:
                try:
                    tensor = sess.graph.get_tensor_by_name(key + ':0')
                    if any((key in name for name in all_partition_variable_names)):
                        has_partition_var = True
                except KeyError:
                    continue
                var_list[key] = tensor
            try:
                saver = saver_lib.Saver(var_list=var_list, write_version=checkpoint_version)
            except TypeError as e:
                if has_partition_var:
                    raise ValueError('Models containing partition variables cannot be converted from checkpoint files. Please pass in a SavedModel using the flag --input_saved_model_dir.')
                elif _has_no_variables(sess):
                    raise ValueError('No variables were found in this model. It is likely the model was frozen previously. You cannot freeze a graph twice.')
                else:
                    raise e
            saver.restore(sess, input_checkpoint)
            if initializer_nodes:
                sess.run(initializer_nodes.replace(' ', '').split(','))
        variable_names_whitelist = variable_names_whitelist.replace(' ', '').split(',') if variable_names_whitelist else None
        variable_names_denylist = variable_names_denylist.replace(' ', '').split(',') if variable_names_denylist else None
        if input_meta_graph_def:
            output_graph_def = convert_to_constants.convert_variables_to_constants(sess, input_meta_graph_def.graph_def, output_node_names.replace(' ', '').split(','), variable_names_whitelist=variable_names_whitelist, variable_names_blacklist=variable_names_denylist)
        else:
            output_graph_def = convert_to_constants.convert_variables_to_constants(sess, input_graph_def, output_node_names.replace(' ', '').split(','), variable_names_whitelist=variable_names_whitelist, variable_names_blacklist=variable_names_denylist)
    if output_graph:
        with gfile.GFile(output_graph, 'wb') as f:
            f.write(output_graph_def.SerializeToString(deterministic=True))
    return output_graph_def

Converts all variables in a graph and checkpoint into constants.

Args:
  input_graph_def: A `GraphDef`.
  input_saver_def: A `SaverDef` (optional).
  input_checkpoint: The prefix of a V1 or V2 checkpoint, with V2 taking
    priority.  Typically the result of `Saver.save()` or that of
    `tf.train.latest_checkpoint()`, regardless of sharded/non-sharded or
    V1/V2.
  output_node_names: The name(s) of the output nodes, comma separated.
  restore_op_name: Unused.
  filename_tensor_name: Unused.
  output_graph: String where to write the frozen `GraphDef`.
  clear_devices: A Bool whether to remove device specifications.
  initializer_nodes: Comma separated string of initializer nodes to run before
    freezing.
  variable_names_whitelist: The set of variable names to convert (optional, by
    default, all variables are converted).
  variable_names_denylist: The set of variable names to omit converting to
    constants (optional).
  input_meta_graph_def: A `MetaGraphDef` (optional),
  input_saved_model_dir: Path to the dir with TensorFlow 'SavedModel' file and
    variables (optional).
  saved_model_tags: Group of comma separated tag(s) of the MetaGraphDef to
    load, in string format (optional).
  checkpoint_version: Tensorflow variable file format (saver_pb2.SaverDef.V1
    or saver_pb2.SaverDef.V2)

Returns:
  Location of the output_graph_def.


def index_add(x, idx, y):
    return _index_update_helper(tf_np.ndarray._with_index_add, x, idx, y)

Pure equivalent of `x[idx] += y`.

Returns the value of x that would result from the NumPy-style indexed
assignment `x[idx] += y`. Because it's a pure function, `x` itself won't be
changed.

Args:
  x: an array with the values to be updated.
  idx: a Numpy-style index, consisting of `None`, integers, slice objects,
    ellipses, ndarrays with integer dtypes, or a tuple of the above.
  y: the array of updates. `y` must be broadcastable to the shape of the array
    that would be returned by `x[idx]`.

Returns:
  The updated version of `x`.


class DataframeTransform(transforms.PTransform):

    def __init__(self, func, proxy=None, yield_elements='schemas', include_indexes=False):
        self._func = func
        self._proxy = proxy
        self._yield_elements = yield_elements
        self._include_indexes = include_indexes

    def expand(self, input_pcolls):
        from apache_beam.dataframe import convert
        input_dict: dict[Any, PCollection] = _flatten(input_pcolls)
        proxies = _flatten(self._proxy) if self._proxy is not None else {tag: None for tag in input_dict}
        input_frames: dict[Any, frame_base.DeferredFrame] = {k: convert.to_dataframe(pc, proxies[k]) for k, pc in input_dict.items()}
        frames_input = _substitute(input_pcolls, input_frames)
        if isinstance(frames_input, dict):
            result_frames = self._func(**frames_input)
        elif isinstance(frames_input, tuple):
            result_frames = self._func(*frames_input)
        else:
            result_frames = self._func(frames_input)
        result_frames_dict = _flatten(result_frames)
        keys = list(result_frames_dict.keys())
        result_frames_tuple = tuple((result_frames_dict[key] for key in keys))
        result_pcolls_tuple = convert.to_pcollection(*result_frames_tuple, label='Eval', always_return_tuple=True, yield_elements=self._yield_elements, include_indexes=self._include_indexes)
        result_pcolls_dict = dict(zip(keys, result_pcolls_tuple))
        return _substitute(result_frames, result_pcolls_dict)

A PTransform for applying function that takes and returns dataframes
to one or more PCollections.

:class:`DataframeTransform` will accept a PCollection with a `schema`_ and
batch it into :class:`~pandas.DataFrame` instances if necessary::

    (pcoll | beam.Select(key=..., foo=..., bar=...)
           | DataframeTransform(lambda df: df.group_by('key').sum()))

It is also possible to process a PCollection of :class:`~pandas.DataFrame`
instances directly, in this case a "proxy" must be provided. For example, if
``pcoll`` is a PCollection of DataFrames, one could write::

    pcoll | DataframeTransform(lambda df: df.group_by('key').sum(), proxy=...)

To pass multiple PCollections, pass a tuple of PCollections wich will be
passed to the callable as positional arguments, or a dictionary of
PCollections, in which case they will be passed as keyword arguments.

Args:
  yield_elements: (optional, default: "schemas") If set to ``"pandas"``,
      return PCollection(s) containing the raw Pandas objects
      (:class:`~pandas.DataFrame` or :class:`~pandas.Series` as appropriate).
      If set to ``"schemas"``, return an element-wise PCollection, where
      DataFrame and Series instances are expanded to one element per row.
      DataFrames are converted to `schema-aware`_ PCollections, where column
      values can be accessed by attribute.
  include_indexes: (optional, default: False) When
     ``yield_elements="schemas"``, if ``include_indexes=True``, attempt to
     include index columns in the output schema for expanded DataFrames.
     Raises an error if any of the index levels are unnamed (name=None), or if
     any of the names are not unique among all column and index names.
  proxy: (optional) An empty :class:`~pandas.DataFrame` or
      :class:`~pandas.Series` instance with the same ``dtype`` and ``name``
      as the elements of the input PCollection. Required when input
      PCollection :class:`~pandas.DataFrame` or :class:`~pandas.Series`
      elements. Ignored when input PCollection has a `schema`_.

.. _schema:
  https://beam.apache.org/documentation/programming-guide/#what-is-a-schema
.. _schema-aware:
  https://beam.apache.org/documentation/programming-guide/#what-is-a-schema


def reraise(e: Exception, prefix: Optional[_Str]=None, suffix: Optional[_Str]=None) -> NoReturn:
    __tracebackhide__ = True
    prefix = prefix() if callable(prefix) else prefix
    suffix = suffix() if callable(suffix) else suffix
    prefix = prefix or ''
    suffix = '\n' + suffix if suffix else ''
    msg = f'{prefix}{e}{suffix}'

    class WrappedException(type(e)):
        """Exception proxy with additional message."""

        def __init__(self, msg):
            Exception.__init__(self, msg)

        def __getattr__(self, name: str):
            return getattr(e, name)
        __repr__ = BaseException.__repr__
        __str__ = BaseException.__str__
    WrappedException.__name__ = type(e).__name__
    WrappedException.__qualname__ = type(e).__qualname__
    WrappedException.__module__ = type(e).__module__
    new_exception = WrappedException(msg)
    raise new_exception.with_traceback(e.__traceback__) from e.__cause__

Reraise an exception with an additional message.

Benefit: Contrary to `raise ... from ...` and
`raise Exception().with_traceback(tb)`, this function will:

* Keep the original exception type, attributes,...
* Avoid multi-nested `During handling of the above exception, another
  exception occurred`. Only the single original stacktrace is displayed.

This result in cleaner and more compact error messages.

Usage:

```
try:
  fn(x)
except Exception as e:
  epy.reraise(e, prefix=f'Error for {x}: ')
```

Args:
  e: Exception to reraise
  prefix: Prefix to add to the exception message.
  suffix: Suffix to add to the exception message.


def batch_encode_candidates(self, text, **kwargs):
    kwargs['padding'] = PaddingStrategy.MAX_LENGTH
    batch_text = text
    batch_text_pair = kwargs.pop('text_pair', None)
    return_tensors = kwargs.pop('return_tensors', None)
    output_data = {'input_ids': [], 'attention_mask': [], 'token_type_ids': []}
    for idx, candidate_text in enumerate(batch_text):
        if batch_text_pair is not None:
            candidate_text_pair = batch_text_pair[idx]
        else:
            candidate_text_pair = None
        encoded_candidates = super().__call__(candidate_text, candidate_text_pair, return_tensors=None, **kwargs)
        encoded_input_ids = encoded_candidates.get('input_ids')
        encoded_attention_mask = encoded_candidates.get('attention_mask')
        encoded_token_type_ids = encoded_candidates.get('token_type_ids')
        if encoded_input_ids is not None:
            output_data['input_ids'].append(encoded_input_ids)
        if encoded_attention_mask is not None:
            output_data['attention_mask'].append(encoded_attention_mask)
        if encoded_token_type_ids is not None:
            output_data['token_type_ids'].append(encoded_token_type_ids)
    output_data = {key: item for key, item in output_data.items() if len(item) != 0}
    return BatchEncoding(output_data, tensor_type=return_tensors)

Encode a batch of text or text pair. This method is similar to regular __call__ method but has the following
differences:

    1. Handle additional num_candidate axis. (batch_size, num_candidates, text)
    2. Always pad the sequences to *max_length*.
    3. Must specify *max_length* in order to stack packs of candidates into a batch.

    - single sequence: `[CLS] X [SEP]`
    - pair of sequences: `[CLS] A [SEP] B [SEP]`

Args:
    text (`List[List[str]]`):
        The batch of sequences to be encoded. Each sequence must be in this format: (batch_size,
        num_candidates, text).
    text_pair (`List[List[str]]`, *optional*):
        The batch of sequences to be encoded. Each sequence must be in this format: (batch_size,
        num_candidates, text).
    **kwargs:
        Keyword arguments of the __call__ method.

Returns:
    [`BatchEncoding`]: Encoded text or text pair.

Example:

```python
>>> from transformers import RealmTokenizerFast

>>> # batch_size = 2, num_candidates = 2
>>> text = [["Hello world!", "Nice to meet you!"], ["The cute cat.", "The adorable dog."]]

>>> tokenizer = RealmTokenizerFast.from_pretrained("google/realm-cc-news-pretrained-encoder")
>>> tokenized_text = tokenizer.batch_encode_candidates(text, max_length=10, return_tensors="pt")
```


def create_gpu_capa_map(match_list, generate_csv=False, filename='compute_capability'):
    gpu_capa = collections.OrderedDict()
    include = False
    gpu = ''
    cnt = 0
    mismatch_cnt = 0
    for match in match_list:
        if 'Products' in match:
            if not include:
                include = True
            continue
        elif 'www' in match:
            include = False
            break
        if include:
            if gpu:
                if gpu in gpu_capa:
                    gpu_capa[gpu].append(match)
                else:
                    gpu_capa[gpu] = [match]
                gpu = ''
                cnt += 1
                if len(list(gpu_capa.keys())) < cnt:
                    mismatch_cnt += 1
                    cnt = len(list(gpu_capa.keys()))
            else:
                gpu = match
    if generate_csv:
        f_name = filename + '.csv'
        write_csv_from_dict(f_name, gpu_capa)
    return gpu_capa

Generates a map between GPU types and corresponding compute capability.

This method is used for retrieving CUDA compute capability from the web only.

Args:
  match_list: List of all CUDA compute capability detected from the webpage.
  generate_csv: Boolean for creating csv file to store results.
  filename: String that is the name of the csv file (without `.csv` ending).

Returns:
  OrderedDict that lists in the incoming order of all CUDA compute capability
  provided as `match_list`.


def __call__(self, image: Union[str, 'Image.Image', List[Dict[str, Any]]], candidate_labels: Optional[Union[str, List[str]]]=None, **kwargs: Any) -> Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]]:
    if 'text_queries' in kwargs:
        candidate_labels = kwargs.pop('text_queries')
    if isinstance(image, (str, Image.Image)):
        inputs = {'image': image, 'candidate_labels': candidate_labels}
    elif isinstance(image, (list, tuple)) and valid_images(image):
        return list(super().__call__(({'image': img, 'candidate_labels': labels} for img, labels in zip(image, candidate_labels)), **kwargs))
    else:
        '\n            Supports the following format\n            - {"image": image, "candidate_labels": candidate_labels}\n            - [{"image": image, "candidate_labels": candidate_labels}]\n            - Generator and datasets\n            This is a common pattern in other multimodal pipelines, so we support it here as well.\n            '
        inputs = image
    results = super().__call__(inputs, **kwargs)
    return results

Detect objects (bounding boxes & classes) in the image(s) passed as inputs.

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

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

        You can use this parameter to send directly a list of images, or a dataset or a generator like so:

        ```python
        >>> from transformers import pipeline

        >>> detector = pipeline(model="google/owlvit-base-patch32", task="zero-shot-object-detection")
        >>> detector(
        ...     [
        ...         {
        ...             "image": "http://images.cocodataset.org/val2017/000000039769.jpg",
        ...             "candidate_labels": ["cat", "couch"],
        ...         },
        ...         {
        ...             "image": "http://images.cocodataset.org/val2017/000000039769.jpg",
        ...             "candidate_labels": ["cat", "couch"],
        ...         },
        ...     ]
        ... )
        [[{'score': 0.287, 'label': 'cat', 'box': {'xmin': 324, 'ymin': 20, 'xmax': 640, 'ymax': 373}}, {'score': 0.25, 'label': 'cat', 'box': {'xmin': 1, 'ymin': 55, 'xmax': 315, 'ymax': 472}}, {'score': 0.121, 'label': 'couch', 'box': {'xmin': 4, 'ymin': 0, 'xmax': 642, 'ymax': 476}}], [{'score': 0.287, 'label': 'cat', 'box': {'xmin': 324, 'ymin': 20, 'xmax': 640, 'ymax': 373}}, {'score': 0.254, 'label': 'cat', 'box': {'xmin': 1, 'ymin': 55, 'xmax': 315, 'ymax': 472}}, {'score': 0.121, 'label': 'couch', 'box': {'xmin': 4, 'ymin': 0, 'xmax': 642, 'ymax': 476}}]]
        ```

    candidate_labels (`str` or `List[str]` or `List[List[str]]`):
        What the model should recognize in the image.

    threshold (`float`, *optional*, defaults to 0.1):
        The probability necessary to make a prediction.

    top_k (`int`, *optional*, defaults to None):
        The number of top predictions that will be returned by the pipeline. If the provided number is `None`
        or higher than the number of predictions available, it will default to the number of predictions.

    timeout (`float`, *optional*, defaults to None):
        The maximum time in seconds to wait for fetching images from the web. If None, no timeout is set and
        the call may block forever.

Return:
    A list of lists containing prediction results, one list per input image. Each list contains dictionaries
    with the following keys:

    - **label** (`str`) -- Text query corresponding to the found object.
    - **score** (`float`) -- Score corresponding to the object (between 0 and 1).
    - **box** (`Dict[str,int]`) -- Bounding box of the detected object in image's original size. It is a
      dictionary with `x_min`, `x_max`, `y_min`, `y_max` keys.


class TestResultRecord:

    def __init__(self, t_name, t_class=None):
        self.test_name = t_name
        self.test_class = t_class
        self.begin_time = None
        self.end_time = None
        self.uid = None
        self.signature = None
        self.retry_parent = None
        self.parent = None
        self.termination_signal = None
        self.extra_errors = collections.OrderedDict()
        self.result = None

    @property
    def details(self):
        """String description of the cause of the test's termination.

    Note a passed test can have this as well due to the explicit pass
    signal. If the test passed implicitly, this field would be None.
    """
        if self.termination_signal:
            return self.termination_signal.details

    @property
    def termination_signal_type(self):
        """Type name of the signal that caused the test's termination.

    Note a passed test can have this as well due to the explicit pass
    signal. If the test passed implicitly, this field would be None.
    """
        if self.termination_signal:
            return self.termination_signal.type

    @property
    def stacktrace(self):
        """The stacktrace string for the exception that terminated the test."""
        if self.termination_signal:
            return self.termination_signal.stacktrace

    @property
    def extras(self):
        """User defined extra information of the test result.

    Must be serializable.
    """
        if self.termination_signal:
            return self.termination_signal.extras

    def test_begin(self):
        """Call this when the test begins execution.

    Sets the begin_time of this record.
    """
        self.begin_time = utils.get_current_epoch_time()
        self.signature = '%s-%s' % (self.test_name, self.begin_time)

    def _test_end(self, result, e):
        """Marks the end of the test logic.

    Args:
      result: One of the TEST_RESULT enums in TestResultEnums.
      e: A test termination signal (usually an exception object). It can
        be any exception instance or of any subclass of
        mobly.signals.TestSignal.
    """
        if self.begin_time is not None:
            self.end_time = utils.get_current_epoch_time()
        self.result = result
        if e:
            self.termination_signal = ExceptionRecord(e)

    def update_record(self):
        """Updates the content of a record.

    Several display fields like "details" and "stacktrace" need to be
    updated based on the content of the record object.

    As the content of the record change, call this method to update all
    the appropirate fields.
    """
        if self.extra_errors:
            if self.result != TestResultEnums.TEST_RESULT_FAIL:
                self.result = TestResultEnums.TEST_RESULT_ERROR
        if not self.termination_signal and self.extra_errors:
            _, self.termination_signal = self.extra_errors.popitem(last=False)

    def test_pass(self, e=None):
        """To mark the test as passed in this record.

    Args:
      e: An instance of mobly.signals.TestPass.
    """
        self._test_end(TestResultEnums.TEST_RESULT_PASS, e)

    def test_fail(self, e=None):
        """To mark the test as failed in this record.

    Only test_fail does instance check because we want 'assert xxx' to also
    fail the test same way assert_true does.

    Args:
      e: An exception object. It can be an instance of AssertionError or
        mobly.base_test.TestFailure.
    """
        self._test_end(TestResultEnums.TEST_RESULT_FAIL, e)

    def test_skip(self, e=None):
        """To mark the test as skipped in this record.

    Args:
      e: An instance of mobly.signals.TestSkip.
    """
        self._test_end(TestResultEnums.TEST_RESULT_SKIP, e)

    def test_error(self, e=None):
        """To mark the test as error in this record.

    Args:
      e: An exception object.
    """
        self._test_end(TestResultEnums.TEST_RESULT_ERROR, e)

    def add_error(self, position, e):
        """Add extra error happened during a test.

    If the test has passed or skipped, this will mark the test result as
    ERROR.

    If an error is added the test record, the record's result is equivalent
    to the case where an uncaught exception happened.

    If the test record has not recorded any error, the newly added error
    would be the main error of the test record. Otherwise the newly added
    error is added to the record's extra errors.

    Args:
      position: string, where this error occurred, e.g. 'teardown_test'.
      e: An exception or a `signals.ExceptionRecord` object.
    """
        if self.result != TestResultEnums.TEST_RESULT_FAIL:
            self.result = TestResultEnums.TEST_RESULT_ERROR
        if position in self.extra_errors:
            raise Error('An exception is already recorded with position "%s", cannot reuse.' % position)
        if isinstance(e, ExceptionRecord):
            self.extra_errors[position] = e
        else:
            self.extra_errors[position] = ExceptionRecord(e, position=position)

    def __str__(self):
        d = self.to_dict()
        kv_pairs = ['%s = %s' % (k, v) for k, v in d.items()]
        s = ', '.join(kv_pairs)
        return s

    def __repr__(self):
        """This returns a short string representation of the test record."""
        t = utils.epoch_to_human_time(self.begin_time)
        return f'{t} {self.test_name} {self.result}'

    def to_dict(self):
        """Gets a dictionary representating the content of this class.

    Returns:
      A dictionary representating the content of this class.
    """
        d = {}
        d[TestResultEnums.RECORD_NAME] = self.test_name
        d[TestResultEnums.RECORD_CLASS] = self.test_class
        d[TestResultEnums.RECORD_BEGIN_TIME] = self.begin_time
        d[TestResultEnums.RECORD_END_TIME] = self.end_time
        d[TestResultEnums.RECORD_RESULT] = self.result
        d[TestResultEnums.RECORD_UID] = self.uid
        d[TestResultEnums.RECORD_SIGNATURE] = self.signature
        d[TestResultEnums.RECORD_RETRY_PARENT] = self.retry_parent.signature if self.retry_parent else None
        d[TestResultEnums.RECORD_PARENT] = {'parent': self.parent[0].signature, 'type': self.parent[1].value} if self.parent else None
        d[TestResultEnums.RECORD_EXTRAS] = self.extras
        d[TestResultEnums.RECORD_DETAILS] = self.details
        d[TestResultEnums.RECORD_TERMINATION_SIGNAL_TYPE] = self.termination_signal_type
        d[TestResultEnums.RECORD_EXTRA_ERRORS] = {key: value.to_dict() for key, value in self.extra_errors.items()}
        d[TestResultEnums.RECORD_STACKTRACE] = self.stacktrace
        return d

A record that holds the information of a single test.

The record object holds all information of a test, including all the
exceptions occurred during the test.

A test can terminate for two reasons:
  1. the test function executes to the end and completes naturally.
  2. the test is terminated by an exception, which we call
   "termination signal".

The termination signal is treated differently. Its content are extracted
into first-tier attributes of the record object, like `details` and
`stacktrace`, for easy consumption.

Note the termination signal is not always an error, it can also be explicit
pass signal or abort/skip signals.

Attributes:
  test_name: string, the name of the test.
  begin_time: Epoch timestamp of when the test started.
  end_time: Epoch timestamp of when the test ended.
  uid: User-defined unique identifier of the test.
  signature: string, unique identifier of a test record, the value is
    generated by Mobly.
  retry_parent: [DEPRECATED] Use the `parent` field instead.
  parent: tuple[TestResultRecord, TestParentType], set for multiple iterations
    of a test. This is the test result record of the previous iteration.
    Parsers can use this field to construct the chain of execution for each test.
  termination_signal: ExceptionRecord, the main exception of the test.
  extra_errors: OrderedDict, all exceptions occurred during the entire
    test lifecycle. The order of occurrence is preserved.
  result: TestResultEnum.TEST_RESULT_*, PASS/FAIL/SKIP.


def _build_ragged_tensor_from_value_ranges(starts, limits, step, values):
    if step is None:
        step = 1
    step = ops.convert_to_tensor(step, name='step')
    if step.dtype.is_integer:
        step = math_ops.cast(step, starts.dtype)
    else:
        raise TypeError('slice strides must be integers or None')
    value_indices = ragged_math_ops.range(starts, limits, step, row_splits_dtype=starts.dtype)
    if isinstance(values, ragged_tensor.RaggedTensor):
        gathered_values = ragged_gather_ops.gather(params=values, indices=value_indices.values)
    else:
        gathered_values = array_ops.gather(params=values, indices=value_indices.values)
    return value_indices.with_values(gathered_values)

Returns a `RaggedTensor` containing the specified sequences of values.

Returns a RaggedTensor `output` where:

```python
output.shape[0] = starts.shape[0]
output[i] = values[starts[i]:limits[i]:step]
```

Requires that `starts.shape == limits.shape` and
`0 <= starts[i] <= limits[i] <= values.shape[0]`.

Args:
  starts: 1D integer Tensor specifying the start indices for the sequences of
    values to include.
  limits: 1D integer Tensor specifying the limit indices for the sequences of
    values to include.
  step: Integer value specifying the step size for strided slices.
  values: The set of values to select from.

Returns:
  A `RaggedTensor`.

Raises:
  ValueError: Until the prerequisite ops are checked in.


def compress(element):
    element_spec = structure.type_spec_from_value(element)
    tensor_list = structure.to_tensor_list(element_spec, element)
    return ged_ops.compress_element(tensor_list)

Compress a dataset element.

Args:
  element: A nested structure of types supported by Tensorflow.

Returns:
  A variant tensor representing the compressed element. This variant can be
  passed to `uncompress` to get back the original element.


def _get_resized_embeddings(self, old_embeddings, new_num_tokens=None) -> tf.Variable:
    old_embedding_dim = shape_list(old_embeddings)[1]
    init_range = getattr(self.config, 'initializer_range', 0.02)
    embeddings_mask, current_embeddings = init_copy_embeddings(old_embeddings, new_num_tokens)
    new_embeddings = self.add_weight(name=old_embeddings.name.split(':')[0], shape=[new_num_tokens, old_embedding_dim], initializer=get_initializer(init_range), dtype=tf.float32)
    init_embeddings = tf.where(embeddings_mask, current_embeddings, new_embeddings.value())
    new_embeddings.assign(init_embeddings)
    return new_embeddings

Build a resized Embedding weights from a provided token Embedding weights. Increasing the size will add newly
initialized vectors at the end. Reducing the size will remove vectors from the end

Args:
    old_embeddings (`tf.Variable`):
        Old embeddings to be resized.
    new_num_tokens (`int`, *optional*):
        New number of tokens in the embedding matrix.

        Increasing the size will add newly initialized vectors at the end. Reducing the size will remove
        vectors from the end. If not provided or `None`, just returns a pointer to the input tokens
        `tf.Variable` module of the model without doing anything.

Return:
    `tf.Variable`: Pointer to the resized Embedding Module or the old Embedding Module if `new_num_tokens` is
    `None`


def create_dummy_object(name: str, backend_name: str) -> str:
    if name.isupper():
        return DUMMY_CONSTANT.format(name)
    elif name.islower():
        return DUMMY_FUNCTION.format(name, backend_name)
    else:
        return DUMMY_CLASS.format(name, backend_name)

Create the code for a dummy object.

Args:
    name (`str`): The name of the object.
    backend_name (`str`): The name of the backend required for that object.

Returns:
    `str`: The code of the dummy object.


def GetExcludePatternsForDir(dirname):
    ignore_patterns = []
    yapfignore_file = os.path.join(dirname, '.yapfignore')
    if os.path.exists(yapfignore_file):
        ignore_patterns += _GetExcludePatternsFromYapfIgnore(yapfignore_file)
    pyproject_toml_file = os.path.join(dirname, 'pyproject.toml')
    if os.path.exists(pyproject_toml_file):
        ignore_patterns += _GetExcludePatternsFromPyprojectToml(pyproject_toml_file)
    return ignore_patterns

Return patterns of files to exclude from ignorefile in a given directory.

Looks for .yapfignore in the directory dirname.

Arguments:
  dirname: (unicode) The name of the directory.

Returns:
  A List of file patterns to exclude if ignore file is found, otherwise empty
  List.


def densifying_unary(func):

    @functools.wraps(func)
    def sparse_wrapper(x, *args, **kwargs):
        if isinstance(x, jax_sparse.JAXSparse):
            x = x.todense()
        return func(x, *args, **kwargs)
    return sparse_wrapper

Decorator to add support for `JAXSparse` tensors (including `BCOO`) to a
non-zero-preserving element-wise unary operator.

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

- The operator must be element-wise
- The operator must be unary (one input tensor and one output tensor)
- The operator must return a tensor of the same shape.

Additional arguments to the function (besides the input tensor) are
supported. The returned result is a dense tensor.

Args:
    func: The unary operator to wrap.
Returns:
    Wrapped function that supports `JAXSparse` tensors.


def forward(self, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.Tensor]=None, position_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Cache]=None, inputs_embeds: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, cache_position: Optional[torch.LongTensor]=None, logits_to_keep: Union[int, torch.Tensor]=0, **kwargs: Unpack[KwargsForCausalLM]) -> CausalLMOutputWithPast:
    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
    outputs: BaseModelOutputWithPast = self.model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, cache_position=cache_position, **kwargs)
    hidden_states = outputs.last_hidden_state
    slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
    logits = self.lm_head(hidden_states[:, slice_indices, :])
    loss = None
    if labels is not None:
        loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
    return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions)

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

Example:

```python
>>> from transformers import AutoTokenizer, Starcoder2ForCausalLM

>>> model = Starcoder2ForCausalLM.from_pretrained("meta-starcoder2/Starcoder2-2-7b-hf")
>>> tokenizer = AutoTokenizer.from_pretrained("meta-starcoder2/Starcoder2-2-7b-hf")

>>> prompt = "Hey, are you conscious? Can you talk to me?"
>>> inputs = tokenizer(prompt, return_tensors="pt")

>>> # Generate
>>> generate_ids = model.generate(inputs.input_ids, max_length=30)
>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
"Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
```


def abs(x, name=None):
    with ops.name_scope(name, 'Abs', [x]) as name:
        x = ops.convert_to_tensor(x, name='x')
        if x.dtype.is_complex:
            return gen_math_ops.complex_abs(x, Tout=x.dtype.real_dtype, name=name)
        return gen_math_ops._abs(x, name=name)

Computes the absolute value of a tensor.

Given a tensor of integer or floating-point values, this operation returns a
tensor of the same type, where each element contains the absolute value of the
corresponding element in the input.

Given a tensor `x` of complex numbers, this operation returns a tensor of type
`float32` or `float64` that is the absolute value of each element in `x`. For
a complex number \\(a + bj\\), its absolute value is computed as
\\(\sqrt{a^2 + b^2}\\).

For example:

>>> # real number
>>> x = tf.constant([-2.25, 3.25])
>>> tf.abs(x)
<tf.Tensor: shape=(2,), dtype=float32,
numpy=array([2.25, 3.25], dtype=float32)>

>>> # complex number
>>> x = tf.constant([[-2.25 + 4.75j], [-3.25 + 5.75j]])
>>> tf.abs(x)
<tf.Tensor: shape=(2, 1), dtype=float64, numpy=
array([[5.25594901],
       [6.60492241]])>

Args:
  x: A `Tensor` or `SparseTensor` of type `float16`, `float32`, `float64`,
    `int32`, `int64`, `complex64` or `complex128`.
  name: A name for the operation (optional).

Returns:
  A `Tensor` or `SparseTensor` of the same size, type and sparsity as `x`,
    with absolute values. Note, for `complex64` or `complex128` input, the
    returned `Tensor` will be of type `float32` or `float64`, respectively.


def GetControlSequenceIndicator(self):
    return self._csi

Returns the control sequence indicator string.

Returns:
  The control sequence indicator string or None if control sequences are not
  supported.


def device(self, name):
    if isinstance(name, LogicalDevice):
        name = name.name
    elif pydev.is_device_spec(name):
        name = name.to_string()
    return _EagerDeviceContext(self, name)

Context-manager to force placement of operations and Tensors on a device.

Args:
  name: Name of the device or None to get default placement.

Returns:
  Context manager that forces device placement.

Raises:
  ValueError: If name is not a string or is an invalid device name.
  RuntimeError: If device scopes are not properly nested.


class Conv3D(keras_layers.Conv3D, base.Layer):

    def __init__(self, filters, kernel_size, strides=(1, 1, 1), padding='valid', data_format='channels_last', dilation_rate=(1, 1, 1), activation=None, use_bias=True, kernel_initializer=None, bias_initializer=init_ops.zeros_initializer(), kernel_regularizer=None, bias_regularizer=None, activity_regularizer=None, kernel_constraint=None, bias_constraint=None, trainable=True, name=None, **kwargs):
        super(Conv3D, self).__init__(filters=filters, kernel_size=kernel_size, strides=strides, padding=padding, data_format=data_format, dilation_rate=dilation_rate, activation=activation, use_bias=use_bias, kernel_initializer=kernel_initializer, bias_initializer=bias_initializer, kernel_regularizer=kernel_regularizer, bias_regularizer=bias_regularizer, activity_regularizer=activity_regularizer, kernel_constraint=kernel_constraint, bias_constraint=bias_constraint, trainable=trainable, name=name, **kwargs)

3D convolution layer (e.g. spatial convolution over volumes).

This layer creates a convolution kernel that is convolved
(actually cross-correlated) with the layer input to produce a tensor of
outputs. If `use_bias` is True (and a `bias_initializer` is provided),
a bias vector is created and added to the outputs. Finally, if
`activation` is not `None`, it is applied to the outputs as well.

Args:
  filters: Integer, the dimensionality of the output space (i.e. the number
    of filters in the convolution).
  kernel_size: An integer or tuple/list of 3 integers, specifying the
    depth, height and width of the 3D convolution window.
    Can be a single integer to specify the same value for
    all spatial dimensions.
  strides: An integer or tuple/list of 3 integers,
    specifying the strides of the convolution along the depth,
    height and width.
    Can be a single integer to specify the same value for
    all spatial dimensions.
    Specifying any stride value != 1 is incompatible with specifying
    any `dilation_rate` value != 1.
  padding: One of `"valid"` or `"same"` (case-insensitive).
    `"valid"` means no padding. `"same"` results in padding evenly to
    the left/right or up/down of the input such that output has the same
    height/width dimension as the input.
  data_format: A string, one of `channels_last` (default) or `channels_first`.
    The ordering of the dimensions in the inputs.
    `channels_last` corresponds to inputs with shape
    `(batch, depth, height, width, channels)` while `channels_first`
    corresponds to inputs with shape
    `(batch, channels, depth, height, width)`.
  dilation_rate: An integer or tuple/list of 3 integers, specifying
    the dilation rate to use for dilated convolution.
    Can be a single integer to specify the same value for
    all spatial dimensions.
    Currently, specifying any `dilation_rate` value != 1 is
    incompatible with specifying any stride value != 1.
  activation: Activation function. Set it to None to maintain a
    linear activation.
  use_bias: Boolean, whether the layer uses a bias.
  kernel_initializer: An initializer for the convolution kernel.
  bias_initializer: An initializer for the bias vector. If None, the default
    initializer will be used.
  kernel_regularizer: Optional regularizer for the convolution kernel.
  bias_regularizer: Optional regularizer for the bias vector.
  activity_regularizer: Optional regularizer function for the output.
  kernel_constraint: Optional projection function to be applied to the
      kernel after being updated by an `Optimizer` (e.g. used to implement
      norm constraints or value constraints for layer weights). The function
      must take as input the unprojected variable and must return the
      projected variable (which must have the same shape). Constraints are
      not safe to use when doing asynchronous distributed training.
  bias_constraint: Optional projection function to be applied to the
      bias after being updated by an `Optimizer`.
  trainable: Boolean, if `True` also add variables to the graph collection
    `GraphKeys.TRAINABLE_VARIABLES` (see `tf.Variable`).
  name: A string, the name of the layer.


def trace(src, options=None):
    options = options or config.Options.create()
    with config.verbosity_from(options):
        loader = load_pytd.create_loader(options)
        ret = analyze.infer_types(src=src, options=options, loader=loader)
        pytd_module = ret.ast
        raw_traces = []
        for op, symbol, data in ret.context.vm.opcode_traces:
            raw_traces.append((op, symbol, tuple((_to_pytd(d, loader, pytd_module) for d in data))))
    return source.Code(src, raw_traces, TypeTrace, options.input)

Generates type traces for the given source code.

Args:
  src: The source text.
  options: A pytype.config.Options object that can be used to specify options
    such as the target Python version.

Returns:
  A source.Code object.


class _Common:
    cache: dict[int, _MutableProxyImpl] = dataclasses.field(default_factory=dict)
    resolved: dict[int, Any] = dataclasses.field(default_factory=dict)
    is_frozen: bool = False

    def get_proxy(self, value: Any) -> _MutableProxyImpl:
        """Returns the proxy associated with the given value, or create it."""
        id_ = id(value)
        if id_ not in self.cache:
            self.cache[id_] = _MutableProxyImpl(obj=value, common=self)
        return self.cache[id_]

Shared variable across all nested childs of an `unfrozen()` object.

Attributes:
  cache: Global mapping `id(_MutableProxyImpl) -> _MutableProxyImpl` to avoid
    duplicating the same object proxy

    ```python
    a = a.unfrozen()
    a.x = x
    a.y = x  # `a.x` and `a.y` point to the same object
    a = a.frozen()
    assert a.x is a.y
    ```

  resolved: Cache of the objects after they have been frozen (to avoid
    evaluating the object twice.
  is_frozen: Become `True` after `.frozen()` is called. After which all
    Mutable are invalid


def _column_name_with_class_name(fc):
    return fc.__class__.__name__ + ':' + fc.name

Returns a unique name for the feature column used during deduping.

Without this two FeatureColumns that have the same name and where
one wraps the other, such as an IndicatorColumn wrapping a
SequenceCategoricalColumn, will fail to deserialize because they will have the
same name in columns_by_name, causing the wrong column to be returned.

Args:
  fc: A FeatureColumn.

Returns:
  A unique name as a string.


def to_json_string(self) -> str:
    dictionary = self.to_dict()
    return json.dumps(dictionary, indent=2, sort_keys=True) + '\n'

Serializes this instance to a JSON string.

Returns:
    `str`: String containing all the attributes that make up this feature_extractor instance in JSON format.


def _replace_row_partitions(value, new_partitions):
    if isinstance(value, tensor.Tensor) or not new_partitions:
        return value
    elif isinstance(value, ragged_tensor.RaggedTensor):
        return ragged_tensor.RaggedTensor._from_row_partition(values=_replace_row_partitions(value.values, new_partitions[1:]), row_partition=new_partitions[0])
    else:
        assert isinstance(value, StructuredTensor)
        new_fields = dict(((k, _replace_row_partitions(v, new_partitions)) for k, v in value._fields.items()))
        return StructuredTensor._old_init(fields=new_fields, shape=value.shape, nrows=value.nrows(), row_partitions=tuple(new_partitions) + tuple(value.row_partitions[len(new_partitions):]))

Updates `value` to use `new_partitions` as its (outer) row partitions.

This is used to ensure that all fields in a `StructuredTensor` use identical
`RowPartition` objects for the shared dimensions.  In particular,
`StructuredTensor.from_fields` first merges all of the row partitions from
any fields, and then replaces the outer row partitions of all fields with
the merged row partitions (using this function).

Args:
  value: A `Tensor`, `RaggedTensor`, or `StructuredTensor`.
  new_partitions: A list of row-partitions that should be used by `value`.
    Must be equivalent to `value`'s current row partitions.

Returns:
  A value that is equivalent to `value`, where outer row partitions have been
  replaced by `new_partitions`.


class TruePositives(_ConfusionMatrixConditionCount):

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

Calculates the number of true positives.

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

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

Args:
    thresholds: (Optional) Defaults to `0.5`. A float value, or a Python
        list/tuple of float threshold values in `[0, 1]`. A threshold is
        compared with prediction values to determine the truth value of
        predictions (i.e., above the threshold is `True`, below is `False`).
        If used with a loss function that sets `from_logits=True` (i.e. no
        sigmoid applied to predictions), `thresholds` should be set to 0.
        One metric value is generated for each threshold value.
    name: (Optional) string name of the metric instance.
    dtype: (Optional) data type of the metric result.

Example:

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

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


def __init__(self, num_groups=2):
    if num_groups < 1:
        raise ValueError(f'Argument `num_groups` must be a positive integer. Received: num_groups={num_groups}')
    self._ready = threading.Condition(threading.Lock())
    self._num_groups = num_groups
    self._group_member_counts = [0] * self._num_groups

Initialize a group lock.

Args:
  num_groups: The number of groups that will be accessing the resource under
    consideration. Should be a positive number.

Returns:
  A group lock that can then be used to synchronize code.

Raises:
  ValueError: If num_groups is less than 1.