def map_flat_values(op, *args, **kwargs):
    partition_lists = []
    flat_values_nrows = []
    inner_args = _replace_ragged_with_flat_values(args, partition_lists, flat_values_nrows)
    inner_kwargs = _replace_ragged_with_flat_values(kwargs, partition_lists, flat_values_nrows)
    if not partition_lists:
        return op(*args, **kwargs)
    if flat_values_nrows:
        flat_values_nrows = set(flat_values_nrows)
        if len(flat_values_nrows) != 1:
            raise ValueError("Input RaggedTensors' flat_values must all have the same outer-dimension size.  Got sizes: %s" % flat_values_nrows)
        flat_values_nrows = flat_values_nrows.pop()
    else:
        flat_values_nrows = None
    partition_dtypes = set((p[0].dtype for p in partition_lists))
    if len(partition_dtypes) > 1:
        if not ragged_config.auto_cast_partition_dtype():
            raise ValueError('Input RaggedTensors have mismatched row partition dtypes; use RaggedTensor.with_row_splits_dtype() to convert them to compatible dtypes.')
        partition_lists = [[p.with_dtype(dtypes.int64) for p in partition_list] for partition_list in partition_lists]
    op_output = op(*inner_args, **inner_kwargs)
    if flat_values_nrows is not None:
        if not op_output.shape[:1].is_compatible_with([flat_values_nrows]):
            raise ValueError('tf.ragged.map_flat_values requires that the output of `op` have the same outer-dimension size as flat_values of any ragged inputs. (output shape: %s; expected outer dimension size: %s)' % (op_output.shape, flat_values_nrows))
    return ragged_tensor.RaggedTensor._from_nested_row_partitions(op_output, _merge_partition_lists(partition_lists), validate=False)

Applies `op` to the `flat_values` of one or more RaggedTensors.

Replaces any `RaggedTensor` in `args` or `kwargs` with its `flat_values`
tensor (which collapses all ragged dimensions), and then calls `op`.  Returns
a `RaggedTensor` that is constructed from the input `RaggedTensor`s'
`nested_row_splits` and the value returned by the `op`.

If the input arguments contain multiple `RaggedTensor`s, then they must have
identical `nested_row_splits`.

This operation is generally used to apply elementwise operations to each value
in a `RaggedTensor`.

Warning: `tf.ragged.map_flat_values` does *not* apply `op` to each row of a
ragged tensor.  This difference is important for non-elementwise operations,
such as `tf.reduce_sum`.  If you wish to apply a non-elementwise operation to
each row of a ragged tensor, use `tf.map_fn` instead.  (You may need to
specify an `output_signature` when using `tf.map_fn` with ragged tensors.)

Examples:

>>> rt = tf.ragged.constant([[1, 2, 3], [], [4, 5], [6]])
>>> tf.ragged.map_flat_values(tf.ones_like, rt)
<tf.RaggedTensor [[1, 1, 1], [], [1, 1], [1]]>
>>> tf.ragged.map_flat_values(tf.multiply, rt, rt)
<tf.RaggedTensor [[1, 4, 9], [], [16, 25], [36]]>
>>> tf.ragged.map_flat_values(tf.add, rt, 5)
<tf.RaggedTensor [[6, 7, 8], [], [9, 10], [11]]>

Example with a non-elementwise operation (note that `map_flat_values` and
`map_fn` return different results):

>>> rt = tf.ragged.constant([[1.0, 3.0], [], [3.0, 6.0, 3.0]])
>>> def normalized(x):
...   return x / tf.reduce_sum(x)
>>> tf.ragged.map_flat_values(normalized, rt)
<tf.RaggedTensor [[0.0625, 0.1875], [], [0.1875, 0.375, 0.1875]]>
>>> tf.map_fn(normalized, rt)
<tf.RaggedTensor [[0.25, 0.75], [], [0.25, 0.5, 0.25]]>

Args:
  op: The operation that should be applied to the RaggedTensor `flat_values`.
    `op` is typically an element-wise operation (such as math_ops.add), but
    any operation that preserves the size of the outermost dimension can be
    used.  I.e., `shape[0]` of the value returned by `op` must match
    `shape[0]` of the `RaggedTensor`s' `flat_values` tensors.
  *args: Arguments for `op`.
  **kwargs: Keyword arguments for `op`.

Returns:
  A `RaggedTensor` whose `ragged_rank` matches the `ragged_rank` of all
  input `RaggedTensor`s.
Raises:
  ValueError: If args contains no `RaggedTensors`, or if the `nested_splits`
    of the input `RaggedTensor`s are not identical.


def do_import(self, keys, values, name=None):
    with tf.name_scope(name or '%s_lookup_table_import' % self._name):
        op = gen_simple_hash_table_op.examples_simple_hash_table_import(self.resource_handle, keys, values)
        return op

Import all `key` and `value` pairs.

(Note that "import" is a python reserved word, so it cannot be the name of
a method.)

Args:
  keys: Tensor of all keys.
  values: Tensor of all values.
  name: A name for the operation (optional).

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


def _find_shape_dtype(fields: Mapping[str, _FieldValue], nrows: Optional[tensor.Tensor], row_partitions: Optional[Sequence[RowPartition]]) -> dtypes.DType:
    field_dtypes = [_field_shape_dtype(v) for v in fields.values()]
    nrows_dtypes = [nrows.dtype] if isinstance(nrows, tensor.Tensor) else []
    rp_dtypes = [] if row_partitions is None else [rp.dtype for rp in row_partitions]
    all_dtypes = field_dtypes + nrows_dtypes + rp_dtypes
    if dtypes.int64 in all_dtypes:
        return dtypes.int64
    if dtypes.int32 in all_dtypes:
        return dtypes.int32
    return dtypes.int64

Return a consistent dtype for fields, nrows, & row_partitions.

In the future, the default will switch from int64 to int32, but for now,
we stick with int64.

Args:
  fields: the fields of the StructuredTensor.
  nrows: the nrows of the StructuredTensor
  row_partitions: the row_partitions of the StructuredTensor.

Returns:
  If anything requires int64, then return int64.
  If int32 is explicitly specified, return int32. Otherwise, return int64.


class PrecisionAtRecall(SensitivitySpecificityBase):

    def __init__(self, recall, num_thresholds=200, class_id=None, name=None, dtype=None):
        if recall < 0 or recall > 1:
            raise ValueError(f'Argument `recall` must be in the range [0, 1]. Received: recall={recall}')
        self.recall = recall
        self.num_thresholds = num_thresholds
        super().__init__(value=recall, num_thresholds=num_thresholds, class_id=class_id, name=name, dtype=dtype)

    def result(self):
        recalls = ops.divide_no_nan(self.true_positives, ops.add(self.true_positives, self.false_negatives))
        precisions = ops.divide_no_nan(self.true_positives, ops.add(self.true_positives, self.false_positives))
        return self._find_max_under_constraint(recalls, precisions, ops.greater_equal)

    def get_config(self):
        config = {'num_thresholds': self.num_thresholds, 'recall': self.recall}
        base_config = super().get_config()
        return {**base_config, **config}

Computes best precision where recall is >= specified value.

This metric creates four local variables, `true_positives`,
`true_negatives`, `false_positives` and `false_negatives` that are used to
compute the precision at the given recall. The threshold for the given
recall value is computed and used to evaluate the corresponding precision.

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

If `class_id` is specified, we calculate precision by considering only the
entries in the batch for which `class_id` is above the threshold
predictions, and computing the fraction of them for which `class_id` is
indeed a correct label.

Args:
    recall: A scalar value in range `[0, 1]`.
    num_thresholds: (Optional) Defaults to 200. The number of thresholds to
        use for matching the given recall.
    class_id: (Optional) Integer class ID for which we want binary metrics.
        This must be in the half-open interval `[0, num_classes)`, where
        `num_classes` is the last dimension of predictions.
    name: (Optional) string name of the metric instance.
    dtype: (Optional) data type of the metric result.

Example:

>>> m = keras.metrics.PrecisionAtRecall(0.5)
>>> m.update_state([0, 0, 0, 1, 1], [0, 0.3, 0.8, 0.3, 0.8])
>>> m.result()
0.5

>>> m.reset_state()
>>> m.update_state([0, 0, 0, 1, 1], [0, 0.3, 0.8, 0.3, 0.8],
...                sample_weight=[2, 2, 2, 1, 1])
>>> m.result()
0.33333333

Usage with `compile()` API:

```python
model.compile(
    optimizer='sgd',
    loss='binary_crossentropy',
    metrics=[keras.metrics.PrecisionAtRecall(recall=0.8)])
```


def get_extended_attention_mask(self, attention_mask: torch.Tensor, input_shape: Tuple[int], device: torch.device, has_query: bool=False) -> torch.Tensor:
    if attention_mask.dim() == 3:
        extended_attention_mask = attention_mask[:, None, :, :]
    elif attention_mask.dim() == 2:
        extended_attention_mask = attention_mask[:, None, None, :]
    else:
        raise ValueError(f'Wrong shape for input_ids (shape {input_shape}) or attention_mask (shape {attention_mask.shape})')
    extended_attention_mask = extended_attention_mask.to(dtype=self.dtype)
    extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
    return extended_attention_mask

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

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

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


class DonutSwinModelOutput(ModelOutput):
    last_hidden_state: Optional[torch.FloatTensor] = None
    pooler_output: Optional[torch.FloatTensor] = None
    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
    reshaped_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None

DonutSwin model's outputs that also contains a pooling of the last hidden states.

Args:
    last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
        Sequence of hidden-states at the output of the last layer of the model.
    pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`, *optional*, returned when `add_pooling_layer=True` is passed):
        Average pooling of the last layer hidden-state.
    hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) 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(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
        Tuple of `torch.FloatTensor` (one for each stage) 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.
    reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
        shape `(batch_size, hidden_size, height, width)`.

        Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
        include the spatial dimensions.


def spatial_2d_padding(x, padding=((1, 1), (1, 1)), data_format=None):
    assert len(padding) == 2
    assert len(padding[0]) == 2
    assert len(padding[1]) == 2
    if data_format is None:
        data_format = image_data_format()
    if data_format not in {'channels_first', 'channels_last'}:
        raise ValueError('Unknown data_format: ' + str(data_format))
    if data_format == 'channels_first':
        pattern = [[0, 0], [0, 0], list(padding[0]), list(padding[1])]
    else:
        pattern = [[0, 0], list(padding[0]), list(padding[1]), [0, 0]]
    return array_ops.pad(x, pattern)

Pads the 2nd and 3rd dimensions of a 4D tensor.

Args:
    x: Tensor or variable.
    padding: Tuple of 2 tuples, padding pattern.
    data_format: One of `channels_last` or `channels_first`.

Returns:
    A padded 4D tensor.

Raises:
    ValueError: if `data_format` is neither
        `channels_last` or `channels_first`.


class OmDetTurboConfig(PretrainedConfig):
    model_type = 'omdet-turbo'
    attribute_map = {'encoder_hidden_dim': 'd_model', 'num_attention_heads': 'encoder_attention_heads'}

    def __init__(self, text_config=None, backbone_config=None, use_timm_backbone=True, backbone='swin_tiny_patch4_window7_224', backbone_kwargs=None, use_pretrained_backbone=False, apply_layernorm_after_vision_backbone=True, image_size=640, disable_custom_kernels=False, layer_norm_eps=1e-05, batch_norm_eps=1e-05, init_std=0.02, text_projection_in_dim=512, text_projection_out_dim=512, task_encoder_hidden_dim=1024, class_embed_dim=512, class_distance_type='cosine', num_queries=900, csp_activation='silu', conv_norm_activation='gelu', encoder_feedforward_activation='relu', encoder_feedforward_dropout=0.0, encoder_dropout=0.0, hidden_expansion=1, vision_features_channels=[256, 256, 256], encoder_hidden_dim=256, encoder_in_channels=[192, 384, 768], encoder_projection_indices=[2], encoder_attention_heads=8, encoder_dim_feedforward=2048, encoder_layers=1, positional_encoding_temperature=10000, num_feature_levels=3, decoder_hidden_dim=256, decoder_num_heads=8, decoder_num_layers=6, decoder_activation='relu', decoder_dim_feedforward=2048, decoder_num_points=4, decoder_dropout=0.0, eval_size=None, learn_initial_query=False, cache_size=100, is_encoder_decoder=True, **kwargs):
        if use_timm_backbone:
            if backbone_config is None:
                backbone_kwargs = {'out_indices': [1, 2, 3], 'img_size': image_size, 'always_partition': True}
        elif backbone_config is None:
            logger.info('`backbone_config` is `None`. Initializing the config with the default `swin` vision config.')
            backbone_config = CONFIG_MAPPING['swin'](window_size=7, image_size=image_size, embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], out_indices=[2, 3, 4])
        elif isinstance(backbone_config, dict):
            backbone_model_type = backbone_config.get('model_type')
            config_class = CONFIG_MAPPING[backbone_model_type]
            backbone_config = config_class.from_dict(backbone_config)
        verify_backbone_config_arguments(use_timm_backbone=use_timm_backbone, use_pretrained_backbone=use_pretrained_backbone, backbone=backbone, backbone_config=backbone_config, backbone_kwargs=backbone_kwargs)
        if text_config is None:
            logger.info('`text_config` is `None`. Initializing the config with the default `clip_text_model` text config.')
            text_config = CONFIG_MAPPING['clip_text_model']()
        elif isinstance(text_config, dict):
            text_model_type = text_config.get('model_type')
            text_config = CONFIG_MAPPING[text_model_type](**text_config)
        if class_distance_type not in ['cosine', 'dot']:
            raise ValueError(f'Invalid `class_distance_type`. It should be either `cosine` or `dot`, but got {class_distance_type}.')
        self.text_config = text_config
        self.backbone_config = backbone_config
        self.use_timm_backbone = use_timm_backbone
        self.backbone = backbone
        self.backbone_kwargs = backbone_kwargs
        self.use_pretrained_backbone = use_pretrained_backbone
        self.apply_layernorm_after_vision_backbone = apply_layernorm_after_vision_backbone
        self.image_size = image_size
        self.disable_custom_kernels = disable_custom_kernels
        self.layer_norm_eps = layer_norm_eps
        self.batch_norm_eps = batch_norm_eps
        self.init_std = init_std
        self.text_projection_in_dim = text_projection_in_dim
        self.text_projection_out_dim = text_projection_out_dim
        self.task_encoder_hidden_dim = task_encoder_hidden_dim
        self.class_embed_dim = class_embed_dim
        self.class_distance_type = class_distance_type
        self.num_queries = num_queries
        self.csp_activation = csp_activation
        self.conv_norm_activation = conv_norm_activation
        self.encoder_feedforward_activation = encoder_feedforward_activation
        self.encoder_feedforward_dropout = encoder_feedforward_dropout
        self.encoder_dropout = encoder_dropout
        self.hidden_expansion = hidden_expansion
        self.vision_features_channels = vision_features_channels
        self.encoder_hidden_dim = encoder_hidden_dim
        self.encoder_in_channels = encoder_in_channels
        self.encoder_projection_indices = encoder_projection_indices
        self.encoder_attention_heads = encoder_attention_heads
        self.encoder_dim_feedforward = encoder_dim_feedforward
        self.encoder_layers = encoder_layers
        self.positional_encoding_temperature = positional_encoding_temperature
        self.num_feature_levels = num_feature_levels
        self.decoder_hidden_dim = decoder_hidden_dim
        self.decoder_num_heads = decoder_num_heads
        self.decoder_num_layers = decoder_num_layers
        self.decoder_activation = decoder_activation
        self.decoder_dim_feedforward = decoder_dim_feedforward
        self.decoder_num_points = decoder_num_points
        self.decoder_dropout = decoder_dropout
        self.eval_size = eval_size
        self.learn_initial_query = learn_initial_query
        self.cache_size = cache_size
        self.is_encoder_decoder = is_encoder_decoder
        super().__init__(is_encoder_decoder=is_encoder_decoder, **kwargs)

This is the configuration class to store the configuration of a [`OmDetTurboForObjectDetection`].
It is used to instantiate a OmDet-Turbo model according to the specified arguments, defining the model architecture
Instantiating a configuration with the defaults will yield a similar configuration to that of the OmDet-Turbo
[omlab/omdet-turbo-swin-tiny-hf](https://huggingface.co/omlab/omdet-turbo-swin-tiny-hf) architecture.

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

Args:
    text_config (`PretrainedConfig`, *optional*):
        The configuration of the text backbone.
    backbone_config (`PretrainedConfig`, *optional*):
        The configuration of the vision backbone.
    use_timm_backbone (`bool`, *optional*, defaults to `True`):
        Whether to use the timm for the vision backbone.
    backbone (`str`, *optional*, defaults to `"swin_tiny_patch4_window7_224"`):
        The name of the pretrained vision backbone to use. If `use_pretrained_backbone=False` a randomly initialized
        backbone with the same architecture `backbone` is used.
    backbone_kwargs (`dict`, *optional*):
        Additional kwargs for the vision backbone.
    use_pretrained_backbone (`bool`, *optional*, defaults to `False`):
        Whether to use a pretrained vision backbone.
    apply_layernorm_after_vision_backbone (`bool`, *optional*, defaults to `True`):
        Whether to apply layer normalization on the feature maps of the vision backbone output.
    image_size (`int`, *optional*, defaults to 640):
        The size (resolution) of each image.
    disable_custom_kernels (`bool`, *optional*, defaults to `False`):
        Whether to disable custom kernels.
    layer_norm_eps (`float`, *optional*, defaults to 1e-05):
        The epsilon value for layer normalization.
    batch_norm_eps (`float`, *optional*, defaults to 1e-05):
        The epsilon value for batch normalization.
    init_std (`float`, *optional*, defaults to 0.02):
        The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
    text_projection_in_dim (`int`, *optional*, defaults to 512):
        The input dimension for the text projection.
    text_projection_out_dim (`int`, *optional*, defaults to 512):
        The output dimension for the text projection.
    task_encoder_hidden_dim (`int`, *optional*, defaults to 1024):
        The feedforward dimension for the task encoder.
    class_embed_dim (`int`, *optional*, defaults to 512):
        The dimension of the classes embeddings.
    class_distance_type (`str`, *optional*, defaults to `"cosine"`):
        The type of of distance to compare predicted classes to projected classes embeddings.
        Can be `"cosine"` or `"dot"`.
    num_queries (`int`, *optional*, defaults to 900):
        The number of queries.
    csp_activation (`str`, *optional*, defaults to `"silu"`):
        The activation function of the Cross Stage Partial (CSP) networks of the encoder.
    conv_norm_activation (`str`, *optional*, defaults to `"gelu"`):
        The activation function of the ConvNormLayer layers of the encoder.
    encoder_feedforward_activation (`str`, *optional*, defaults to `"relu"`):
        The activation function for the feedforward network of the encoder.
    encoder_feedforward_dropout (`float`, *optional*, defaults to 0.0):
        The dropout rate following the activation of the encoder feedforward network.
    encoder_dropout (`float`, *optional*, defaults to 0.0):
        The dropout rate of the encoder multi-head attention module.
    hidden_expansion (`int`, *optional*, defaults to 1):
        The hidden expansion of the CSP networks in the encoder.
    vision_features_channels (`tuple(int)`, *optional*, defaults to `[256, 256, 256]`):
        The projected vision features channels used as inputs for the decoder.
    encoder_hidden_dim (`int`, *optional*, defaults to 256):
        The hidden dimension of the encoder.
    encoder_in_channels (`List(int)`, *optional*, defaults to `[192, 384, 768]`):
        The input channels for the encoder.
    encoder_projection_indices (`List(int)`, *optional*, defaults to `[2]`):
        The indices of the input features projected by each layers.
    encoder_attention_heads (`int`, *optional*, defaults to 8):
        The number of attention heads for the encoder.
    encoder_dim_feedforward (`int`, *optional*, defaults to 2048):
        The feedforward dimension for the encoder.
    encoder_layers (`int`, *optional*, defaults to 1):
        The number of layers in the encoder.
    positional_encoding_temperature (`int`, *optional*, defaults to 10000):
        The positional encoding temperature in the encoder.
    num_feature_levels (`int`, *optional*, defaults to 3):
        The number of feature levels for the multi-scale deformable attention module of the decoder.
    decoder_hidden_dim (`int`, *optional*, defaults to 256):
        The hidden dimension of the decoder.
    decoder_num_heads (`int`, *optional*, defaults to 8):
        The number of heads for the decoder.
    decoder_num_layers (`int`, *optional*, defaults to 6):
        The number of layers for the decoder.
    decoder_activation (`str`, *optional*, defaults to `"relu"`):
        The activation function for the decoder.
    decoder_dim_feedforward (`int`, *optional*, defaults to 2048):
        The feedforward dimension for the decoder.
    decoder_num_points (`int`, *optional*, defaults to 4):
        The number of points sampled in the decoder multi-scale deformable attention module.
    decoder_dropout (`float`, *optional*, defaults to 0.0):
        The dropout rate for the decoder.
    eval_size (`Tuple[int, int]`, *optional*):
        Height and width used to computes the effective height and width of the position embeddings after taking
        into account the stride (see RTDetr).
    learn_initial_query (`bool`, *optional*, defaults to `False`):
        Whether to learn the initial query.
    cache_size (`int`, *optional*, defaults to 100):
        The cache size for the classes and prompts caches.
    is_encoder_decoder (`bool`, *optional*, defaults to `True`):
        Whether the model is used as an encoder-decoder model or not.
    kwargs (`Dict[str, Any]`, *optional*):
        Additional parameters from the architecture. The values in kwargs will be saved as part of the configuration
        and can be used to control the model outputs.

Examples:

```python
>>> from transformers import OmDetTurboConfig, OmDetTurboForObjectDetection

>>> # Initializing a OmDet-Turbo omlab/omdet-turbo-swin-tiny-hf style configuration
>>> configuration = OmDetTurboConfig()

>>> # Initializing a model (with random weights) from the omlab/omdet-turbo-swin-tiny-hf style configuration
>>> model = OmDetTurboForObjectDetection(configuration)

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


def most_specific_common_supertype(self, others: Sequence[trace.TraceType]) -> Optional['TypeSpec']:
    if any((type(self) is not type(other) for other in others)):
        return None
    has_supertype = True

    def make_supertype_attribute(attribute_self, *attribute_others):
        nonlocal has_supertype
        if not has_supertype:
            return
        if isinstance(attribute_self, trace.TraceType):
            attribute_supertype = attribute_self.most_specific_common_supertype(attribute_others)
            if attribute_supertype is None:
                has_supertype = False
                return
            return attribute_supertype
        else:
            if not all((attribute_self == attribute_other for attribute_other in attribute_others)):
                has_supertype = False
                return
            return attribute_self
    try:
        serialized_supertype = nest.map_structure(make_supertype_attribute, self._serialize(), *(o._serialize() for o in others))
    except (ValueError, TypeError):
        return None
    return self._deserialize(serialized_supertype) if has_supertype else None

Returns the most specific supertype TypeSpec  of `self` and `others`.

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

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

Args:
  others: A sequence of TraceTypes.


def softplus(features, name=None):
    return gen_nn_ops.softplus(features, name)

Computes elementwise softplus: `softplus(x) = log(exp(x) + 1)`.

`softplus` is a smooth approximation of `relu`. Like `relu`, `softplus` always
takes on positive values.

<img style="width:100%" src="https://www.tensorflow.org/images/softplus.png">

Example:

>>> import tensorflow as tf
>>> tf.math.softplus(tf.range(0, 2, dtype=tf.float32)).numpy()
array([0.6931472, 1.3132616], dtype=float32)

Args:
  features: `Tensor`
  name: Optional: name to associate with this operation.
Returns:
  `Tensor`


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

Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
adding special tokens. A XLMProphetNet sequence has the following format:

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

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

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


def is_dtensor(tensor) -> bool:
    return _dtensor_device().is_dtensor(tensor)

Check whether the input tensor is a DTensor.

In Python, a DTensor has the same type as a `tf.Tensor`. This method will
let you check and handle the tensor differently if a tf.Tensor is a DTensor.

Args:
  tensor: an object to be checked.

Returns:
  bool, True if the given tensor is a DTensor.


def trainable_variables(self):
    return tuple((v for v in self.variables if v.trainable))

A sequence of trainable variables accessed by this FuncGraph.

Note that functions keep only weak references to variables. Calling the
function after a variable it accesses has been deleted is an error.

Returns:
  Sequence of trainable variables for this func graph.


def __init__(self, name: Union[str, bytes], bound_context: context.Context, function_type: function_type_lib.FunctionType, children: Optional[List['AtomicFunction']]=None, call_options: CallOptions=CallOptions(), cached_graph: Optional[func_graph_module.FuncGraph]=None):
    self._name = compat.as_bytes(name)
    self._bound_context = bound_context
    self._function_type = function_type
    self._children = children if children else []
    self._call_options = call_options
    self._cached_definition = None
    self._cached_graph = cached_graph
    self._generated_graph = None
    ref_key = (self._bound_context.function_scope_id, self.name)
    if ref_key not in RUNTIME_FUNCTION_REFS:
        RUNTIME_FUNCTION_REFS[ref_key] = 1
    else:
        RUNTIME_FUNCTION_REFS[ref_key] += 1

Construct a new AtomicFunction.

Args:
  name: str/bytes name of the runtime function in the bound context.
  bound_context: interface to the runtime for the AtomicFunction.
  function_type: input/output contract for the AtomicFunction
  children: list of AtomicFunctions that are needed to call this one.
  call_options: extra configuration options for the call.
  cached_graph: FuncGraph that this AtomicFunction was generated from (if
    known). Otherwise it will lazily construct a new corresponding FuncGraph
    if ever needed.


def __init__(self, host: str, port: int, command: Optional[str]=None, batch_size: int=100, embedded_columns: list=[]):
    self._host = host
    self._port = port
    self._command = command
    self._batch_size = batch_size
    self.embedded_columns = embedded_columns

Args:
host (str): The redis host
port (int): The redis port
command (str): command to be executed with redis client
batch_size (int): Number of key, values pairs to write at once
embedded_columns (list): list of column whose embedding needs to be generated

Returns:
:class:`~apache_beam.transforms.ptransform.PTransform`


def _get_weighted_mean_squared_error(self, quant_min, quant_max) -> tuple[float, float, float]:
    dequantized_hist_mids = self._get_dequantized_hist_mids_after_quantize(quant_min, quant_max)
    squared_error = (self._hist_mids - dequantized_hist_mids) ** 2
    weighted_error = np.sum(squared_error * self._hist_freq)
    return (weighted_error, quant_min, quant_max)

Gets mean squared error between hist_mids and dequantized hist_mids.

Quantization converts the range of numbers from [quant_min, quant_max] to
[0, 2^num_bits - 1]. Values less than quant_min are converted to 0, and
values greater than quant_max are converted to 2^num_bits - 1.

Args:
  quant_min: The minimum real value that can be represented by a quantized
    value.
  quant_max: The maximum real value that can be represented by a quantized
    value.

Returns:
  (error, quant_min, quant_max): Tuple of weighted mean squared error.
  error = (hist_mids - dequantized_hist_mids)**2 * hist_freq


def from_config(cls, config):
    return cls(**config)

Creates a quantizer from its config.

This method is the reverse of `get_config`,
capable of instantiating the same quantizer from the config
dictionary.

This method is used by Keras `model_to_estimator`, saving and
loading models to HDF5 formats, Keras model cloning, some visualization
utilities, and exporting models to and from JSON.

Args:
    config: A Python dictionary, typically the output of get_config.

Returns:
    A quantizer instance.


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.


def _get_num_multimodal_tokens(self, image_sizes=None, **kwargs):
    vision_data = {}
    if image_sizes is not None:
        num_image_tokens = [self.image_seq_length] * len(image_sizes)
        num_image_patches = [1] * len(image_sizes)
        vision_data.update({'num_image_tokens': num_image_tokens, 'num_image_patches': num_image_patches})
    return MultiModalData(**vision_data)

Computes the number of placeholder tokens needed for multimodal inputs with the given sizes.

Args:
    image_sizes (List[List[str]], *optional*):
        The input sizes formatted as (height, width) per each image.
Returns:
    Dict[str, List[int]]: A dictionary mapping each modality ("image", "video", "audio")
    to a list containing the number of placeholder tokens required. If the model doesn't accept
    a certain modality or no input sizes are provided, the dict value is set to an empty list.


def sobol_sample(dim: types.IntTensor, num_results: types.IntTensor, sequence_indices: types.IntTensor=None, digital_shift: types.IntTensor=None, scrambling_matrices: types.IntTensor=None, apply_tent_transform: bool=False, validate_args: bool=False, dtype: tf.DType=None, name: str=None) -> types.RealTensor:
    with tf.name_scope(name or 'sobol_sample'):
        dtype = dtype or tf.float32
        num_digits = tf.cast(tf.math.ceil(utils.log2(tf.cast(num_results, dtype=tf.float32))), tf.int32)
        generating_matrices = sobol_generating_matrices(dim, num_results, num_digits, validate_args=validate_args, dtype=tf.int32)
        if scrambling_matrices is not None:
            generating_matrices = digital_net.scramble_generating_matrices(generating_matrices, scrambling_matrices, num_digits, validate_args=validate_args)
        return digital_net.digital_net_sample(generating_matrices, num_results, num_digits, sequence_indices=sequence_indices, digital_shift=digital_shift, apply_tent_transform=apply_tent_transform, validate_args=validate_args, dtype=dtype)

Samples points from the Sobol sequence.

#### Examples

```python
import tf_quant_finance as tff

# Example: Sampling 1,000 points from the 2D Sobol sequence.

dim = 2
num_results = 1000

tff.math.qmc.sobol_sample(dim, num_results)
# ==> tf.Tensor([
#             [0.,         0.        ],
#             [0.5,        0.5       ],
#             [0.25,       0.75      ],
#             ...
#             [0.65527344, 0.9736328 ],
#             [0.40527344, 0.7236328 ],
#             [0.90527344, 0.22363281],
#         ], shape=(1000, 2), dtype=float32)
```

Args:
  dim: Positive scalar `Tensor` of integers with rank 0. The event size of the
    sampled points.
  num_results: Positive scalar `Tensor` of integers with rank 0. The number of
    points to sample.
  sequence_indices: Optional positive scalar `Tensor` of integers with rank 1.
    The elements of the sequence to return specified by their position in the
    sequence.
    Default value: `None` which corresponds to the `[0, num_results)` range.
  digital_shift: Optional digital shift to be applied to all the points via a
    bitwise xor.
    Default value: `None`.
  scrambling_matrices: Positive scalar `Tensor` with the same `shape` and
    `dtype` as `generating_matrices`. Used to randomize `generating_matrices`.
    Default value: `None`.
  apply_tent_transform: Python `bool` indicating whether to apply a tent
    transform to the sampled points.
    Default value: `False`.
  validate_args: Python `bool` indicating whether to validate arguments.
    Default value: `False`.
  dtype: Optional `dtype`. The `dtype` of the output `Tensor` (either
    `float32` or `float64`).
    Default value: `None` which maps to `float32`.
  name: Python `str` name prefixed to ops created by this function.
    Default value: `None` which maps to `sobol_sample`.

Returns:
  A `Tensor` of samples from  the Sobol sequence with `shape`
  `(num_samples, dim)` where `num_samples = min(num_results,
  size(sequence_indices))` and `dim = tf.shape(generating_matrices)[0]`.


class Cohere2Config(PretrainedConfig):
    model_type = 'cohere2'
    keys_to_ignore_at_inference = ['past_key_values']
    base_model_tp_plan = {'layers.*.self_attn.q_proj': 'colwise', 'layers.*.self_attn.k_proj': 'colwise', 'layers.*.self_attn.v_proj': 'colwise', 'layers.*.self_attn.o_proj': 'rowwise', 'layers.*.mlp.gate_proj': 'colwise', 'layers.*.mlp.up_proj': 'colwise', 'layers.*.mlp.down_proj': 'rowwise'}
    base_model_pp_plan = {'embed_tokens': (['input_ids'], ['inputs_embeds']), 'layers': (['hidden_states', 'attention_mask'], ['hidden_states']), 'norm': (['hidden_states'], ['hidden_states'])}

    def __init__(self, vocab_size=256000, hidden_size=8192, intermediate_size=22528, logit_scale=0.0625, num_hidden_layers=40, num_attention_heads=64, num_key_value_heads=None, hidden_act='silu', max_position_embeddings=8192, initializer_range=0.02, layer_norm_eps=1e-05, use_cache=True, pad_token_id=0, bos_token_id=5, eos_token_id=255001, tie_word_embeddings=True, rope_theta=10000.0, rope_scaling=None, attention_bias=False, attention_dropout=0.0, sliding_window=4096, layer_types=None, **kwargs):
        self.vocab_size = vocab_size
        self.max_position_embeddings = max_position_embeddings
        self.hidden_size = hidden_size
        self.logit_scale = logit_scale
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        if num_key_value_heads is None:
            num_key_value_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.hidden_act = hidden_act
        self.initializer_range = initializer_range
        self.layer_norm_eps = layer_norm_eps
        self.use_cache = use_cache
        self.rope_theta = rope_theta
        self.rope_scaling = rope_scaling
        self.attention_bias = attention_bias
        self.attention_dropout = attention_dropout
        self.sliding_window = sliding_window
        self.layer_types = layer_types
        self.head_dim = hidden_size // num_attention_heads
        rope_config_validation(self)
        super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs)
        if self.layer_types is None:
            sliding_window_pattern = getattr(self, 'sliding_window_pattern', 4)
            self.layer_types = ['sliding_attention' if bool((i + 1) % sliding_window_pattern) else 'full_attention' for i in range(self.num_hidden_layers)]
        layer_type_validation(self.layer_types)

This is the configuration class to store the configuration of a [`CohereModel`]. It is used to instantiate an Cohere
model according to the specified arguments, defining the model architecture.

Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
documentation from [`PretrainedConfig`] for more information. Instantiating a configuration
with the defaults will yield a similar configuration to that of the [CohereForAI/c4ai-command-r-v01](https://huggingface.co/CohereForAI/c4ai-command-r-v01) model.

Args:
    vocab_size (`int`, *optional*, defaults to 256000):
        Vocabulary size of the Cohere model. Defines the number of different tokens that can be represented by the
        `inputs_ids` passed when calling [`CohereModel`]
    hidden_size (`int`, *optional*, defaults to 8192):
        Dimension of the hidden representations.
    intermediate_size (`int`, *optional*, defaults to 22528):
        Dimension of the MLP representations.
    logit_scale (`float`, *optional*, defaults to 0.0625):
        The scaling factor for the output logits.
    num_hidden_layers (`int`, *optional*, defaults to 40):
        Number of hidden layers in the Transformer decoder.
    num_attention_heads (`int`, *optional*, defaults to 64):
        Number of attention heads for each attention layer in the Transformer decoder.
    num_key_value_heads (`int`, *optional*):
        This is the number of key_value heads that should be used to implement Grouped Query Attention. If
        `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
        `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
        converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
        by meanpooling all the original heads within that group. For more details, check out [this
        paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to
        `num_attention_heads`.
    hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
        The non-linear activation function (function or string) in the decoder.
    max_position_embeddings (`int`, *optional*, defaults to 8192):
        The maximum sequence length that this model might ever be used with.
    initializer_range (`float`, *optional*, defaults to 0.02):
        The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
    layer_norm_eps (`float`, *optional*, defaults to 1e-05):
        The epsilon used by the layer normalization.
    use_cache (`bool`, *optional*, defaults to `True`):
        Whether or not the model should return the last key/values attentions (not used by all models). Only
        relevant if `config.is_decoder=True`.
    pad_token_id (`int`, *optional*, defaults to 0):
        Padding token id.
    bos_token_id (`int`, *optional*, defaults to 5):
        Beginning of stream token id.
    eos_token_id (`int`, *optional*, defaults to 255001):
        End of stream token id.
    tie_word_embeddings (`bool`, *optional*, defaults to `True`):
        Whether to tie weight embeddings
    rope_theta (`float`, *optional*, defaults to 10000.0):
        The base period of the RoPE embeddings.
    rope_scaling (`Dict`, *optional*):
        Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
        and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
        accordingly.
        Expected contents:
            `rope_type` (`str`):
                The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
                'llama3'], with 'default' being the original RoPE implementation.
            `factor` (`float`, *optional*):
                Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
                most scaling types, a `factor` of x will enable the model to handle sequences of length x *
                original maximum pre-trained length.
            `original_max_position_embeddings` (`int`, *optional*):
                Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
                pretraining.
            `attention_factor` (`float`, *optional*):
                Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
                computation. If unspecified, it defaults to value recommended by the implementation, using the
                `factor` field to infer the suggested value.
            `beta_fast` (`float`, *optional*):
                Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
                ramp function. If unspecified, it defaults to 32.
            `beta_slow` (`float`, *optional*):
                Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
                ramp function. If unspecified, it defaults to 1.
            `short_factor` (`List[float]`, *optional*):
                Only used with 'longrope'. The scaling factor to be applied to short contexts (<
                `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
                size divided by the number of attention heads divided by 2
            `long_factor` (`List[float]`, *optional*):
                Only used with 'longrope'. The scaling factor to be applied to long contexts (<
                `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
                size divided by the number of attention heads divided by 2
            `low_freq_factor` (`float`, *optional*):
                Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
            `high_freq_factor` (`float`, *optional*):
                Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
    attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
        Whether to use a bias in the query, key, value and output projection layers during self-attention.
    attention_dropout (`float`, *optional*, defaults to 0.0):
        The dropout ratio for the attention probabilities.
    sliding_window (`int`, *optional*, defaults to 4096):
        Size of the sliding window attention context.
    layer_types (`list`, *optional*):
        Attention pattern for each layer.

```python
>>> from transformers import Cohere2Model, Cohere2Config

>>> # Initializing a Cohere Nextmodel configuration
>>> configuration = Cohere2Config()

>>> # Initializing a model from the Cohere2 configuration
>>> model = Cohere2Model(configuration) # doctest: +SKIP

>>> # Accessing the model configuration
>>> configuration = model.config # doctest: +SKIP
```


def get_function_def(self, name):
    if is_oss:
        with c_api_util.tf_buffer() as buffer_:
            pywrap_tfe.TFE_ContextGetFunctionDef(self._handle, name, buffer_)
            proto_data = pywrap_tf_session.TF_GetBuffer(buffer_)
        function_def = function_pb2.FunctionDef()
        function_def.ParseFromString(proto_data)
    else:
        function_def = pywrap_tfe.TFE_ContextGetFunctionDefNoSerialization(self._handle, name)
    return function_def

Get a function definition from the context.

Args:
  name: function signature name.

Returns:
  The requested FunctionDef.

Raises:
  tf.errors.NotFoundError: if name is not the name of a registered function.


def location_variable_convolution(self, hidden_states: torch.FloatTensor, kernel: torch.FloatTensor, bias: torch.FloatTensor, dilation: int=1, hop_size: int=256):
    batch, _, in_length = hidden_states.shape
    batch, _, out_channels, kernel_size, kernel_length = kernel.shape
    if in_length != kernel_length * hop_size:
        raise ValueError(f'Dim 2 of `hidden_states` should be {kernel_length * hop_size}) but got {in_length}. Please check `hidden_states` or `kernel` and `hop_size` to make sure they are correct.')
    padding = dilation * int((kernel_size - 1) / 2)
    hidden_states = nn.functional.pad(hidden_states, (padding, padding), 'constant', 0)
    hidden_states = hidden_states.unfold(2, hop_size + 2 * padding, hop_size)
    if hop_size < dilation:
        hidden_states = nn.functional.pad(hidden_states, (0, dilation), 'constant', 0)
    hidden_states = hidden_states.unfold(3, dilation, dilation)
    hidden_states = hidden_states[:, :, :, :, :hop_size]
    hidden_states = hidden_states.transpose(3, 4)
    hidden_states = hidden_states.unfold(4, kernel_size, 1)
    output_hidden_states = torch.einsum('bildsk,biokl->bolsd', hidden_states, kernel)
    output_hidden_states = output_hidden_states.to(memory_format=torch.channels_last_3d)
    bias = bias.unsqueeze(-1).unsqueeze(-1).to(memory_format=torch.channels_last_3d)
    output_hidden_states = output_hidden_states + bias
    output_hidden_states = output_hidden_states.contiguous().view(batch, out_channels, -1)
    return output_hidden_states

Performs location-variable convolution operation on the input sequence (hidden_states) using the local
convolution kernel. This was introduced in [LVCNet: Efficient Condition-Dependent Modeling Network for Waveform
Generation](https://huggingface.co/papers/2102.10815) by Zhen Zheng, Jianzong Wang, Ning Cheng, and Jing Xiao.

Time: 414 μs ± 309 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each), test on NVIDIA V100.

Args:
    hidden_states (`torch.FloatTensor` of shape `(batch_size, in_channels, in_length)`):
        The input sequence of shape (batch, in_channels, in_length).
    kernel (`torch.FloatTensor` of shape `(batch_size, in_channels, out_channels, kernel_size, kernel_length)`):
        The local convolution kernel of shape (batch, in_channels, out_channels, kernel_size, kernel_length).
    bias (`torch.FloatTensor` of shape `(batch_size, out_channels, kernel_length)`):
        The bias for the local convolution of shape (batch, out_channels, kernel_length).
    dilation (`int`, *optional*, defaults to 1):
        The dilation of convolution.
    hop_size (`int`, *optional*, defaults to 256):
        The hop_size of the conditioning sequence.
Returns:
    `torch.FloatTensor`: the output sequence after performing local convolution with shape (batch_size,
    out_channels, in_length).


def conv3d(inputs, 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, reuse=None):
    warnings.warn('`tf.layers.conv3d` is deprecated and will be removed in a future version. Please Use `tf.keras.layers.Conv3D` instead.')
    layer = Conv3D(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, _reuse=reuse, _scope=name)
    return layer.apply(inputs)

Functional interface for the 3D convolution layer.

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:
  inputs: Tensor input.
  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.
  reuse: Boolean, whether to reuse the weights of a previous layer
    by the same name.

Returns:
  Output tensor.

Raises:
  ValueError: if eager execution is enabled.


class TvpLoss(nn.Module):

    def __init__(self, losses):
        super().__init__()
        self.loss_map = {'iou': self.loss_iou, 'distance': self.loss_distance, 'duration': self.loss_duration}
        for loss in losses:
            if loss not in self.loss_map:
                raise ValueError(f'Loss {loss} not supported')
        self.losses = losses

    def loss_iou(self, start_time, end_time, candidates_start_time, candidates_end_time, duration):
        """
        Measure the intersection over union.
        """
        inter = torch.min(candidates_end_time, end_time) - torch.max(candidates_start_time, start_time)
        union = torch.max(candidates_end_time, end_time) - torch.min(candidates_start_time, start_time)
        iou = 1 - inter.clamp(min=0) / union
        return iou

    def loss_distance(self, start_time, end_time, candidates_start_time, candidates_end_time, duration):
        """
        Measure the distance of mid points.
        """
        mid_candidates = torch.div(torch.add(candidates_start_time, candidates_end_time), 2.0)
        mid_groundtruth = torch.div(torch.add(start_time, end_time), 2.0)
        distance_diff = torch.div(torch.max(mid_candidates, mid_groundtruth) - torch.min(mid_candidates, mid_groundtruth), duration).clamp(min=0.2)
        return distance_diff

    def loss_duration(self, start_time, end_time, candidates_start_time, candidates_end_time, duration):
        """
        Measure the difference of duration.
        """
        duration_candidates = torch.sub(candidates_end_time, candidates_start_time)
        duration_groundtruth = torch.sub(end_time, start_time)
        duration_diff = torch.square(torch.div(torch.sub(duration_candidates, duration_groundtruth), duration))
        duration_diff = duration_diff.clamp(min=0.4)
        return duration_diff

    def forward(self, logits, labels):
        """
        This performs the loss computation.

        Args:
            logits (`torch.FloatTensor`):
                The output logits of head module.
            labels (`List[torch.FloatTensor]`):
                List of tensors ([start, end, duration]), which contains start time, end time of the video corresponding to the text, and also the duration.
        """
        duration, start_time, end_time = labels
        candidates = torch.mul(logits, duration)
        candidates_start_time, candidates_end_time = (candidates[:, 0].float(), candidates[:, 1].float())
        losses_dict = {}
        for loss in self.losses:
            losses_dict.update({loss: self.loss_map[loss](start_time, end_time, candidates_start_time, candidates_end_time, duration)})
        return losses_dict

This class computes the losses for `TvpForVideoGrounding`. The process happens in two steps: 1) we compute
hungarian assignment between ground truth boxes and the outputs of the model 2) we supervise each pair of matched
ground-truth / prediction (supervise class and box).

Args:
    losses (`List[str]`):
        List of all the losses to be applied.


def load(path, element_spec=None, compression=None, reader_func=None):
    return dataset_ops.Dataset.load(path, element_spec, compression, reader_func)

Loads a previously saved dataset.

Example usage:

>>> import tempfile
>>> path = os.path.join(tempfile.gettempdir(), "saved_data")
>>> # Save a dataset
>>> dataset = tf.data.Dataset.range(2)
>>> tf.data.experimental.save(dataset, path)
>>> new_dataset = tf.data.experimental.load(path)
>>> for elem in new_dataset:
...   print(elem)
tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(1, shape=(), dtype=int64)

If the default option of sharding the saved dataset was used, the element
order of the saved dataset will be preserved when loading it.

The `reader_func` argument can be used to specify a custom order in which
elements should be loaded from the individual shards. The `reader_func` is
expected to take a single argument -- a dataset of datasets, each containing
elements of one of the shards -- and return a dataset of elements. For
example, the order of shards can be shuffled when loading them as follows:

```python
def custom_reader_func(datasets):
  datasets = datasets.shuffle(NUM_SHARDS)
  return datasets.interleave(lambda x: x, num_parallel_calls=AUTOTUNE)

dataset = tf.data.experimental.load(
    path="/path/to/data", ..., reader_func=custom_reader_func)
```

Args:
  path: Required. A path pointing to a previously saved dataset.
  element_spec: Optional. A nested structure of `tf.TypeSpec` objects matching
    the structure of an element of the saved dataset and specifying the type
    of individual element components. If not provided, the nested structure of
    `tf.TypeSpec` saved with the saved dataset is used. Note that this
    argument is required in graph mode.
  compression: Optional. The algorithm to use to decompress the data when
    reading it. Supported options are `GZIP` and `NONE`. Defaults to `NONE`.
  reader_func: Optional. A function to control how to read data from shards.
    If present, the function will be traced and executed as graph computation.

Returns:
  A `tf.data.Dataset` instance.

Raises:
  FileNotFoundError: If `element_spec` is not specified and the saved nested
    structure of `tf.TypeSpec` can not be located with the saved dataset.
  ValueError: If `element_spec` is not specified and the method is executed
    in graph mode.


def GetConfig(self, request, global_params=None):
    config = self.GetMethodConfig('GetConfig')
    return self._RunMethod(config, request, global_params=global_params)

Get encoded debug configuration for component. Not cacheable.

Args:
  request: (DataflowProjectsJobsDebugGetConfigRequest) input message
  global_params: (StandardQueryParameters, default: None) global arguments
Returns:
  (GetDebugConfigResponse) The response message.


def build_inputs_with_special_tokens(self, token_ids_0: List[int], token_ids_1: Optional[List[int]]=None) -> List[int]:
    if token_ids_1 is None:
        return [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
    cls = [self.cls_token_id]
    sep = [self.sep_token_id]
    return cls + token_ids_0 + sep + token_ids_1 + sep

Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
adding special tokens. A REALM sequence has the following format:

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

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

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


def repeat_with_axis(data, repeats, axis, name=None):
    use_optimized_non_xla_implementation = False
    if not isinstance(axis, int):
        raise TypeError(f'Argument `axis` must be an int. Received `axis` = {axis} of type {type(axis).__name__}')
    with ops.name_scope(name, 'Repeat', [data, repeats]):
        data = ops.convert_to_tensor(data, name='data')
        if not use_optimized_non_xla_implementation:
            repeats = convert_to_int_tensor(repeats, name='repeats')
        else:
            repeats = convert_to_int_tensor(repeats, name='repeats', dtype=None)
        repeats.shape.with_rank_at_most(1)
        data = _with_nonzero_rank(data)
        data_shape = shape(data, out_type=repeats.dtype)
        axis = get_positive_axis(axis, data.shape.rank, ndims_name='rank(data)')
        if repeats.shape.num_elements() == 1:
            repeats = reshape(repeats, [])
            expanded = expand_dims(data, axis + 1)
            tiled = tile_one_dimension(expanded, axis + 1, repeats)
            result_shape = concat([data_shape[:axis], [repeats * data_shape[axis]], data_shape[axis + 1:]], axis=0)
            return reshape(tiled, result_shape)
        if repeats.shape.ndims == 1:
            data.shape.dims[axis].assert_is_compatible_with(repeats.shape[0])
        repeats = broadcast_to(repeats, [data_shape[axis]])
        if not use_optimized_non_xla_implementation:
            repeats_original = repeats
            if repeats.shape.ndims != axis + 1:
                repeats_shape = shape(repeats)
                repeats_ndims = rank(repeats)
                broadcast_shape = concat([data_shape[:axis + 1 - repeats_ndims], repeats_shape], axis=0)
                repeats = broadcast_to(repeats, broadcast_shape)
                repeats.set_shape([None] * (axis + 1))
            max_repeat = gen_math_ops._max(repeats, _all_dimensions(repeats))
            max_repeat = gen_math_ops.maximum(ops.convert_to_tensor(0, name='zero', dtype=max_repeat.dtype), max_repeat)
            mask = sequence_mask(repeats, max_repeat)
            expanded = expand_dims(data, axis + 1)
            tiled = tile_one_dimension(expanded, axis + 1, max_repeat)
            masked = boolean_mask(tiled, mask)
            if axis == 0:
                result = masked
            else:
                repeated_dim_size = gen_math_ops._sum(repeats_original, axis=gen_math_ops._range(0, rank(repeats_original), 1))
                result_shape = concat([data_shape[:axis], [repeated_dim_size], data_shape[axis + 1:]], axis=0)
                result = reshape(masked, result_shape)
            if data.shape.ndims is not None:
                new_axis_size = 0 if repeats.shape[0] == 0 else None
                result.set_shape(data.shape[:axis].concatenate([new_axis_size]).concatenate(data.shape[axis + 1:]))
            return result
        else:
            repeats_scan = gen_math_ops.cumsum(repeats)
            output_size = concat([zeros(1, dtype=repeats_scan.dtype), repeats_scan], axis=0)[-1]
            output_indices = gen_math_ops.range(output_size, dtype=repeats.dtype)
            gather_indices = searchsorted(repeats_scan, output_indices, side='right', out_type=repeats.dtype)
            return gather(data, gather_indices, axis=axis)

Repeats elements of `data`.

Args:
  data: An `N`-dimensional tensor.
  repeats: A 1-D integer tensor specifying how many times each element in
    `axis` should be repeated.  `len(repeats)` must equal `data.shape[axis]`.
    Supports broadcasting from a scalar value.
  axis: `int`.  The axis along which to repeat values.  Must be less than
    `max(N, 1)`.
  name: A name for the operation.

Returns:
  A tensor with `max(N, 1)` dimensions.  Has the same shape as `data`,
  except that dimension `axis` has size `sum(repeats)`.

Example usage:

>>> repeat(['a', 'b', 'c'], repeats=[3, 0, 2], axis=0)
<tf.Tensor: shape=(5,), dtype=string,
numpy=array([b'a', b'a', b'a', b'c', b'c'], dtype=object)>
>>> repeat([[1, 2], [3, 4]], repeats=[2, 3], axis=0)
<tf.Tensor: shape=(5, 2), dtype=int32, numpy=
array([[1, 2],
       [1, 2],
       [3, 4],
       [3, 4],
       [3, 4]], dtype=int32)>
>>> repeat([[1, 2], [3, 4]], repeats=[2, 3], axis=1)
<tf.Tensor: shape=(2, 5), dtype=int32, numpy=
array([[1, 1, 2, 2, 2],
       [3, 3, 4, 4, 4]], dtype=int32)>


def sparse_placeholder(dtype, shape=None, name=None):
    if context.executing_eagerly():
        raise RuntimeError('`sparse_placeholder` is not compatible with eager execution.')
    shape_name = name + '/shape' if name is not None else None
    default_shape_name = name + '/shape_default' if name is not None else None
    if shape is None:
        rank = None
        dense_shape = placeholder(dtypes.int64, shape=[rank], name=shape_name)
        dense_shape_default = tensor_util.constant_value_as_shape(dense_shape)
    else:
        if isinstance(shape, tensor_lib.Tensor):
            rank = shape.get_shape()[0]
            dense_shape_default = tensor_util.constant_value_as_shape(shape)
        else:
            rank = len(shape)
            dense_shape_default = tensor_shape.TensorShape(tuple((None if dim == -1 else dim for dim in shape)))
            shape = tuple((tensor_shape.dimension_value(dim) for dim in shape))
            shape = tuple((-1 if dim is None else dim for dim in shape))
            shape = ops.convert_to_tensor(shape, dtype=dtypes.int64, name=default_shape_name)
        dense_shape = placeholder_with_default(shape, shape=shape.shape, name=shape_name)
    result = sparse_tensor.SparseTensor(values=placeholder(dtype, shape=[None], name=name + '/values' if name is not None else None), indices=placeholder(dtypes.int64, shape=[None, rank], name=name + '/indices' if name is not None else None), dense_shape=dense_shape)
    result.set_shape(dense_shape_default)
    return result

Inserts a placeholder for a sparse tensor that will be always fed.

**Important**: This sparse tensor will produce an error if evaluated.
Its value must be fed using the `feed_dict` optional argument to
`Session.run()`, `Tensor.eval()`, or `Operation.run()`.

For example:

```python
x = tf.compat.v1.sparse.placeholder(tf.float32)
y = tf.sparse.reduce_sum(x)

with tf.compat.v1.Session() as sess:
  print(sess.run(y))  # ERROR: will fail because x was not fed.

  indices = np.array([[3, 2, 0], [4, 5, 1]], dtype=np.int64)
  values = np.array([1.0, 2.0], dtype=np.float32)
  shape = np.array([7, 9, 2], dtype=np.int64)
  print(sess.run(y, feed_dict={
    x: tf.compat.v1.SparseTensorValue(indices, values, shape)}))  # Will
    succeed.
  print(sess.run(y, feed_dict={
    x: (indices, values, shape)}))  # Will succeed.

  sp = tf.sparse.SparseTensor(indices=indices, values=values,
                              dense_shape=shape)
  sp_value = sp.eval(session=sess)
  print(sess.run(y, feed_dict={x: sp_value}))  # Will succeed.
```

Args:
  dtype: The type of `values` elements in the tensor to be fed.
  shape: The shape of the tensor to be fed (optional). If the shape is not
    specified, you can feed a sparse tensor of any shape.
  name: A name for prefixing the operations (optional).

Returns:
  A `SparseTensor` that may be used as a handle for feeding a value, but not
  evaluated directly.

Raises:
  RuntimeError: if eager execution is enabled

@compatibility(TF2)
This API is not compatible with eager execution and `tf.function`. To migrate
to TF2, rewrite the code to be compatible with eager execution. Check the
[migration
guide](https://www.tensorflow.org/guide/migrate#1_replace_v1sessionrun_calls)
on replacing `Session.run` calls. In TF2, you can just pass tensors directly
into ops and layers. If you want to explicitly set up your inputs, also see
[Keras functional API](https://www.tensorflow.org/guide/keras/functional) on
how to use `tf.keras.Input` to replace `tf.compat.v1.sparse_placeholder`.
`tf.function` arguments also do the job of `tf.compat.v1.sparse_placeholder`.
For more details please read [Better
performance with tf.function](https://www.tensorflow.org/guide/function).
@end_compatibility


def get_inlined_extension_url(field: descriptor.FieldDescriptor) -> str:
    options = annotation_utils.get_options(field)
    if options.HasExtension(annotations_pb2.fhir_inlined_extension_url):
        return options.Extensions[annotations_pb2.fhir_inlined_extension_url]
    return field.camelcase_name

Returns the FHIR inlined extension URL for a field.

Args:
  field: The FieldDescriptor to examine.

Returns:
  The FHIR inlined extension URL, if one exists, otherwise returns the camel-
  case name of the FieldDescriptor.


def tokenize(self, text, never_split=None):
    never_split = self.never_split.union(set(never_split)) if never_split else self.never_split
    text = self._clean_text(text)
    if self.tokenize_chinese_chars:
        text = self._tokenize_chinese_chars(text)
    unicode_normalized_text = unicodedata.normalize('NFC', text)
    orig_tokens = whitespace_tokenize(unicode_normalized_text)
    split_tokens = []
    for token in orig_tokens:
        if token not in never_split:
            if self.do_lower_case:
                token = token.lower()
                if self.strip_accents is not False:
                    token = self._run_strip_accents(token)
            elif self.strip_accents:
                token = self._run_strip_accents(token)
        split_tokens.extend(self._run_split_on_punc(token, never_split))
    output_tokens = whitespace_tokenize(' '.join(split_tokens))
    return output_tokens

Basic Tokenization of a piece of text. For sub-word tokenization, see WordPieceTokenizer.

Args:
    never_split (`List[str]`, *optional*)
        Kept for backward compatibility purposes. Now implemented directly at the base class level (see
        [`PreTrainedTokenizer.tokenize`]) List of token not to split.


def find_all_build_files(dir: str) -> List[Tuple[str, str]]:
    build_file_dirs = []
    for root, _, files in os.walk(dir):
        for file in files:
            if file in BUILD_FILENAMES:
                root = root.strip('./')
                build_file_dirs.append((root, file))
    return build_file_dirs

List all the BUILD files.

Returns:
    The list of (directory, filename) of all BUILD files.


def __init__(self, dim: int, mean_reversion: types.RealTensor, volatility: Union[types.RealTensor, Callable[..., types.RealTensor]], initial_discount_rate_fn: Callable[..., types.RealTensor], corr_matrix: types.RealTensor=None, validate_args: bool=False, dtype: tf.DType=None, name: str=None):
    self._name = name or 'quasi_gaussian_hjm_model'
    with tf.name_scope(self._name):
        self._dtype = dtype or tf.float32
        self._dim = dim + dim ** 2
        self._factors = dim

        def _instant_forward_rate_fn(t):
            t = tf.convert_to_tensor(t, dtype=self._dtype)

            def _log_zero_coupon_bond(x):
                r = tf.convert_to_tensor(initial_discount_rate_fn(x), dtype=self._dtype)
                return -r * x
            rate = -gradient.fwd_gradient(_log_zero_coupon_bond, t, use_gradient_tape=True, unconnected_gradients=tf.UnconnectedGradients.ZERO)
            return rate

        def _initial_discount_rate_fn(t):
            return tf.convert_to_tensor(initial_discount_rate_fn(t), dtype=self._dtype)
        self._instant_forward_rate_fn = _instant_forward_rate_fn
        self._initial_discount_rate_fn = _initial_discount_rate_fn
        mean_reversion = tf.convert_to_tensor(mean_reversion, dtype=dtype, name='mean_reversion')

        def _infer_batch_shape():
            zero = tf.constant([0], dtype=self._dtype)
            return _initial_discount_rate_fn(zero).shape.as_list()[:-1]
        self._batch_shape = _infer_batch_shape()
        self._batch_rank = len(self._batch_shape)
        self._mean_reversion = mean_reversion
        if callable(volatility):
            self._volatility = volatility
        else:
            volatility = tf.convert_to_tensor(volatility, dtype=dtype)
            if self._batch_rank > 0:
                volatility = tf.expand_dims(volatility, axis=self._batch_rank)

            def _tensor_to_volatility_fn(t, r):
                del t, r
                return volatility
            self._volatility = _tensor_to_volatility_fn
        if corr_matrix is None:
            corr_matrix = tf.eye(dim, dim, batch_shape=self._batch_shape, dtype=self._dtype)
        self._rho = tf.convert_to_tensor(corr_matrix, dtype=self._dtype, name='rho')
        if validate_args:
            try:
                self._sqrt_rho = tf.linalg.cholesky(self._rho)
            except:
                raise ValueError('The input correlation matrix is not positive semidefinite.')
        else:
            self._sqrt_rho = _get_valid_sqrt_matrix(self._rho)

    def _vol_fn(t, state):
        """Volatility function of qG-HJM."""
        x = state[..., :self._factors]
        batch_shape_x = x.shape.as_list()[:-1]
        r_t = self._instant_forward_rate_fn(t) + tf.reduce_sum(x, axis=-1, keepdims=True)
        volatility = self._volatility(t, r_t)
        volatility = tf.expand_dims(volatility, axis=-1)
        diffusion_x = tf.broadcast_to(tf.expand_dims(self._sqrt_rho, axis=self._batch_rank) * volatility, batch_shape_x + [self._factors, self._factors])
        paddings = tf.constant([[0, 0]] * len(batch_shape_x) + [[0, self._factors ** 2], [0, self._factors ** 2]], dtype=tf.int32)
        diffusion = tf.pad(diffusion_x, paddings)
        return diffusion

    def _drift_fn(t, state):
        """Drift function of qG-HJM."""
        x = state[..., :self._factors]
        y = state[..., self._factors:]
        batch_shape_x = x.shape.as_list()[:-1]
        y = tf.reshape(y, batch_shape_x + [self._factors, self._factors])
        r_t = self._instant_forward_rate_fn(t) + tf.reduce_sum(x, axis=-1, keepdims=True)
        volatility = self._volatility(t, r_t)
        volatility = tf.expand_dims(volatility, axis=-1)
        volatility_squared = tf.linalg.matmul(volatility, volatility, transpose_b=True)
        mr2 = tf.expand_dims(self._mean_reversion, axis=-1)
        perm = list(range(self._batch_rank)) + [self._batch_rank + 1, self._batch_rank]
        mr2 = mr2 + tf.transpose(mr2, perm=perm)
        mr2 = tf.expand_dims(mr2, axis=self._batch_rank)
        mr = self._mean_reversion
        if self._batch_rank > 0:
            mr = tf.expand_dims(self._mean_reversion, axis=1)
        drift_x = tf.math.reduce_sum(y, axis=-1) - mr * x
        drift_y = tf.expand_dims(self._rho, axis=self._batch_rank) * volatility_squared - mr2 * y
        drift_y = tf.reshape(drift_y, batch_shape_x + [self._factors * self._factors])
        drift = tf.concat([drift_x, drift_y], axis=-1)
        return drift
    super(QuasiGaussianHJM, self).__init__(self._dim, _drift_fn, _vol_fn, self._dtype, self._name)

Initializes a batch of HJM models.

Args:
  dim: A Python scalar which corresponds to the number of factors
    comprising the model.
  mean_reversion: A real positive `Tensor` of shape `batch_shape + [dim]`.
    `batch_shape` denotes the shape of independent HJM models within the
    batch. Corresponds to the mean reversion rate of each factor.
  volatility: A real positive `Tensor` of the same `dtype` and shape as
    `mean_reversion` or a callable with the following properties:
    (a)  The callable should accept a scalar `Tensor` `t` and a `Tensor`
    `r(t)` of shape `batch_shape + [num_samples]` and returns a `Tensor` of
    shape compatible with `batch_shape + [num_samples, dim]`. The variable
    `t`  stands for time and `r(t)` is the short rate at time `t`. The
    function returns instantaneous volatility `sigma(t) = sigma(t, r(t))`.
    When `volatility` is specified as a real `Tensor`, each factor is
    assumed to have a constant instantaneous volatility  and the  model is
    effectively a Gaussian HJM model.
    Corresponds to the instantaneous volatility of each factor.
  initial_discount_rate_fn: A Python callable that accepts expiry time as
    a real `Tensor` of the same `dtype` as `mean_reversion` and returns a
    `Tensor` of shape `batch_shape + input_shape`.
    Corresponds to the zero coupon bond yield at the present time for the
    input expiry time.
  corr_matrix: A `Tensor` of shape `batch_shape + [dim, dim]` and the same
    `dtype` as `mean_reversion`.
    Corresponds to the correlation matrix `Rho`.
  validate_args: Optional boolean flag to enable validation of the input
    correlation matrix. If the flag is enabled and the input correlation
    matrix is not positive semidefinite, an error is raised.
    Default value: False.
  dtype: The default dtype to use when converting values to `Tensor`s.
    Default value: `None` which maps to `tf.float32`.
  name: Python string. The name to give to the ops created by this class.
    Default value: `None` which maps to the default name
    `quasi_gaussian_hjm_model`.


def convert(self):
    if not _jit:
        raise ImportError('Cannot import jit from jax.')
    if not self._serving_funcs:
        raise ValueError('No serving func is specified.')
    if not self._inputs:
        raise ValueError('Input tensors are not specified.')
    if len(self._inputs) != len(self._serving_funcs):
        msg = 'Input tensor mapping len {} does not match serving func len {}.'.format(len(self._inputs), len(self._serving_funcs))
        raise ValueError(msg)
    if not isinstance(self._inputs, (tuple, list)):
        raise ValueError('Input tensors should be pass in a tuple list wrapped in an array.')
    if len(self._serving_funcs) > 1:
        raise ValueError('Currently only support single serving function.')
    if not isinstance(self._inputs[0], (tuple, list)):
        raise ValueError('The input placeholders are not a dictionary.')
    input_names = []
    ordered_inputs = []
    for input_name, tensor in self._inputs[0]:
        input_names.append(input_name)
        ordered_inputs.append(tensor)
    try:
        hlo_proto = _jit(self._serving_funcs[0]).trace(*ordered_inputs).lower(lowering_platforms=('cpu',)).compiler_ir('hlo').as_serialized_hlo_module_proto()
    except Exception:
        raise ValueError('Failed to convert the given Jax function to hlo.')
    converter_kwargs = {'input_content': hlo_proto, 'input_names': input_names, 'is_proto_format': True}
    converter_kwargs.update(self._get_base_converter_args())
    quant_mode = QuantizationMode(self.optimizations, self.target_spec, self.representative_dataset, None, experimental_qdq_annotation=self._experimental_strict_qdq)
    self._validate_inference_input_output_types(quant_mode)
    converter_kwargs.update(quant_mode.converter_flags())
    result = _convert_jax_hlo(**converter_kwargs)
    return self._optimize_tflite_model(result, quant_mode, _build_conversion_flags(**converter_kwargs).debug_options, quant_io=self.experimental_new_quantizer)

Converts a Jax serving func based on instance variables.

Returns:
  The converted data in serialized format.

Raises:
  ImportError:
    If cannot import the jit from jax.
  ValueError:
    No serving function is specified.
    Input tensors are not specified.
    The truth value of an array with more than one element is ambiguous.
    Failed to convert the given Jax function to hlo.


def forward(self, data: torch.Tensor, observed_indicator: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    denominator = observed_indicator.sum(self.dim, keepdim=self.keepdim)
    denominator = denominator.clamp_min(1.0)
    loc = (data * observed_indicator).sum(self.dim, keepdim=self.keepdim) / denominator
    variance = (((data - loc) * observed_indicator) ** 2).sum(self.dim, keepdim=self.keepdim) / denominator
    scale = torch.sqrt(variance + self.minimum_scale)
    return ((data - loc) / scale, loc, scale)

Parameters:
    data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`):
        input for Batch norm calculation
    observed_indicator (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`):
        Calculating the scale on the observed indicator.
Returns:
    tuple of `torch.Tensor` of shapes
        (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`,
        `(batch_size, 1, num_input_channels)`)


def are_compatible(spec1, spec2):
    try:
        nest.assert_same_structure(spec1, spec2)
    except TypeError:
        return False
    except ValueError:
        return False
    for s1, s2 in zip(nest.flatten(spec1), nest.flatten(spec2)):
        if not s1.is_compatible_with(s2) or not s2.is_compatible_with(s1):
            return False
    return True

Indicates whether two type specifications are compatible.

Two type specifications are compatible if they have the same nested structure
and the their individual components are pair-wise compatible.

Args:
  spec1: A `tf.TypeSpec` object to compare.
  spec2: A `tf.TypeSpec` object to compare.

Returns:
  `True` if the two type specifications are compatible and `False` otherwise.


def __getitem__(self, key):
    if key is None:
        key = self._key()
    value = self._get_recursive(key)
    if value is None:
        value = self[key] = self.default_factory()
    return value

Gets the value at key (or current context), or sets default value.

Args:
  key: May be `None` or `Graph`object. When `None`, the key is set to the
    current context.

Returns:
  Either the cached or default value.


def train_on_batch(self, x, y=None, sample_weight=None, class_weight=None, return_dict=False):
    raise NotImplementedError

Runs a single gradient update on a single batch of data.

Args:
    x: Input data. Must be array-like.
    y: Target data. Must be array-like.
    sample_weight: Optional array of the same length as x, containing
        weights to apply to the model's loss for each sample.
        In the case of temporal data, you can pass a 2D array
        with shape `(samples, sequence_length)`, to apply a different
        weight to every timestep of every sample.
    class_weight: Optional dictionary mapping class indices (integers)
        to a weight (float) to apply to the model's loss for the samples
        from this class during training. This can be useful to tell the
        model to "pay more attention" to samples from an
        under-represented class. When `class_weight` is specified
        and targets have a rank of 2 or greater, either `y` must
        be one-hot encoded, or an explicit final dimension of 1
        must be included for sparse class labels.
    return_dict: If `True`, loss and metric results are returned as a
        dict, with each key being the name of the metric. If `False`,
        they are returned as a list.

Returns:
    A scalar loss value (when no metrics and `return_dict=False`),
    a list of loss and metric values
    (if there are metrics and `return_dict=False`), or a dict of
    metric and loss values (if `return_dict=True`).


def run(argv=None, save_main_session=True):
    known_args, pipeline_args = parse_known_args(argv)
    pipeline_options = PipelineOptions(pipeline_args)
    pipeline_options.view_as(SetupOptions).save_main_session = save_main_session
    engine_handler = KeyedModelHandler(TensorRTEngineHandlerNumPy(min_batch_size=1, max_batch_size=1, engine_path=known_args.engine_path))
    with beam.Pipeline(options=pipeline_options) as p:
        filename_value_pair = p | 'ReadImageNames' >> beam.io.ReadFromText(known_args.input) | 'ReadImageData' >> beam.Map(lambda image_name: read_image(image_file_name=image_name, path_to_dir=known_args.images_dir)) | 'AttachImageSizeToKey' >> beam.Map(attach_im_size_to_key) | 'PreprocessImages' >> beam.MapTuple(lambda file_name, data: (file_name, preprocess_image(data)))
        predictions = filename_value_pair | 'TensorRTRunInference' >> RunInference(engine_handler) | 'ProcessOutput' >> beam.ParDo(PostProcessor())
        _ = predictions | 'WriteOutputToGCS' >> beam.io.WriteToText(known_args.output, shard_name_template='', append_trailing_newlines=True)

Args:
  argv: Command line arguments defined for this example.


def reverse_bettertransformer(self):
    if not is_optimum_available():
        raise ImportError('The package `optimum` is required to use Better Transformer.')
    from optimum.version import __version__ as optimum_version
    if version.parse(optimum_version) < version.parse('1.7.0'):
        raise ImportError(f'Please install optimum>=1.7.0 to use Better Transformer. The version {optimum_version} was found.')
    from optimum.bettertransformer import BetterTransformer
    return BetterTransformer.reverse(self)

Reverts the transformation from [`~PreTrainedModel.to_bettertransformer`] so that the original modeling is
used, for example in order to save the model.

Returns:
    [`PreTrainedModel`]: The model converted back to the original modeling.


def from_backbone_and_decoder_configs(cls, backbone_config: PretrainedConfig, decoder_config: PretrainedConfig, **kwargs):
    return cls(backbone_config=backbone_config, decoder_config=decoder_config, **kwargs)

Instantiate a [`MaskFormerConfig`] (or a derived class) from a pre-trained backbone model configuration and DETR model
configuration.

    Args:
        backbone_config ([`PretrainedConfig`]):
            The backbone configuration.
        decoder_config ([`PretrainedConfig`]):
            The transformer decoder configuration to use.

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


def serialize_keras_object(obj):
    if obj is None:
        return obj
    if isinstance(obj, PLAIN_TYPES):
        return obj
    if isinstance(obj, (list, tuple)):
        config_arr = [serialize_keras_object(x) for x in obj]
        return tuple(config_arr) if isinstance(obj, tuple) else config_arr
    if isinstance(obj, dict):
        return serialize_dict(obj)
    if isinstance(obj, bytes):
        return {'class_name': '__bytes__', 'config': {'value': obj.decode('utf-8')}}
    if isinstance(obj, slice):
        return {'class_name': '__slice__', 'config': {'start': serialize_keras_object(obj.start), 'stop': serialize_keras_object(obj.stop), 'step': serialize_keras_object(obj.step)}}
    if isinstance(obj, type(Ellipsis)):
        return {'class_name': '__ellipsis__', 'config': {}}
    if isinstance(obj, backend.KerasTensor):
        history = getattr(obj, '_keras_history', None)
        if history:
            history = list(history)
            history[0] = history[0].name
        return {'class_name': '__keras_tensor__', 'config': {'shape': obj.shape, 'dtype': obj.dtype, 'keras_history': history}}
    if tf.available and isinstance(obj, tf.TensorShape):
        return obj.as_list() if obj._dims is not None else None
    if backend.is_tensor(obj):
        return {'class_name': '__tensor__', 'config': {'value': backend.convert_to_numpy(obj).tolist(), 'dtype': backend.standardize_dtype(obj.dtype)}}
    if type(obj).__module__ == np.__name__:
        if isinstance(obj, np.ndarray) and obj.ndim > 0:
            return {'class_name': '__numpy__', 'config': {'value': obj.tolist(), 'dtype': backend.standardize_dtype(obj.dtype)}}
        else:
            return obj.item()
    if tf.available and isinstance(obj, tf.DType):
        return obj.name
    if isinstance(obj, types.FunctionType) and obj.__name__ == '<lambda>':
        warnings.warn(f'The object being serialized includes a `lambda`. This is unsafe. In order to reload the object, you will have to pass `safe_mode=False` to the loading function. Please avoid using `lambda` in the future, and use named Python functions instead. This is the `lambda` being serialized: {inspect.getsource(obj)}', stacklevel=2)
        return {'class_name': '__lambda__', 'config': {'value': python_utils.func_dump(obj)}}
    if tf.available and isinstance(obj, tf.TypeSpec):
        ts_config = obj._serialize()
        ts_config = list(map(lambda x: x.as_list() if isinstance(x, tf.TensorShape) else x.name if isinstance(x, tf.DType) else x, ts_config))
        return {'class_name': '__typespec__', 'spec_name': obj.__class__.__name__, 'module': obj.__class__.__module__, 'config': ts_config, 'registered_name': None}
    inner_config = _get_class_or_fn_config(obj)
    config_with_public_class = serialize_with_public_class(obj.__class__, inner_config)
    if config_with_public_class is not None:
        get_build_and_compile_config(obj, config_with_public_class)
        record_object_after_serialization(obj, config_with_public_class)
        return config_with_public_class
    if isinstance(obj, types.FunctionType):
        module = obj.__module__
    else:
        module = obj.__class__.__module__
    class_name = obj.__class__.__name__
    if module == 'builtins':
        registered_name = None
    elif isinstance(obj, types.FunctionType):
        registered_name = object_registration.get_registered_name(obj)
    else:
        registered_name = object_registration.get_registered_name(obj.__class__)
    config = {'module': module, 'class_name': class_name, 'config': inner_config, 'registered_name': registered_name}
    get_build_and_compile_config(obj, config)
    record_object_after_serialization(obj, config)
    return config

Retrieve the config dict by serializing the Keras object.

`serialize_keras_object()` serializes a Keras object to a python dictionary
that represents the object, and is a reciprocal function of
`deserialize_keras_object()`. See `deserialize_keras_object()` for more
information about the config format.

Args:
    obj: the Keras object to serialize.

Returns:
    A python dict that represents the object. The python dict can be
    deserialized via `deserialize_keras_object()`.


class BlipTextVisionModelOutput(ModelOutput):
    loss: Optional[torch.FloatTensor] = None
    image_embeds: Optional[torch.FloatTensor] = None
    last_hidden_state: Optional[torch.FloatTensor] = None
    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None

Adapted from the base class for vision model's outputs that also contains image embeddings of the pooling of the
last hidden states. This class also adds the loss term from the text decoder.

Args:
    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
        Language modeling loss from the text decoder.
    image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`):
        The image embeddings obtained by applying the projection layer to the pooler_output.
    last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
        Sequence of hidden-states at the output of the last layer of the model.
    hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
        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 optional initial embedding outputs.
    attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
        Tuple of `torch.FloatTensor` (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.


def dup(node, copy_map, field_name='___pyct_anno'):
    for n in gast.walk(node):
        for k in copy_map:
            if hasanno(n, k, field_name):
                setanno(n, copy_map[k], getanno(n, k, field_name), field_name)

Recursively copies annotations in an AST tree.

Args:
  node: ast.AST
  copy_map: Dict[Hashable, Hashable], maps a source anno key to a destination
      key. All annotations with the source key will be copied to identical
      annotations with the destination key.
  field_name: str


class Zeros(Initializer):

    def __call__(self, shape, dtype=None):
        """Returns a tensor object initialized as specified by the initializer.

        Args:
            shape: Shape of the tensor.
            dtype: Optional dtype of the tensor. Only numeric or boolean dtypes
                are supported. If not specified, `keras.backend.floatx()`
                is used, which default to `float32` unless you configured it
                otherwise (via `keras.backend.set_floatx(float_dtype)`).
        """
        dtype = standardize_dtype(dtype)
        return ops.zeros(shape, dtype=dtype)

Initializer that generates tensors initialized to 0.

Examples:

>>> # Standalone usage:
>>> initializer = Zeros()
>>> values = initializer(shape=(2, 2))

>>> # Usage in a Keras layer:
>>> initializer = Zeros()
>>> layer = Dense(units=3, kernel_initializer=initializer)


def trivial_reward(example):
    return example

Reward for the trivial search space.

The reward (i.e. fitness) is the value itself. The goal of the search,
therefore, is to find the value 1.

Args:
  example: a materialized value.

Returns:
  The corresponding reward.


def _dynamic_range_quantize(src_saved_model_path: str, dst_saved_model_path: str, quantization_options: _QuantizationOptions) -> autotrackable.AutoTrackable:
    mode_str = 'dynamic-range quantization'
    if _is_qat_saved_model(src_saved_model_path):
        raise ValueError('The models trained with quantization-aware training (QAT) is not supported for %s.' % mode_str)
    logging.info('Running post-training %s on model: %s', mode_str, src_saved_model_path)
    logging.info('QuantizationOptions: \n%s', quantization_options)
    signature_def_map = save_model.get_signatures_from_saved_model(src_saved_model_path, quantization_options.signature_keys, quantization_options.tags)
    pywrap_quantize_model.quantize_ptq_dynamic_range(src_saved_model_path, dst_saved_model_path, quantization_options_serialized=quantization_options.SerializeToString(), signature_keys=list(quantization_options.signature_keys), signature_def_map_serialized=_serialize_signature_def_map(signature_def_map), py_function_library=py_function_lib.PyFunctionLibrary())
    return saved_model_load.load(dst_saved_model_path)

Quantizes the given SavedModel via post-training dynamic range quantization.

Args:
  src_saved_model_path: Path to the saved model.
  dst_saved_model_path: The path to save the output SavedModel. The directory
    will be overwritten if not empty.
  quantization_options: QuantizationOptions proto describing quantization
    related config.

Returns:
  A SavedModel object with TF quantization applied.

Raises:
  ValueError: when the model is QAT model.


def connect_raise_node(self, node, except_guards):
    for guard in except_guards:
        if guard in self.raises:
            self.raises[guard].append(node)
        else:
            self.raises[guard] = [node]

Adds extra connection between a raise node and containing except guards.

The node is a graph node, not an ast node.

Args:
  node: Node
  except_guards: Tuple[ast.AST, ...], the except sections that guard node


def _prepare_4d_attention_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int]=None):
    return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)

Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
`(batch_size, key_value_length)`

Args:
    mask (`torch.Tensor`):
        A 2D attention mask of shape `(batch_size, key_value_length)`
    dtype (`torch.dtype`):
        The torch dtype the created mask shall have.
    tgt_len (`int`):
        The target length or query length the created mask shall have.