def _create_variables(self, num_clusters):
    init_value = array_ops.placeholder_with_default([], shape=None)
    cluster_centers = variable_v1.VariableV1(init_value, name=CLUSTERS_VAR_NAME, validate_shape=False)
    cluster_centers_initialized = variable_v1.VariableV1(False, dtype=dtypes.bool, name='initialized')
    if self._use_mini_batch and self._mini_batch_steps_per_iteration > 1:
        cluster_centers_updated = variable_v1.VariableV1(init_value, name='clusters_updated', validate_shape=False)
        update_in_steps = variable_v1.VariableV1(self._mini_batch_steps_per_iteration, dtype=dtypes.int64, name='update_in_steps')
        cluster_counts = variable_v1.VariableV1(array_ops.zeros([num_clusters], dtype=dtypes.int64))
    else:
        cluster_centers_updated = cluster_centers
        update_in_steps = None
        cluster_counts = variable_v1.VariableV1(array_ops.ones([num_clusters], dtype=dtypes.int64)) if self._use_mini_batch else None
    return (cluster_centers, cluster_centers_initialized, cluster_counts, cluster_centers_updated, update_in_steps)

Creates variables.

Args:
  num_clusters: an integer Tensor providing the number of clusters.

Returns:
  Tuple with following elements:
  - cluster_centers: a Tensor for storing cluster centers
  - cluster_centers_initialized: bool Variable indicating whether clusters
        are initialized.
  - cluster_counts: a Tensor for storing counts of points assigned to this
        cluster. This is used by mini-batch training.
  - cluster_centers_updated: Tensor representing copy of cluster centers
        that are updated every step.
  - update_in_steps: numbers of steps left before we sync
        cluster_centers_updated back to cluster_centers.


def _pad_image(self, images: 'torch.tensor', size_divisibility: int=32) -> 'torch.tensor':
    height, width = get_image_size(images, channel_dim=ChannelDimension.FIRST)
    pad_height = 0 if height % size_divisibility == 0 else size_divisibility - height % size_divisibility
    pad_width = 0 if width % size_divisibility == 0 else size_divisibility - width % size_divisibility
    if pad_width + pad_height > 0:
        padding = (0, 0, pad_width, pad_height)
        images = F.pad(images, padding)
    return images

Pads an image or batched images constantly so that width and height are divisible by size_divisibility

Args:
    image (`torch,tensor`):
        Image to pad.
    size_divisibility (`int`, *optional*, defaults to 32):
        The width and height of the image will be padded to be divisible by this number.


def _remove_ancillary_layers(model, layer_map, layers):
    ancillary_layers = []
    if not model._is_graph_network:
        return (layers, ancillary_layers)
    depths = [depth for depth in model._nodes_by_depth.keys() if depth < 0]
    depths.sort(reverse=True)
    for depth in depths:
        for node in model._nodes_by_depth[depth]:
            ancillary_layers.append(layer_map[node.outbound_layer])
    return ([l for l in layers if l not in ancillary_layers], ancillary_layers)

Removes and returns any ancillary layers from `layers` based on `model`.

Ancillary layers are part of the model topology but not used to compute the
model outputs, e.g., layers from `add_loss` and `add_metric`.

Args:
  model: A Keras Model.
  layer_map: A map to from layers in the `model` to those in `layers`.
  layers: A list of all layers.

Returns:
  Two lists of layers: (1) `layers` with the ancillary layers removed, and (2)
  the ancillary layers.


def set_number_of_shards(self, number_of_shards):
    for policy in self._sharding_policies:
        policy.set_number_of_shards(number_of_shards)
        policy.set_number_of_partitions(self._number_of_partitions)
    self._validate()

Sets the number of shards to use for the InfeedQueue.

Args:
  number_of_shards: number of ways to shard the InfeedQueue.

Raises:
  ValueError: if number_of_shards is not > 0; or the policies have
    been frozen and number_of_shards was already set to something
    else.


def _save_and_write_assets(self, assets_collection_to_add=None):
    asset_filename_map = _maybe_save_assets(_add_asset_to_collection, assets_collection_to_add)
    if not asset_filename_map:
        tf_logging.info('No assets to write.')
        return
    copy_assets_to_destination_dir(asset_filename_map, self._export_dir, self._saved_asset_files)

Saves asset to the meta graph and writes asset files to disk.

Args:
  assets_collection_to_add: The collection where the asset paths are setup.


def oneof(candidates: Iterable[Any], *, name: Optional[str]=None, hints: Optional[Any]=None) -> Any:
    return OneOf(candidates=list(candidates), name=name, hints=hints)

N choose 1.

Example::

  @pg.members([
    ('x', pg.typing.Int())
  ])
  class A(pg.Object):
    pass

  # A single categorical choice:
  v = pg.oneof([1, 2, 3])

  # A complex type as candidate.
  v1 = pg.oneof(['a', {'x': 1}, A(1)])

  # A hierarchical categorical choice:
  v2 = pg.oneof([
      'foo',
      'bar',
      A(pg.oneof([1, 2, 3]))
  ])

See also:

  * :class:`pyglove.hyper.OneOf`
  * :func:`pyglove.manyof`
  * :func:`pyglove.floatv`
  * :func:`pyglove.permutate`
  * :func:`pyglove.evolve`

.. note::

  Under symbolic mode (by default), `pg.oneof` returns a ``pg.hyper.OneOf``
  object. Under dynamic evaluation mode, which is called under the context of
  :meth:`pyglove.hyper.DynamicEvaluationContext.collect` or
  :meth:`pyglove.hyper.DynamicEvaluationContext.apply`, it evaluates to
  a concrete candidate value.

  To use conditional search space in dynamic evaluation mode, the candidate
  should be wrapped with a `lambda` function, which is not necessary under
  symbolic mode. For example::

    pg.oneof([lambda: pg.oneof([0, 1], name='sub'), 2], name='root')

Args:
  candidates: Candidates to select from. Items of candidate can be any type,
    therefore it can have nested hyper primitives, which forms a hierarchical
    search space.
  name: A name that can be used to identify a decision point in the search
    space. This is needed when the code to instantiate the same hyper
    primitive may be called multiple times under a
    `pg.DynamicEvaluationContext.collect` context or under a
    `pg.DynamicEvaluationContext.apply` context.
  hints: An optional value which acts as a hint for the controller.

Returns:
  In symbolic mode, this function returns a `ChoiceValue`.
  In dynamic evaluation mode, this function returns one of the items in
  `candidates`.
  If evaluated under a `pg.DynamicEvaluationContext.apply` scope,
  this function will return the selected candidate.
  If evaluated under a `pg.DynamicEvaluationContext.collect`
  scope, it will return the first candidate.


def _validate_snapshot(path: str, metadata: snapshot_pb2.DistributedSnapshotMetadata, element_spec: Any, compression: str) -> None:
    error_file = _pywrap_snapshot_utils.TF_DATA_SnapshotErrorFilePath(path)
    if gfile.Exists(error_file):
        with gfile.GFile(error_file, 'r') as f:
            raise ValueError(f'Failed to load tf.data snapshot at {path}. The save job failed to write it. Status: {f.read()}')
    snapshot_element_spec = _parse_element_spec(metadata.element_spec)
    if element_spec and element_spec != snapshot_element_spec:
        raise ValueError(f'Failed to load tf.data snapshot at {path}. User specified element_spec {element_spec}, but the actual element_spec is {snapshot_element_spec}.')
    if compression and compression != metadata.compression:
        raise ValueError(f'Failed to load tf.data snapshot at {path}. User specified compression {compression}, but the actual compression is {metadata.compression}.')

Validates a tf.data distributed snapshot.

Args:
  path: Root path of the distributed snapshot.
  metadata: The DistributedSnapshotMetadata of the snapshot.
  element_spec: Dataset element_spec.
  compression: Compression method used for saving.

Raises:
  ValueError if the snapshot is invalid.


class TorchModuleWrapper(Layer):

    def __init__(self, module, name=None, output_shape=None, **kwargs):
        super().__init__(name=name, **kwargs)
        import torch.nn as nn
        from keras.src.backend.torch.core import get_device
        if isinstance(module, nn.modules.lazy.LazyModuleMixin) and module.has_uninitialized_params():
            raise ValueError(f'LazyModules are not supported unless they are already initialized. Received uninitialized LazyModule: module={module}')
        self.module = module.to(get_device())
        self._track_module_parameters()
        self.output_shape = output_shape

    def parameters(self, recurse=True):
        return self.module.parameters(recurse=recurse)

    def _track_module_parameters(self):
        for param in self.module.parameters():
            variable = backend.Variable(initializer=param, trainable=param.requires_grad)
            self._track_variable(variable)
        self.built = True

    def call(self, *args, training=None, **kwargs):
        if training is False:
            self.eval()
        else:
            self.train()
        return self.module(*args, **kwargs)

    def save_own_variables(self, store):
        """Saves model's state from `state_dict`.
        `model.parameters` excludes some of model's state like
        `BatchNorm` mean and variance. So, use `state_dict` to obtain
        all of model's state.
        """
        state_dict = self.module.state_dict()
        for key in state_dict.keys():
            store[key] = convert_to_numpy(state_dict[key])

    def load_own_variables(self, store):
        """Loads model's state via `state_dict`."""
        state_dict = {}
        for key in store.keys():
            if isinstance(key, bytes):
                key = key.decode()
            state_dict[key] = convert_to_tensor(store[key])
        self.module.load_state_dict(state_dict)

    def compute_output_shape(self, input_shape):
        if self.output_shape is None:
            return super().compute_output_shape(input_shape)
        return self.output_shape

    def get_config(self):
        base_config = super().get_config()
        import torch
        buffer = io.BytesIO()
        torch.save(self.module, buffer)
        config = {'module': buffer.getvalue(), 'output_shape': self.output_shape}
        return {**base_config, **config}

    @classmethod
    def from_config(cls, config):
        import torch
        if 'module' in config:
            buffer = io.BytesIO(config['module'])
            config['module'] = torch.load(buffer, weights_only=False)
        return cls(**config)

Torch module wrapper layer.

`TorchModuleWrapper` is a wrapper class that can turn any
`torch.nn.Module` into a Keras layer, in particular by making its
parameters trackable by Keras.

`TorchModuleWrapper` is only compatible with the PyTorch backend and
cannot be used with the TensorFlow or JAX backends.

Args:
    module: `torch.nn.Module` instance. If it's a `LazyModule`
        instance, then its parameters must be initialized before
        passing the instance to `TorchModuleWrapper` (e.g. by calling
        it once).
    output_shape :The shape of the output of this layer. It helps Keras
        perform automatic shape inference.
    name: The name of the layer (string).

Example:

Here's an example of how the `TorchModuleWrapper` can be used with vanilla
PyTorch modules.

```python
import torch
import torch.nn as nn
import torch.nn.functional as F

import keras
from keras.layers import TorchModuleWrapper

class Classifier(keras.Model):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        # Wrap `torch.nn.Module`s with `TorchModuleWrapper`
        # if they contain parameters
        self.conv1 = TorchModuleWrapper(
            nn.Conv2d(in_channels=1, out_channels=32, kernel_size=(3, 3))
        )
        self.conv2 = TorchModuleWrapper(
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=(3, 3))
        )
        self.pool = nn.MaxPool2d(kernel_size=(2, 2))
        self.flatten = nn.Flatten()
        self.dropout = nn.Dropout(p=0.5)
        self.fc = TorchModuleWrapper(nn.Linear(1600, 10))

    def call(self, inputs):
        x = F.relu(self.conv1(inputs))
        x = self.pool(x)
        x = F.relu(self.conv2(x))
        x = self.pool(x)
        x = self.flatten(x)
        x = self.dropout(x)
        x = self.fc(x)
        return F.softmax(x, dim=1)

model = Classifier()
model.build((1, 28, 28))
print("Output shape:", model(torch.ones(1, 1, 28, 28).to("cuda")).shape)

model.compile(
    loss="sparse_categorical_crossentropy",
    optimizer="adam",
    metrics=["accuracy"]
)
model.fit(train_loader, epochs=5)
```


class TFBaseModelOutputWithCrossAttentions(ModelOutput):
    last_hidden_state: Optional[tf.Tensor] = None
    hidden_states: Tuple[tf.Tensor] | None = None
    attentions: Tuple[tf.Tensor] | None = None
    cross_attentions: Tuple[tf.Tensor] | None = None

Base class for model's outputs, with potential hidden states and attentions.

Args:
    last_hidden_state (`tf.Tensor` 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(tf.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape
        `(batch_size, sequence_length, hidden_size)`.

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

        Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
        heads.
    cross_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
        Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
        sequence_length)`.

        Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the
        weighted average in the cross-attention heads.


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

Create a mask from the two sequences passed to be used in a sequence-pair classification task. LUKE does not
make use of token type ids, therefore a list of zeros is returned.

Args:
    token_ids_0 (`List[int]`):
        List of IDs.
    token_ids_1 (`List[int]`, *optional*):
        Optional second list of IDs for sequence pairs.

Returns:
    `List[int]`: List of zeros.


def write(tag, tensor, step=None, metadata=None, name=None):
    with ops.name_scope(name, 'write_summary') as scope:
        if _summary_state.writer is None:
            return constant_op.constant(False)
        if step is None:
            step = get_step()
        if metadata is None:
            serialized_metadata = b''
        elif hasattr(metadata, 'SerializeToString'):
            serialized_metadata = metadata.SerializeToString()
        else:
            serialized_metadata = metadata

        def record():
            """Record the actual summary and return True."""
            if step is None:
                raise ValueError('No step set. Please specify one either through the `step` argument or through tf.summary.experimental.set_step()')
            with ops.device('cpu:0'):
                summary_tensor = tensor() if callable(tensor) else array_ops.identity(tensor)
                writer = _summary_state.writer
                summary_value = _maybe_convert_tensor_to_dtensor(writer, summary_tensor)
                step_value = _maybe_convert_tensor_to_dtensor(writer, step)
                write_summary_op = gen_summary_ops.write_summary(writer._resource, step_value, summary_value, tag, serialized_metadata, name=scope)
                with ops.control_dependencies([write_summary_op]):
                    return constant_op.constant(True)
        op = smart_cond.smart_cond(should_record_summaries(), record, _nothing, name='summary_cond')
        if not context.executing_eagerly():
            ops.add_to_collection(ops.GraphKeys._SUMMARY_COLLECTION, op)
        return op

Writes a generic summary to the default SummaryWriter if one exists.

This exists primarily to support the definition of type-specific summary ops
like scalar() and image(), and is not intended for direct use unless defining
a new type-specific summary op.

Args:
  tag: string tag used to identify the summary (e.g. in TensorBoard), usually
    generated with `tf.summary.summary_scope`
  tensor: the Tensor holding the summary data to write or a callable that
    returns this Tensor. If a callable is passed, it will only be called when
    a default SummaryWriter exists and the recording condition specified by
    `record_if()` is met.
  step: Explicit `int64`-castable monotonic step value for this summary. If
    omitted, this defaults to `tf.summary.experimental.get_step()`, which must
    not be None.
  metadata: Optional SummaryMetadata, as a proto or serialized bytes
  name: Optional string name for this op.

Returns:
  True on success, or false if no summary was written because no default
  summary writer was available.

Raises:
  ValueError: if a default writer exists, but no step was provided and
    `tf.summary.experimental.get_step()` is None.


def create(path, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO):
    filesystem = FileSystems.get_filesystem(path)
    return filesystem.create(path, mime_type, compression_type)

Returns a write channel for the given file path.

Args:
  path: string path of the file object to be written to the system
  mime_type: MIME type to specify the type of content in the file object
  compression_type: Type of compression to be used for this object. See
    ``CompressionTypes`` for possible values.

Returns: file handle with a ``close`` function for the user to use.


def _load_saved_model(self, saved_model_dir, saved_model_tags):
    graph = _ops.Graph()
    saved_model = _loader_impl.SavedModelLoader(saved_model_dir)
    saved_model.load_graph(graph, tags=saved_model_tags)
    meta_graph = saved_model.get_meta_graph_def_from_tags(saved_model_tags)
    graph_def = meta_graph.graph_def
    signature_def = meta_graph.signature_def[_signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
    input_tensors = [graph.get_tensor_by_name(signature_def.inputs[key].name) for key in signature_def.inputs]
    output_tensors = [graph.get_tensor_by_name(signature_def.outputs[key].name) for key in signature_def.outputs]
    return (graph_def, input_tensors, output_tensors)

Load graph_def from saved model with the default serving signature key.

Args:
  saved_model_dir: Directory of the SavedModel.
  saved_model_tags: Set of tags identifying the MetaGraphDef within the
    SavedModel to analyze.

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


def _save_model_and_copy_assets(exported_model: exported_model_pb2.ExportedModel, src_saved_model_path: str, dst_saved_model_path: str, signature_def_map: Mapping[str, meta_graph_pb2.SignatureDef], tags: Collection[str]) -> bool:
    save_model.save_model_v1(exported_model.graph_def, dst_saved_model_path, signature_def_map, tags, init_op_name=exported_model.init_node_name, saver_def=_get_saver_def_or_none(exported_model), checkpoint_dir=exported_model.checkpoint_dir, function_aliases=exported_model.function_aliases, asset_file_defs=exported_model.asset_file_defs)
    _copy_assets(src_saved_model_path, dst_saved_model_path)
    return True

Saves the model and copies the assets from the source model.

Args:
  exported_model: ExportedModel to save.
  src_saved_model_path: Path to the source SavedModel. This will be used to
    copy the asset files to `dst_saved_model_path`.
  dst_saved_model_path: Destination path to save the exported model.
  signature_def_map: Signature key -> SignatureDef mapping.
  tags: Tags to attach to the saved MetaGraphDef.

Returns:
  `True` upon successfully saving the model.


def implied_vol(*, prices, strikes, expiries, spots=None, forwards=None, discount_factors=None, is_call_options=None, validate_args=False, polya_factor=2 / np.pi, dtype=None, name=None):
    if (spots is None) == (forwards is None):
        raise ValueError('Either spots or forwards must be supplied but not both.')
    name = name or 'implied_vol'
    with tf.name_scope(name):
        prices = tf.convert_to_tensor(prices, dtype=dtype, name='prices')
        dtype = prices.dtype
        strikes = tf.convert_to_tensor(strikes, dtype=dtype, name='strikes')
        expiries = tf.convert_to_tensor(expiries, dtype=dtype, name='expiries')
        if discount_factors is None:
            discount_factors = tf.convert_to_tensor(1.0, dtype=dtype, name='discount_factors')
        else:
            discount_factors = tf.convert_to_tensor(discount_factors, dtype=dtype, name='discount_factors')
        if forwards is not None:
            forwards = tf.convert_to_tensor(forwards, dtype=dtype, name='forwards')
        else:
            spots = tf.convert_to_tensor(spots, dtype=dtype, name='spots')
            forwards = spots / discount_factors
        control_inputs = None
        if validate_args:
            control_inputs = _validate_args_control_deps(prices, forwards, strikes, expiries, discount_factors, is_call_options)
        with tf.compat.v1.control_dependencies(control_inputs):
            adjusted_strikes = strikes * discount_factors
            normalized_prices = prices / adjusted_strikes
            normalized_forwards = forwards / strikes
            return _approx_implied_vol_polya(normalized_prices, normalized_forwards, expiries, is_call_options, polya_factor)

Approximates the implied vol using the Stefanica-Radiocic algorithm.

Finds an approximation to the implied vol using the Polya approximation for
the Normal CDF. This algorithm was described by Stefanica and Radiocic in
ref [1]. They show that if the Normal CDFs appearing in the Black Scholes
formula for the option price are replaced with Polya's approximation, the
implied vol can be solved for analytically. The Polya approximation produces
absolute errors of less than 0.003 and the resulting implied vol is fairly
close to the true value. For practical purposes, this may not be accurate
enough so this result should be used as a starting point for some method with
controllable tolerance (e.g. a root finder).

#### References:
[1]: Dan Stefanica and Rados Radoicic. An explicit implied volatility formula.
  International Journal of Theoretical and Applied Finance,
  Vol. 20, no. 7, 2017.
  https://papers.ssrn.com/sol3/papers.cfm?abstract_id=2908494
[2]: Omar Eidous, Samar Al-Salman. One-term approximation for Normal
  distribution function. Mathematics and Statistics 4(1), 2016.
  http://www.hrpub.org/download/20160229/MS2-13405192.pdf

Args:
  prices: A real `Tensor` of any shape. The prices of the options whose
    implied vol is to be calculated.
  strikes: A real `Tensor` of the same dtype as `prices` and a shape that
    broadcasts with `prices`. The strikes of the options.
  expiries: A real `Tensor` of the same dtype as `prices` and a shape that
    broadcasts with `prices`. The expiry for each option. The units should
    be such that `expiry * volatility**2` is dimensionless.
  spots: A real `Tensor` of any shape that broadcasts to the shape
    of the `prices`. The current spot price of the underlying. Either this
    argument or the `forwards` (but not both) must be supplied.
  forwards: A real `Tensor` of any shape that broadcasts to the shape of
    `prices`. The forwards to maturity. Either this argument or the `spots`
    must be supplied but both must not be supplied.
  discount_factors: An optional real `Tensor` of same dtype as the `prices`.
    If not None, these are the discount factors to expiry (i.e. e^(-rT)).
    If None, no discounting is applied (i.e. it is assumed that the
    undiscounted option prices are provided ). If `spots` is supplied and
    `discount_factors` is not None then this is also used to compute the
    forwards to expiry.
    Default value: None, equivalent to discount factors = 1.
  is_call_options: A boolean `Tensor` of a shape compatible with `prices`.
    Indicates whether the option is a call (if True) or a put (if False).
    If not supplied, call options are assumed.
  validate_args: A Python bool. If True, indicates that arguments should be
    checked for correctness before performing the computation. The checks
    performed are: (1) Forwards/spots and strikes are positive. (2) The prices
      satisfy the arbitrage bounds (i.e. for call options, checks the
      inequality `max(F-K, 0) <= Price <= F` and for put options, checks that
      `max(K-F, 0) <= Price <= K`.). (3) Checks that the prices are not too
      close to the bounds. It is numerically unstable to compute the implied
      vols from options too far in the money or out of the money.
    Default value: False
  polya_factor: A real scalar. The coefficient to use in the
    approximation for the Normal CDF. The approximation is: `N(x) ~ 0.5 + 0.5
      * sign(x) * sqrt[ 1 - exp(-k * x**2) ]` where `k` is the coefficient
      supplied with `polya_factor`. The original Polya approximation has the
      value `2 / pi` and this is approximation used in Ref [1]. However, as
      described in Ref [2], a slightly more accurate approximation is achieved
      if we use the value of `k=5/8`).
  dtype: `tf.Dtype` to use when converting arguments to `Tensor`s. If not
    supplied, the default TensorFlow conversion will take place. Note that
    this argument does not do any casting for `Tensor`s or numpy arrays.
    Default value: None.
  name: (Optional) Python str. The name prefixed to the ops created by this
    function. If not supplied, the default name 'implied_vol' is
    used.
    Default value: None

Returns:
  implied_vols: A `Tensor` of the same dtype as `prices` and shape as the
    common broadcasted shape of `(prices, spots/forwards, strikes, expiries)`.
    The approximate implied total volatilities computed using the Polya
    approximation method.

Raises:
  ValueError: If both `forwards` and `spots` are supplied or if neither is
    supplied.


class DonutFastImageProcessorKwargs(DefaultFastImageProcessorKwargs):
    do_thumbnail: Optional[bool]
    do_align_long_axis: Optional[bool]
    do_pad: Optional[bool]

Args:
    do_thumbnail (`bool`, *optional*, defaults to `self.do_thumbnail`):
        Whether to resize the image using thumbnail method.
    do_align_long_axis (`bool`, *optional*, defaults to `self.do_align_long_axis`):
        Whether to align the long axis of the image with the long axis of `size` by rotating by 90 degrees.
    do_pad (`bool`, *optional*, defaults to `self.do_pad`):
        Whether to pad the image. If `random_padding` is set to `True`, each image is padded with a random
        amount of padding on each size, up to the largest image size in the batch. Otherwise, all images are
        padded to the largest image size in the batch.


class TestResult:

    def __init__(self):
        self.requested = []
        self.failed = []
        self.executed = []
        self.passed = []
        self.skipped = []
        self.error = []
        self.controller_info = []

    def __add__(self, r):
        """Overrides '+' operator for TestResult class.

    The add operator merges two TestResult objects by concatenating all of
    their lists together.

    Args:
      r: another instance of TestResult to be added

    Returns:
      A TestResult instance that's the sum of two TestResult instances.
    """
        if not isinstance(r, TestResult):
            raise TypeError('Operand %s of type %s is not a TestResult.' % (r, type(r)))
        sum_result = TestResult()
        for name in sum_result.__dict__:
            r_value = getattr(r, name)
            l_value = getattr(self, name)
            if isinstance(r_value, list):
                setattr(sum_result, name, l_value + r_value)
        return sum_result

    def add_record(self, record):
        """Adds a test record to test result.

    A record is considered executed once it's added to the test result.

    Adding the record finalizes the content of a record, so no change
    should be made to the record afterwards.

    Args:
      record: A test record object to add.
    """
        record.update_record()
        if record.result == TestResultEnums.TEST_RESULT_SKIP:
            self.skipped.append(record)
            return
        self.executed.append(record)
        if record.result == TestResultEnums.TEST_RESULT_FAIL:
            self.failed.append(record)
        elif record.result == TestResultEnums.TEST_RESULT_PASS:
            self.passed.append(record)
        else:
            self.error.append(record)

    def add_controller_info_record(self, controller_info_record):
        """Adds a controller info record to results.

    This can be called multiple times for each test class.

    Args:
      controller_info_record: ControllerInfoRecord object to be added to
        the result.
    """
        self.controller_info.append(controller_info_record)

    def add_class_error(self, test_record):
        """Add a record to indicate a test class has failed before any test
    could execute.

    This is only called before any test is actually executed. So it only
    adds an error entry that describes why the class failed to the tally
    and does not affect the total number of tests requrested or exedcuted.

    Args:
      test_record: A TestResultRecord object for the test class.
    """
        test_record.update_record()
        self.error.append(test_record)

    def is_test_executed(self, test_name):
        """Checks if a specific test has been executed.

    Args:
      test_name: string, the name of the test to check.

    Returns:
      True if the test has been executed according to the test result,
      False otherwise.
    """
        for record in self.executed:
            if record.test_name == test_name:
                return True
        return False

    def _count_eventually_passing_retries(self):
        """Counts the number of retry iterations that eventually passed.

    If a test is retried and eventually passed, all the associated non-passing
    iterations should not be considered when devising the final state of the
    test run.

    Returns:
      Int, the number that should be subtracted from the result altering error
      counts.
    """
        count = 0
        for record in self.passed:
            r = record
            while r.parent is not None and r.parent[1] == TestParentType.RETRY:
                count += 1
                r = r.parent[0]
        return count

    @property
    def is_all_pass(self):
        """True if no tests failed or threw errors, False otherwise."""
        num_of_result_altering_errors = len(self.failed) + len(self.error) - self._count_eventually_passing_retries()
        if num_of_result_altering_errors == 0:
            return True
        return False

    def requested_test_names_dict(self):
        """Gets the requested test names of a test run in a dict format.

    Note a test can be requested multiple times, so there can be duplicated
    values

    Returns:
      A dict with a key and the list of strings.
    """
        return {'Requested Tests': copy.deepcopy(self.requested)}

    def summary_str(self):
        """Gets a string that summarizes the stats of this test result.

    The summary provides the counts of how many tests fall into each
    category, like 'Passed', 'Failed' etc.

    Format of the string is:
      Requested <int>, Executed <int>, ...

    Returns:
      A summary string of this test result.
    """
        kv_pairs = ['%s %d' % (k, v) for k, v in self.summary_dict().items()]
        msg = ', '.join(sorted(kv_pairs))
        return msg

    def summary_dict(self):
        """Gets a dictionary that summarizes the stats of this test result.

    The summary provides the counts of how many tests fall into each
    category, like 'Passed', 'Failed' etc.

    Returns:
      A dictionary with the stats of this test result.
    """
        d = {}
        d['Requested'] = len(self.requested)
        d['Executed'] = len(self.executed)
        d['Passed'] = len(self.passed)
        d['Failed'] = len(self.failed)
        d['Skipped'] = len(self.skipped)
        d['Error'] = len(self.error)
        return d

A class that contains metrics of a test run.

This class is essentially a container of TestResultRecord objects.

Attributes:
  requested: A list of strings, each is the name of a test requested
    by user.
  failed: A list of records for tests failed.
  executed: A list of records for tests that were actually executed.
  passed: A list of records for tests passed.
  skipped: A list of records for tests skipped.
  error: A list of records for tests with error result token.
  controller_info: list of ControllerInfoRecord.


def _get_traces(self, lineno, ops, symbol, maxmatch=-1, num_lines=1):
    if not isinstance(symbol, _SymbolMatcher):
        symbol = _SymbolMatcher.from_one_match(symbol)
    for tr in itertools.chain.from_iterable((self.source.traces[line] for line in range(lineno, lineno + num_lines))):
        if maxmatch == 0:
            break
        m_matched = self._matched
        assert m_matched is not None
        if id(tr) not in m_matched and tr.op in ops and symbol.match(tr.symbol):
            maxmatch -= 1
            m_matched.add(id(tr))
            yield tr

Yields matching traces.

Args:
  lineno: A starting line number.
  ops: A list of opcode names to match on.
  symbol: A symbol or _SymbolMatcher instance to match on.
  maxmatch: The maximum number of traces to yield. -1 for no maximum.
  num_lines: The number of consecutive lines to search.


def forward(self, input_ids: torch.Tensor, cache_position: torch.Tensor):
    _, seqlen = input_ids.shape
    position_ids = cache_position.unsqueeze(0)
    past_key_values = self.static_cache
    outs = self.model(input_ids=input_ids, attention_mask=None, position_ids=position_ids, cache_position=cache_position, past_key_values=past_key_values, use_cache=True)
    return outs.logits

Forward pass of the module, which is compatible with the ExecuTorch runtime.

Args:
    input_ids (`torch.Tensor`): Tensor representing current input token id to the module.
    cache_position (`torch.Tensor`): Tensor representing current input position in the cache.

Returns:
    torch.Tensor: Logits output from the model.

This forward adapter serves two primary purposes:

1. **Making the Model `torch.export`-Compatible**:
    The adapter hides unsupported objects, such as the `Cache`, from the graph inputs and outputs,
    enabling the model to be exportable using `torch.export` without encountering issues.

2. **Ensuring Compatibility with `ExecuTorch` runtime**:
    The adapter matches the model's forward signature with that in `executorch/extension/llm/runner`,
    ensuring that the exported model can be executed in `ExecuTorch` out-of-the-box.


def get_arg_value(node, arg_name, arg_pos=None):
    if arg_name is not None:
        for kw in node.keywords:
            if kw.arg == arg_name:
                return (True, kw.value)
    if arg_pos is not None:
        idx = 0
        for arg in node.args:
            if sys.version_info[:2] >= (3, 5) and isinstance(arg, ast.Starred):
                continue
            if idx == arg_pos:
                return (True, arg)
            idx += 1
    return (False, None)

Get the value of an argument from a ast.Call node.

This function goes through the positional and keyword arguments to check
whether a given argument was used, and if so, returns its value (the node
representing its value).

This cannot introspect *args or **args, but it safely handles *args in
Python3.5+.

Args:
  node: The ast.Call node to extract arg values from.
  arg_name: The name of the argument to extract.
  arg_pos: The position of the argument (in case it's passed as a positional
    argument).

Returns:
  A tuple (arg_present, arg_value) containing a boolean indicating whether
  the argument is present, and its value in case it is.


def get_metadata(self, key: str, per_trial: bool=True) -> Optional[Any]:

Gets metadata for current trial or current sampling.

Args:
  key: A string as key to metadata.
  per_trial: If True, the key is retrieved per curent trial. Otherwise, it
    is retrieved per current sampling.

Returns:
  A value that can be deserialized by `pg.from_json_str`.


def rescale(self, image: np.ndarray, rescale_factor: float, data_format: Optional[Union[str, ChannelDimension]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None) -> np.ndarray:
    return rescale(image, rescale_factor, data_format=data_format, input_data_format=input_data_format)

Rescale the image by the given factor. image = image * rescale_factor.

Args:
    image (`np.ndarray`):
        Image to rescale.
    rescale_factor (`float`):
        The value to use for rescaling.
    data_format (`str` or `ChannelDimension`, *optional*):
        The channel dimension format for the output image. If unset, the channel dimension format of the input
        image is used. Can be one of:
        - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
        - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
    input_data_format (`str` or `ChannelDimension`, *optional*):
        The channel dimension format for the input image. If unset, is inferred from the input image. Can be
        one of:
        - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
        - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.


def _batch_prepare_for_model(self, batch_ids_pairs: List[Tuple[List[int], None]], batch_entity_ids_pairs: List[Tuple[Optional[List[int]], Optional[List[int]]]], batch_entity_token_spans_pairs: List[Tuple[Optional[List[Tuple[int, int]]], Optional[List[Tuple[int, int]]]]], add_special_tokens: bool=True, padding_strategy: PaddingStrategy=PaddingStrategy.DO_NOT_PAD, truncation_strategy: TruncationStrategy=TruncationStrategy.DO_NOT_TRUNCATE, max_length: Optional[int]=None, max_entity_length: Optional[int]=None, stride: int=0, pad_to_multiple_of: Optional[int]=None, padding_side: Optional[str]=None, return_tensors: Optional[str]=None, return_token_type_ids: Optional[bool]=None, return_attention_mask: Optional[bool]=None, return_overflowing_tokens: bool=False, return_special_tokens_mask: bool=False, return_length: bool=False, verbose: bool=True) -> BatchEncoding:
    batch_outputs = {}
    for input_ids, entity_ids, entity_token_span_pairs in zip(batch_ids_pairs, batch_entity_ids_pairs, batch_entity_token_spans_pairs):
        first_ids, second_ids = input_ids
        first_entity_ids, second_entity_ids = entity_ids
        first_entity_token_spans, second_entity_token_spans = entity_token_span_pairs
        outputs = self.prepare_for_model(first_ids, second_ids, entity_ids=first_entity_ids, pair_entity_ids=second_entity_ids, entity_token_spans=first_entity_token_spans, pair_entity_token_spans=second_entity_token_spans, add_special_tokens=add_special_tokens, padding=PaddingStrategy.DO_NOT_PAD.value, truncation=truncation_strategy.value, max_length=max_length, max_entity_length=max_entity_length, stride=stride, pad_to_multiple_of=None, padding_side=None, return_attention_mask=False, return_token_type_ids=return_token_type_ids, return_overflowing_tokens=return_overflowing_tokens, return_special_tokens_mask=return_special_tokens_mask, return_length=return_length, return_tensors=None, prepend_batch_axis=False, verbose=verbose)
        for key, value in outputs.items():
            if key not in batch_outputs:
                batch_outputs[key] = []
            batch_outputs[key].append(value)
    batch_outputs = self.pad(batch_outputs, padding=padding_strategy.value, max_length=max_length, pad_to_multiple_of=pad_to_multiple_of, padding_side=padding_side, return_attention_mask=return_attention_mask)
    batch_outputs = BatchEncoding(batch_outputs, tensor_type=return_tensors)
    return batch_outputs

Prepares a sequence of input id, or a pair of sequences of inputs ids so that it can be used by the model. It
adds special tokens, truncates sequences if overflowing while taking into account the special tokens and
manages a moving window (with user defined stride) for overflowing tokens

Args:
    batch_ids_pairs: list of tokenized input ids or input ids pairs
    batch_entity_ids_pairs: list of entity ids or entity ids pairs
    batch_entity_token_spans_pairs: list of entity spans or entity spans pairs
    max_entity_length: The maximum length of the entity sequence.


def sparse_to_dense(sparse_indices, output_shape, sparse_values, default_value=0, validate_indices=True, name=None):
    return gen_sparse_ops.sparse_to_dense(sparse_indices, output_shape, sparse_values, default_value=default_value, validate_indices=validate_indices, name=name)

Converts a sparse representation into a dense tensor.

Builds an array `dense` with shape `output_shape` such that

```python
# If sparse_indices is scalar
dense[i] = (i == sparse_indices ? sparse_values : default_value)

# If sparse_indices is a vector, then for each i
dense[sparse_indices[i]] = sparse_values[i]

# If sparse_indices is an n by d matrix, then for each i in [0, n)
dense[sparse_indices[i][0], ..., sparse_indices[i][d-1]] = sparse_values[i]
```

All other values in `dense` are set to `default_value`.  If `sparse_values`
is a scalar, all sparse indices are set to this single value.

Indices should be sorted in lexicographic order, and indices must not
contain any repeats. If `validate_indices` is True, these properties
are checked during execution.

Args:
  sparse_indices: A 0-D, 1-D, or 2-D `Tensor` of type `int32` or `int64`.
    `sparse_indices[i]` contains the complete index where `sparse_values[i]`
    will be placed.
  output_shape: A 1-D `Tensor` of the same type as `sparse_indices`.  Shape
    of the dense output tensor.
  sparse_values: A 0-D or 1-D `Tensor`.  Values corresponding to each row of
    `sparse_indices`, or a scalar value to be used for all sparse indices.
  default_value: A 0-D `Tensor` of the same type as `sparse_values`.  Value
    to set for indices not specified in `sparse_indices`.  Defaults to zero.
  validate_indices: A boolean value.  If True, indices are checked to make
    sure they are sorted in lexicographic order and that there are no repeats.
  name: A name for the operation (optional).

Returns:
  Dense `Tensor` of shape `output_shape`.  Has the same type as
  `sparse_values`.


def _slice_length(value_length, slice_key):
    zeros = array_ops.zeros(value_length, dtype=dtypes.bool)
    return array_ops.size(zeros[slice_key], out_type=value_length.dtype)

Computes the number of elements in a slice of a value with a given length.

Returns the equivalent of: `len(range(value_length)[slice_key])`

Args:
  value_length: Scalar int `Tensor`: the length of the value being sliced.
  slice_key: A `slice` object used to slice elements from the value.

Returns:
  The number of elements in the sliced value.


def _get_compile_args(self, user_metrics=True):
    self._assert_compile_was_called()
    saved_metrics = self.compiled_metrics._user_metrics
    saved_weighted_metrics = self.compiled_metrics._user_weighted_metrics
    if not user_metrics:
        if saved_metrics is not None:
            saved_metrics = self.compiled_metrics._metrics
        if saved_weighted_metrics is not None:
            saved_weighted_metrics = self.compiled_metrics._weighted_metrics
    compile_args = {'optimizer': self.optimizer, 'loss': self.compiled_loss._user_losses, 'metrics': saved_metrics, 'weighted_metrics': saved_weighted_metrics, 'loss_weights': self.compiled_loss._user_loss_weights}
    return compile_args

Used for saving or cloning a Model.

Args:
  user_metrics: Whether to return user-supplied metrics or `Metric` objects.
    Defaults to returning the user-supplied metrics.

Returns:
  Dictionary of arguments that were used when compiling the model.


class InputSpec(object):

    def __init__(self, dtype=None, shape=None, ndim=None, max_ndim=None, min_ndim=None, axes=None, allow_last_axis_squeeze=False, name=None):
        self.dtype = dtypes.as_dtype(dtype).name if dtype is not None else None
        shape = tensor_shape.TensorShape(shape)
        if shape.rank is None:
            shape = None
        else:
            shape = tuple(shape.as_list())
        if shape is not None:
            self.ndim = len(shape)
            self.shape = shape
        else:
            self.ndim = ndim
            self.shape = None
        self.max_ndim = max_ndim
        self.min_ndim = min_ndim
        self.name = name
        self.allow_last_axis_squeeze = allow_last_axis_squeeze
        try:
            axes = axes or {}
            self.axes = {int(k): axes[k] for k in axes}
        except (ValueError, TypeError):
            raise TypeError('The keys in axes must be integers.')
        if self.axes and (self.ndim is not None or self.max_ndim is not None):
            max_dim = (self.ndim if self.ndim else self.max_ndim) - 1
            max_axis = max(self.axes)
            if max_axis > max_dim:
                raise ValueError('Axis {} is greater than the maximum allowed value: {}'.format(max_axis, max_dim))

    def __repr__(self):
        spec = ['dtype=' + str(self.dtype) if self.dtype else '', 'shape=' + str(self.shape) if self.shape else '', 'ndim=' + str(self.ndim) if self.ndim else '', 'max_ndim=' + str(self.max_ndim) if self.max_ndim else '', 'min_ndim=' + str(self.min_ndim) if self.min_ndim else '', 'axes=' + str(self.axes) if self.axes else '']
        return 'InputSpec(%s)' % ', '.join((x for x in spec if x))

    def get_config(self):
        return {'dtype': self.dtype, 'shape': self.shape, 'ndim': self.ndim, 'max_ndim': self.max_ndim, 'min_ndim': self.min_ndim, 'axes': self.axes}

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

Specifies the rank, dtype and shape of every input to a layer.

Layers can expose (if appropriate) an `input_spec` attribute:
an instance of `InputSpec`, or a nested structure of `InputSpec` instances
(one per input tensor). These objects enable the layer to run input
compatibility checks for input structure, input rank, input shape, and
input dtype.

A None entry in a shape is compatible with any dimension,
a None shape is compatible with any shape.

Args:
  dtype: Expected DataType of the input.
  shape: Shape tuple, expected shape of the input
    (may include None for unchecked axes). Includes the batch size.
  ndim: Integer, expected rank of the input.
  max_ndim: Integer, maximum rank of the input.
  min_ndim: Integer, minimum rank of the input.
  axes: Dictionary mapping integer axes to
    a specific dimension value.
  allow_last_axis_squeeze: If True, then allow inputs of rank N+1 as long
    as the last axis of the input is 1, as well as inputs of rank N-1
    as long as the last axis of the spec is 1.
  name: Expected key corresponding to this input when passing data as
    a dictionary.

Example:

```python
class MyLayer(Layer):
    def __init__(self):
        super(MyLayer, self).__init__()
        # The layer will accept inputs with shape (?, 28, 28) & (?, 28, 28, 1)
        # and raise an appropriate error message otherwise.
        self.input_spec = InputSpec(
            shape=(None, 28, 28, 1),
            allow_last_axis_squeeze=True)
```


def fake_quant_with_min_max_vars(inputs, min_vals, max_vals, num_bits=8, narrow_range=False, axis=None):
    if any_symbolic_tensors((inputs,)):
        return FakeQuantWithMinMaxVars().symbolic_call(inputs, min_vals, max_vals)
    inputs = ops.convert_to_tensor(inputs)
    min_vals = ops.convert_to_tensor(min_vals)
    max_vals = ops.convert_to_tensor(max_vals)
    num_bits = int(num_bits)
    if axis is not None:
        axis = canonicalize_axis(axis, inputs.ndim)
    if backend.backend() == 'tensorflow':
        import tensorflow as tf
        dtype = backend.standardize_dtype(inputs.dtype)
        if axis is None:
            outputs = tf.quantization.fake_quant_with_min_max_vars(ops.cast(inputs, 'float32'), ops.cast(ops.reshape(min_vals, ()), 'float32'), ops.cast(ops.reshape(max_vals, ()), 'float32'), num_bits=num_bits, narrow_range=narrow_range)
            return ops.cast(outputs, dtype=dtype)
        else:
            last_axis = inputs.ndim - 1
            inputs = ops.swapaxes(inputs, axis, last_axis)
            outputs = tf.quantization.fake_quant_with_min_max_vars_per_channel(ops.cast(inputs, 'float32'), ops.cast(min_vals, 'float32'), ops.cast(max_vals, 'float32'), num_bits=num_bits, narrow_range=narrow_range)
            outputs = ops.cast(outputs, dtype=dtype)
            return ops.swapaxes(outputs, last_axis, axis)

    @ops.custom_gradient
    def _fake_quant_with_min_max_vars_per_channel(x, min_val, max_val):
        dtype = backend.standardize_dtype(x.dtype)
        nudged_min, nudged_max, scale, inv_scale = adjust_and_nudge(min_val, max_val, num_bits, narrow_range)
        quant_zero = ops.floor(ops.add(ops.multiply(-nudged_min, inv_scale), 0.5))
        x_clamped = ops.clip(x, ops.cast(nudged_min, x.dtype), ops.cast(nudged_max, x.dtype))
        x_clamped_shifted = ops.subtract(x_clamped, nudged_min)
        result = ops.multiply(ops.floor(ops.add(ops.subtract(ops.multiply(x_clamped_shifted, inv_scale), quant_zero), 0.5)), scale)
        result = ops.cast(result, dtype=dtype)
        masks = ops.logical_and(ops.greater_equal(x, nudged_min), ops.less_equal(x, nudged_max))

        def grad(*args, upstream=None):
            if upstream is None:
                upstream, = args
            dx = ops.where(masks, upstream, 0.0)
            axes = [i for i in range(len(dx.shape)) if i != axis]
            min_mask = ops.less_equal(x, nudged_min)
            grad_min = ops.where(min_mask, upstream, 0.0)
            if axis is not None:
                grad_min = ops.sum(grad_min, axis=axes)
            else:
                grad_min = ops.sum(grad_min)
            max_mask = ops.greater_equal(x, nudged_max)
            grad_max = ops.where(max_mask, upstream, 0.0)
            if axis is not None:
                grad_max = ops.sum(grad_max, axis=axes)
            else:
                grad_max = ops.sum(grad_max)
            return (dx, grad_min, grad_max)
        return (result, grad)
    return _fake_quant_with_min_max_vars_per_channel(inputs, min_vals, max_vals)

Perform per-tensor or per-channel fake quantization.

`[min_vals, max_vals]` define the clamping range for the `inputs`.

The `inputs` are quantized into the quantization range:
- `[0, 2^num_bits - 1]` when `narrow_range=False`
- `[1, 2^num_bits - 1]` when `narrow_range=True`

After quantization, the values are dequantized and output as floats within
the `[min_vals, max_vals]` interval.

This operation supports gradient computation, allowing `min_vals` and
`max_vals` to be trained.

Args:
    inputs: Input Keras tensor of float dtype.
    min_vals: A global minimum scalar or a per-channel minimum tensor.
    max_vals: A global maximum scalar or a per-channel maximum tensor.
    num_bits: Quantization bit width (e.g., `8` for int8). Defaults to `8`.
    narrow_range: Whether to use narrow quantization range. Defaults to
        `False`.
    axis: Axis along which to perform per-channel quantization. If `None`,
          per-tensor quantization is performed. Defaults to `None`.

Returns:
    Tensor: A Keras tensor with fake quantization applied.


class ModernBertConfig(PretrainedConfig):
    model_type = 'modernbert'
    keys_to_ignore_at_inference = ['past_key_values']

    def __init__(self, vocab_size=50368, hidden_size=768, intermediate_size=1152, num_hidden_layers=22, num_attention_heads=12, hidden_activation='gelu', max_position_embeddings=8192, initializer_range=0.02, initializer_cutoff_factor=2.0, norm_eps=1e-05, norm_bias=False, pad_token_id=50283, eos_token_id=50282, bos_token_id=50281, cls_token_id=50281, sep_token_id=50282, global_rope_theta=160000.0, attention_bias=False, attention_dropout=0.0, global_attn_every_n_layers=3, local_attention=128, local_rope_theta=10000.0, embedding_dropout=0.0, mlp_bias=False, mlp_dropout=0.0, decoder_bias=True, classifier_pooling: Literal['cls', 'mean']='cls', classifier_dropout=0.0, classifier_bias=False, classifier_activation='gelu', deterministic_flash_attn=False, sparse_prediction=False, sparse_pred_ignore_index=-100, reference_compile=None, repad_logits_with_grad=False, **kwargs):
        super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, cls_token_id=cls_token_id, sep_token_id=sep_token_id, **kwargs)
        self.vocab_size = vocab_size
        self.max_position_embeddings = max_position_embeddings
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.initializer_range = initializer_range
        self.initializer_cutoff_factor = initializer_cutoff_factor
        self.norm_eps = norm_eps
        self.norm_bias = norm_bias
        self.global_rope_theta = global_rope_theta
        self.attention_bias = attention_bias
        self.attention_dropout = attention_dropout
        self.hidden_activation = hidden_activation
        self.global_attn_every_n_layers = global_attn_every_n_layers
        self.local_attention = local_attention
        self.local_rope_theta = local_rope_theta
        self.embedding_dropout = embedding_dropout
        self.mlp_bias = mlp_bias
        self.mlp_dropout = mlp_dropout
        self.decoder_bias = decoder_bias
        self.classifier_pooling = classifier_pooling
        self.classifier_dropout = classifier_dropout
        self.classifier_bias = classifier_bias
        self.classifier_activation = classifier_activation
        self.deterministic_flash_attn = deterministic_flash_attn
        self.sparse_prediction = sparse_prediction
        self.sparse_pred_ignore_index = sparse_pred_ignore_index
        self.reference_compile = reference_compile
        self.repad_logits_with_grad = repad_logits_with_grad
        if self.classifier_pooling not in ['cls', 'mean']:
            raise ValueError(f'Invalid value for `classifier_pooling`, should be either "cls" or "mean", but is {self.classifier_pooling}.')

    def to_dict(self):
        output = super().to_dict()
        output.pop('reference_compile', None)
        return output

This is the configuration class to store the configuration of a [`ModernBertModel`]. It is used to instantiate an ModernBert
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 ModernBERT-base.
e.g. [answerdotai/ModernBERT-base](https://huggingface.co/answerdotai/ModernBERT-base)

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

Args:
    vocab_size (`int`, *optional*, defaults to 50368):
        Vocabulary size of the ModernBert model. Defines the number of different tokens that can be represented by the
        `inputs_ids` passed when calling [`ModernBertModel`]
    hidden_size (`int`, *optional*, defaults to 768):
        Dimension of the hidden representations.
    intermediate_size (`int`, *optional*, defaults to 1152):
        Dimension of the MLP representations.
    num_hidden_layers (`int`, *optional*, defaults to 22):
        Number of hidden layers in the Transformer decoder.
    num_attention_heads (`int`, *optional*, defaults to 12):
        Number of attention heads for each attention layer in the Transformer decoder.
    hidden_activation (`str` or `function`, *optional*, defaults to `"gelu"`):
        The non-linear activation function (function or string) in the decoder. Will default to `"gelu"`
        if not specified.
    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.
    initializer_cutoff_factor (`float`, *optional*, defaults to 2.0):
        The cutoff factor for the truncated_normal_initializer for initializing all weight matrices.
    norm_eps (`float`, *optional*, defaults to 1e-05):
        The epsilon used by the rms normalization layers.
    norm_bias (`bool`, *optional*, defaults to `False`):
        Whether to use bias in the normalization layers.
    pad_token_id (`int`, *optional*, defaults to 50283):
        Padding token id.
    eos_token_id (`int`, *optional*, defaults to 50282):
        End of stream token id.
    bos_token_id (`int`, *optional*, defaults to 50281):
        Beginning of stream token id.
    cls_token_id (`int`, *optional*, defaults to 50281):
        Classification token id.
    sep_token_id (`int`, *optional*, defaults to 50282):
        Separation token id.
    global_rope_theta (`float`, *optional*, defaults to 160000.0):
        The base period of the global RoPE embeddings.
    attention_bias (`bool`, *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.
    global_attn_every_n_layers (`int`, *optional*, defaults to 3):
        The number of layers between global attention layers.
    local_attention (`int`, *optional*, defaults to 128):
        The window size for local attention.
    local_rope_theta (`float`, *optional*, defaults to 10000.0):
        The base period of the local RoPE embeddings.
    embedding_dropout (`float`, *optional*, defaults to 0.0):
        The dropout ratio for the embeddings.
    mlp_bias (`bool`, *optional*, defaults to `False`):
        Whether to use bias in the MLP layers.
    mlp_dropout (`float`, *optional*, defaults to 0.0):
        The dropout ratio for the MLP layers.
    decoder_bias (`bool`, *optional*, defaults to `True`):
        Whether to use bias in the decoder layers.
    classifier_pooling (`str`, *optional*, defaults to `"cls"`):
        The pooling method for the classifier. Should be either `"cls"` or `"mean"`. In local attention layers, the
        CLS token doesn't attend to all tokens on long sequences.
    classifier_dropout (`float`, *optional*, defaults to 0.0):
        The dropout ratio for the classifier.
    classifier_bias (`bool`, *optional*, defaults to `False`):
        Whether to use bias in the classifier.
    classifier_activation (`str`, *optional*, defaults to `"gelu"`):
        The activation function for the classifier.
    deterministic_flash_attn (`bool`, *optional*, defaults to `False`):
        Whether to use deterministic flash attention. If `False`, inference will be faster but not deterministic.
    sparse_prediction (`bool`, *optional*, defaults to `False`):
        Whether to use sparse prediction for the masked language model instead of returning the full dense logits.
    sparse_pred_ignore_index (`int`, *optional*, defaults to -100):
        The index to ignore for the sparse prediction.
    reference_compile (`bool`, *optional*):
        Whether to compile the layers of the model which were compiled during pretraining. If `None`, then parts of
        the model will be compiled if 1) `triton` is installed, 2) the model is not on MPS, 3) the model is not
        shared between devices, and 4) the model is not resized after initialization. If `True`, then the model may
        be faster in some scenarios.
    repad_logits_with_grad (`bool`, *optional*, defaults to `False`):
        When True, ModernBertForMaskedLM keeps track of the logits' gradient when repadding for output. This only
        applies when using Flash Attention 2 with passed labels. Otherwise output logits always have a gradient.

Examples:

```python
>>> from transformers import ModernBertModel, ModernBertConfig

>>> # Initializing a ModernBert style configuration
>>> configuration = ModernBertConfig()

>>> # Initializing a model from the modernbert-base style configuration
>>> model = ModernBertModel(configuration)

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


def _prepare_4d_causal_attention_mask_with_cache_position(attention_mask: torch.Tensor, sequence_length: int, target_length: int, dtype: torch.dtype, cache_position: torch.Tensor, batch_size: int, **kwargs):
    if attention_mask is not None and attention_mask.dim() == 4:
        causal_mask = attention_mask
    else:
        min_dtype = torch.finfo(dtype).min
        causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device)
        if sequence_length != 1:
            causal_mask = torch.triu(causal_mask, diagonal=1)
        causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1)
        causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
        if attention_mask is not None:
            causal_mask = causal_mask.clone()
            mask_length = attention_mask.shape[-1]
            padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(causal_mask.device)
            padding_mask = padding_mask == 0
            causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(padding_mask, min_dtype)
    return causal_mask

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

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


class EncoderNoRepeatNGramLogitsProcessor(LogitsProcessor):

    def __init__(self, encoder_ngram_size: int, encoder_input_ids: torch.LongTensor):
        if not isinstance(encoder_ngram_size, int) or encoder_ngram_size <= 0:
            raise ValueError(f'`encoder_ngram_size` has to be a strictly positive integer, but is {encoder_ngram_size}')
        self.ngram_size = encoder_ngram_size
        if len(encoder_input_ids.shape) == 1:
            encoder_input_ids = encoder_input_ids.unsqueeze(0)
        self.batch_size = encoder_input_ids.shape[0]
        self.generated_ngrams = _get_ngrams(encoder_ngram_size, encoder_input_ids, self.batch_size)

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        num_hypos = scores.shape[0]
        num_beams = num_hypos // self.batch_size
        cur_len = input_ids.shape[-1]
        scores_processed = scores.clone()
        banned_batch_tokens = [_get_generated_ngrams(self.generated_ngrams[hypo_idx // num_beams], input_ids[hypo_idx], self.ngram_size, cur_len) for hypo_idx in range(num_hypos)]
        for i, banned_tokens in enumerate(banned_batch_tokens):
            scores_processed[i, banned_tokens] = -float('inf')
        return scores_processed

[`LogitsProcessor`] that works similarly to [`NoRepeatNGramLogitsProcessor`], but applied exclusively to prevent
the repetition of n-grams present in the prompt.

It was designed to promote chattiness in a language model, by preventing the generation of n-grams present in
previous conversation rounds.

Args:
    encoder_ngram_size (`int`):
        All ngrams of size `ngram_size` can only occur within the encoder input ids.
    encoder_input_ids (`int`):
        The encoder_input_ids that should not be repeated within the decoder ids.

Examples:

```py
>>> from transformers import AutoTokenizer, AutoModelForCausalLM

>>> model = AutoModelForCausalLM.from_pretrained("bigscience/bloomz-560m")
>>> tokenizer = AutoTokenizer.from_pretrained("bigscience/bloomz-560m")

>>> inputs = tokenizer("Alice: I love cats. What do you love?\nBob:", return_tensors="pt")

>>> # With greedy decoding, we see Bob repeating Alice's opinion. If Bob was a chatbot, it would be a poor one.
>>> outputs = model.generate(**inputs)
>>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
Alice: I love cats. What do you love?
Bob: I love cats. What do you

>>> # With this logits processor, we can prevent Bob from repeating Alice's opinion.
>>> outputs = model.generate(**inputs, encoder_no_repeat_ngram_size=2)
>>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
Alice: I love cats. What do you love?
Bob: My cats are very cute.
```


def softmax_v2(logits, axis=None, name=None):
    if axis is None:
        axis = -1
    return _wrap_2d_function(logits, gen_nn_ops.softmax, axis, name)

Computes softmax activations.

Used for multi-class predictions. The sum of all outputs generated by softmax
is 1.

This function performs the equivalent of

```python
softmax = tf.exp(logits) / tf.reduce_sum(tf.exp(logits), axis, keepdims=True)
```
Example usage:

>>> softmax = tf.nn.softmax([-1, 0., 1.])
>>> softmax
<tf.Tensor: shape=(3,), dtype=float32,
numpy=array([0.09003057, 0.24472848, 0.66524094], dtype=float32)>
>>> sum(softmax)
<tf.Tensor: shape=(), dtype=float32, numpy=1.0>

Args:
  logits: A non-empty `Tensor`. Must be one of the following types: `half`,
    `float32`, `float64`.
  axis: The dimension softmax would be performed on. The default is -1 which
    indicates the last dimension.
  name: A name for the operation (optional).

Returns:
  A `Tensor`. Has the same type and shape as `logits`.

Raises:
  InvalidArgumentError: if `logits` is empty or `axis` is beyond the last
    dimension of `logits`.


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

Remove a Bitbucket Server repository from an given BitbucketServerConfig's connected repositories. This API is experimental.

Args:
  request: (CloudbuildProjectsLocationsBitbucketServerConfigsRemoveBitbucketServerConnectedRepositoryRequest) input message
  global_params: (StandardQueryParameters, default: None) global arguments
Returns:
  (Empty) 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.prefix_tokens + token_ids_0 + self.suffix_tokens
    return self.prefix_tokens + token_ids_0 + token_ids_1 + self.suffix_tokens

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

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

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

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

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


def apply_transform(self, data: OperationInputT, output_column_name: str) -> dict[str, OperationOutputT]:

Define any processing logic in the apply_transform() method.
processing logics are applied on inputs and returns a transformed
output.
Args:
  inputs: input data.


def create_position_ids_from_input_ids(input_ids, padding_idx, past_key_values_length=0):
    mask = input_ids.ne(padding_idx).int()
    incremental_indices = (torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length) * mask
    return incremental_indices.long() + padding_idx

Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
are ignored. This is modified from fairseq's `utils.make_positions`.

Args:
    x: torch.Tensor x:

Returns: torch.Tensor


def __init__(self, key_dtype, value_dtype, default_value, name='MutableHashTable', checkpoint=True, experimental_is_anonymous=False):
    self._default_value = ops.convert_to_tensor(default_value, dtype=value_dtype)
    self._value_shape = self._default_value.get_shape()
    self._checkpoint = checkpoint
    self._key_dtype = key_dtype
    self._value_dtype = value_dtype
    self._name = name
    self._is_anonymous = experimental_is_anonymous
    if not self._is_anonymous:
        self._shared_name = None
        if context.executing_eagerly():
            self._shared_name = 'table_%d' % (ops.uid(),)
    super(MutableHashTable, self).__init__(key_dtype, value_dtype)
    self._resource_handle = self._create_resource()
    if checkpoint:
        saveable = MutableHashTable._Saveable(self, name)
        if not context.executing_eagerly():
            ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, saveable)

Creates an empty `MutableHashTable` object.

Creates a table, the type of its keys and values are specified by key_dtype
and value_dtype, respectively.

Args:
  key_dtype: the type of the key tensors.
  value_dtype: the type of the value tensors.
  default_value: The value to use if a key is missing in the table.
  name: A name for the operation (optional).
  checkpoint: if True, the contents of the table are saved to and restored
    from checkpoints. If `shared_name` is empty for a checkpointed table, it
    is shared using the table node name.
  experimental_is_anonymous: Whether to use anonymous mode for the
    table (default is False). In anonymous mode, the table
    resource can only be accessed via a resource handle. It can't
    be looked up by a name. When all resource handles pointing to
    that resource are gone, the resource will be deleted
    automatically.

Returns:
  A `MutableHashTable` object.

Raises:
  ValueError: If checkpoint is True and no name was specified.


def _nested_type_wrapper(fun):

    def wrapper(pickler, obj):
        if _is_nested_class(obj) and obj.__module__ != '__main__':
            containing_class_and_name = _find_containing_class(obj)
            if containing_class_and_name is not None:
                return pickler.save_reduce(getattr, containing_class_and_name, obj=obj)
        try:
            return fun(pickler, obj)
        except dill.dill.PicklingError:
            return pickler.save_reduce(dill.dill._create_type, (type(obj), obj.__name__, obj.__bases__, _dict_from_mappingproxy(obj.__dict__)), obj=obj)
    return wrapper

A wrapper for the standard pickler handler for class objects.

Args:
  fun: Original pickler handler for type objects.

Returns:
  A wrapper for type objects that handles nested classes.

The wrapper detects if an object being pickled is a nested class object.
For nested class object only it will save the containing class object so
the nested structure is recreated during unpickle.


class Mask2FormerPixelLevelModuleOutput(ModelOutput):
    encoder_last_hidden_state: Optional[torch.FloatTensor] = None
    encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    decoder_last_hidden_state: Optional[torch.FloatTensor] = None
    decoder_hidden_states: Tuple[torch.FloatTensor] = None

Mask2Former's pixel level module output. It returns the output of the encoder (optional) and all hidden states
(multi-scale features) from the `decoder`. By default, the `encoder` is a Swin Backbone and the `decoder` is a
Multi-Scale Deformable Attention based decoder.

The `decoder_last_hidden_state` are the **per-pixel embeddings** while `decoder_hidden_states` refer to multi-scale
feature maps produced using **multi-scaling strategy** defined in the paper.

Args:
    encoder_last_hidden_state (`torch.FloatTensor`):
        Last hidden states (final feature map of shape `(batch_size, num_channels, height, width)`) of the last
        stage of the encoder.
    encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*):
        Tuple of `torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`. Hidden states (also
        called feature maps) of the model at the output of each stage. Returned if output_hidden_states is set to
        True.
    decoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)):
        1/4 scale features from the last Pixel Decoder Layer.
    decoder_hidden_states (`tuple(torch.FloatTensor)`):
        Tuple of `torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`. Hidden states (also
        called feature maps) of the model at the output of each stage.


def pair_wise_sigmoid_focal_loss(inputs: Tensor, labels: Tensor, alpha: float=0.25, gamma: float=2.0) -> Tensor:
    if alpha < 0:
        raise ValueError('alpha must be positive')
    height_and_width = inputs.shape[1]
    criterion = nn.BCEWithLogitsLoss(reduction='none')
    prob = inputs.sigmoid()
    cross_entropy_loss_pos = criterion(inputs, torch.ones_like(inputs))
    focal_pos = (1 - prob) ** gamma * cross_entropy_loss_pos
    focal_pos *= alpha
    cross_entropy_loss_neg = criterion(inputs, torch.zeros_like(inputs))
    focal_neg = prob ** gamma * cross_entropy_loss_neg
    focal_neg *= 1 - alpha
    loss = torch.matmul(focal_pos, labels.T) + torch.matmul(focal_neg, (1 - labels).T)
    return loss / height_and_width

A pair wise version of the focal loss, see `sigmoid_focal_loss` for usage.

Args:
    inputs (`torch.Tensor`):
        A tensor representing a mask.
    labels (`torch.Tensor`):
        A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs
        (0 for the negative class and 1 for the positive class).
    alpha (float, *optional*, defaults to 0.25):
        Weighting factor in range (0,1) to balance positive vs negative examples.
    gamma (float, *optional*, defaults to 2.0):
        Exponent of the modulating factor \\(1 - p_t\\) to balance easy vs hard examples.

Returns:
    `torch.Tensor`: The computed loss between each pairs.


def _exec_one_test_with_retry(self, test_name, test_method, max_count):

    def should_retry(record):
        return record.result in [records.TestResultEnums.TEST_RESULT_FAIL, records.TestResultEnums.TEST_RESULT_ERROR]
    previous_record = self.exec_one_test(test_name, test_method)
    if not should_retry(previous_record):
        return
    for i in range(max_count - 1):
        retry_name = f'{test_name}_retry_{i + 1}'
        new_record = records.TestResultRecord(retry_name, self.TAG)
        new_record.retry_parent = previous_record
        new_record.parent = (previous_record, records.TestParentType.RETRY)
        previous_record = self.exec_one_test(retry_name, test_method, new_record)
        if not should_retry(previous_record):
            break

Executes one test and retry the test if needed.

Repeatedly execute a test case until it passes or the maximum count of
iteration has been reached.

Args:
  test_name: string, Name of the test.
  test_method: function, The test method to execute.
  max_count: int, the maximum number of iterations to execute the test for.


def batch_shape(self):
    return self.shape[:-2]

`TensorShape` of batch dimensions of this `LinearOperator`.

If this operator acts like the batch matrix `A` with
`A.shape = [B1,...,Bb, M, N]`, then this returns
`TensorShape([B1,...,Bb])`, equivalent to `A.shape[:-2]`

Returns:
  `TensorShape`, statically determined, may be undefined.


def find_all_hinted_output_nodes(session=None, graph_def=None):
    if session is not None and graph_def is not None:
        raise ValueError('Provide only one of session and graph_def.')
    hinted_outputs_nodes = []
    if session is not None:
        hints = _find_all_hints_in_nodes(session.graph_def.node)
    elif graph_def is not None:
        hints = _find_all_hints_in_nodes(graph_def.node)
    for hint in hints.values():
        _, output_nodes = hint.flattened_inputs_and_outputs()
        hinted_outputs_nodes.extend(output_nodes)
    return hinted_outputs_nodes

Find all Ophints output nodes in the graph.

This is used to get all the output nodes those are ophinted, it is important
for operation like convert_variables_to_constants keep all ophints structure.
Note: only one of session or graph_def should be used, not both.
Why this can be useful? Some TensorFlow ops (e.g. bidirectional rnn), can
generate multiple outputs for unfused subgraph. If not all output nodes are
consumed, graph optimization can potentially drop the unused nodes and cause
ophints in an invalid states (due to missing ophinted output nodes). So it's
important for us to find all those hinted output nodes and make sure they're
not discarded away.

Args:
  session: A TensorFlow session that contains the graph to convert.
  graph_def: A graph def that we should convert.

Returns:
  A list of OpHints output nodes.
Raises:
  ValueError: If both session and graph_def are provided.


def counter(start=0, step=1, dtype=dtypes.int64, name=None) -> 'DatasetV2':
    from tensorflow.python.data.ops import counter_op
    return counter_op._counter(start, step, dtype, name=name)

Creates a `Dataset` that counts from `start` in steps of size `step`.

Unlike `tf.data.Dataset.range`, which stops at some ending number,
`tf.data.Dataset.counter` produces elements indefinitely.

>>> dataset = tf.data.experimental.Counter().take(5)
>>> [a.item() for a in dataset.as_numpy_iterator()]
[0, 1, 2, 3, 4]
>>> dataset.element_spec
TensorSpec(shape=(), dtype=tf.int64, name=None)
>>> dataset = tf.data.experimental.Counter(dtype=tf.int32)
>>> dataset.element_spec
TensorSpec(shape=(), dtype=tf.int32, name=None)
>>> dataset = tf.data.experimental.Counter(start=2).take(5)
>>> [a.item() for a in dataset.as_numpy_iterator()]
[2, 3, 4, 5, 6]
>>> dataset = tf.data.experimental.Counter(start=2, step=5).take(5)
>>> [a.item() for a in dataset.as_numpy_iterator()]
[2, 7, 12, 17, 22]
>>> dataset = tf.data.experimental.Counter(start=10, step=-1).take(5)
>>> [a.item() for a in dataset.as_numpy_iterator()]
[10, 9, 8, 7, 6]

Args:
  start: (Optional.) The starting value for the counter. Defaults to 0.
  step: (Optional.) The step size for the counter. Defaults to 1.
  dtype: (Optional.) The data type for counter elements. Defaults to
    `tf.int64`.
  name: (Optional.) A name for the tf.data operation.

Returns:
  A `Dataset` of scalar `dtype` elements.


def forward(self, input_ids: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None, global_attention_mask: Optional[torch.Tensor]=None, head_mask: Optional[torch.Tensor]=None, token_type_ids: Optional[torch.Tensor]=None, position_ids: Optional[torch.Tensor]=None, inputs_embeds: Optional[torch.Tensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[Tuple, LongformerBaseModelOutputWithPooling]:
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    if input_ids is not None and inputs_embeds is not None:
        raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time')
    elif input_ids is not None:
        self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
        input_shape = input_ids.size()
    elif inputs_embeds is not None:
        input_shape = inputs_embeds.size()[:-1]
    else:
        raise ValueError('You have to specify either input_ids or inputs_embeds')
    device = input_ids.device if input_ids is not None else inputs_embeds.device
    if attention_mask is None:
        attention_mask = torch.ones(input_shape, device=device)
    if token_type_ids is None:
        token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
    if global_attention_mask is not None:
        attention_mask = self._merge_to_attention_mask(attention_mask, global_attention_mask)
    padding_len, input_ids, attention_mask, token_type_ids, position_ids, inputs_embeds = self._pad_to_window_size(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, inputs_embeds=inputs_embeds, pad_token_id=self.config.pad_token_id)
    extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape)[:, 0, 0, :]
    embedding_output = self.embeddings(input_ids=input_ids, position_ids=position_ids, token_type_ids=token_type_ids, inputs_embeds=inputs_embeds)
    encoder_outputs = self.encoder(embedding_output, attention_mask=extended_attention_mask, head_mask=head_mask, padding_len=padding_len, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    sequence_output = encoder_outputs[0]
    pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
    if not return_dict:
        return (sequence_output, pooled_output) + encoder_outputs[1:]
    return LongformerBaseModelOutputWithPooling(last_hidden_state=sequence_output, pooler_output=pooled_output, hidden_states=encoder_outputs.hidden_states, attentions=encoder_outputs.attentions, global_attentions=encoder_outputs.global_attentions)

global_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
    Mask to decide the attention given on each token, local attention or global attention. Tokens with global
    attention attends to all other tokens, and all other tokens attend to them. This is important for
    task-specific finetuning because it makes the model more flexible at representing the task. For example,
    for classification, the <s> token should be given global attention. For QA, all question tokens should also
    have global attention. Please refer to the [Longformer paper](https://huggingface.co/papers/2004.05150) for more
    details. Mask values selected in `[0, 1]`:

    - 0 for local attention (a sliding window attention),
    - 1 for global attention (tokens that attend to all other tokens, and all other tokens attend to them).

Examples:

```python
>>> import torch
>>> from transformers import LongformerModel, AutoTokenizer

>>> model = LongformerModel.from_pretrained("allenai/longformer-base-4096")
>>> tokenizer = AutoTokenizer.from_pretrained("allenai/longformer-base-4096")

>>> SAMPLE_TEXT = " ".join(["Hello world! "] * 1000)  # long input document
>>> input_ids = torch.tensor(tokenizer.encode(SAMPLE_TEXT)).unsqueeze(0)  # batch of size 1

>>> attention_mask = torch.ones(
...     input_ids.shape, dtype=torch.long, device=input_ids.device
... )  # initialize to local attention
>>> global_attention_mask = torch.zeros(
...     input_ids.shape, dtype=torch.long, device=input_ids.device
... )  # initialize to global attention to be deactivated for all tokens
>>> global_attention_mask[
...     :,
...     [
...         1,
...         4,
...         21,
...     ],
... ] = 1  # Set global attention to random tokens for the sake of this example
>>> # Usually, set global attention based on the task. For example,
>>> # classification: the <s> token
>>> # QA: question tokens
>>> # LM: potentially on the beginning of sentences and paragraphs
>>> outputs = model(input_ids, attention_mask=attention_mask, global_attention_mask=global_attention_mask)
>>> sequence_output = outputs.last_hidden_state
>>> pooled_output = outputs.pooler_output
```


class NllbTokenizer(PreTrainedTokenizer):
    vocab_files_names = VOCAB_FILES_NAMES
    model_input_names = ['input_ids', 'attention_mask']
    prefix_tokens: List[int] = []
    suffix_tokens: List[int] = []

    def __init__(self, vocab_file, bos_token='<s>', eos_token='</s>', sep_token='</s>', cls_token='<s>', unk_token='<unk>', pad_token='<pad>', mask_token='<mask>', tokenizer_file=None, src_lang=None, tgt_lang=None, sp_model_kwargs: Optional[Dict[str, Any]]=None, additional_special_tokens=None, legacy_behaviour=False, **kwargs):
        if additional_special_tokens is None:
            additional_special_tokens = FAIRSEQ_LANGUAGE_CODES
        bos_token = AddedToken(bos_token, normalized=False, special=True) if isinstance(bos_token, str) else bos_token
        pad_token = AddedToken(pad_token, normalized=False, special=True) if isinstance(pad_token, str) else pad_token
        eos_token = AddedToken(eos_token, normalized=False, special=True) if isinstance(eos_token, str) else eos_token
        unk_token = AddedToken(unk_token, normalized=False, special=True) if isinstance(unk_token, str) else unk_token
        mask_token = AddedToken(mask_token, normalized=True, lstrip=True, special=True) if isinstance(mask_token, str) else mask_token
        self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
        self.legacy_behaviour = legacy_behaviour
        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
        self.sp_model.Load(str(vocab_file))
        self.vocab_file = vocab_file
        self._added_tokens_decoder = {0: bos_token, 1: pad_token, 2: eos_token, 3: unk_token}
        self.fairseq_offset = 1
        self.sp_model_size = len(self.sp_model)
        super().__init__(bos_token=bos_token, eos_token=eos_token, unk_token=unk_token, sep_token=sep_token, cls_token=cls_token, pad_token=pad_token, mask_token=mask_token, tokenizer_file=tokenizer_file, src_lang=src_lang, tgt_lang=tgt_lang, additional_special_tokens=additional_special_tokens, sp_model_kwargs=self.sp_model_kwargs, legacy_behaviour=legacy_behaviour, **kwargs)
        self._src_lang = src_lang if src_lang is not None else 'eng_Latn'
        self.cur_lang_code_id = self.convert_tokens_to_ids(self._src_lang)
        self.tgt_lang = tgt_lang
        self.set_src_lang_special_tokens(self._src_lang)

    def __getstate__(self):
        state = self.__dict__.copy()
        state['sp_model'] = None
        state['sp_model_proto'] = self.sp_model.serialized_model_proto()
        return state

    def __setstate__(self, d):
        self.__dict__ = d
        if not hasattr(self, 'sp_model_kwargs'):
            self.sp_model_kwargs = {}
        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
        self.sp_model.LoadFromSerializedProto(self.sp_model_proto)

    @property
    def vocab_size(self):
        return len(self.sp_model) + self.fairseq_offset

    @property
    def src_lang(self) -> str:
        return self._src_lang

    @src_lang.setter
    def src_lang(self, new_src_lang: str) -> None:
        self._src_lang = new_src_lang
        self.set_src_lang_special_tokens(self._src_lang)

    def get_special_tokens_mask(self, token_ids_0: List[int], token_ids_1: Optional[List[int]]=None, already_has_special_tokens: bool=False) -> List[int]:
        """
        Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
        special tokens using the tokenizer `prepare_for_model` method.

        Args:
            token_ids_0 (`List[int]`):
                List of IDs.
            token_ids_1 (`List[int]`, *optional*):
                Optional second list of IDs for sequence pairs.
            already_has_special_tokens (`bool`, *optional*, defaults to `False`):
                Whether or not the token list is already formatted with special tokens for the model.

        Returns:
            `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
        """
        if already_has_special_tokens:
            return super().get_special_tokens_mask(token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True)
        prefix_ones = [1] * len(self.prefix_tokens)
        suffix_ones = [1] * len(self.suffix_tokens)
        if token_ids_1 is None:
            return prefix_ones + [0] * len(token_ids_0) + suffix_ones
        return prefix_ones + [0] * len(token_ids_0) + [0] * len(token_ids_1) + suffix_ones

    def build_inputs_with_special_tokens(self, token_ids_0: List[int], token_ids_1: Optional[List[int]]=None) -> List[int]:
        """
        Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
        adding special tokens. An NLLB sequence has the following format, where `X` represents the sequence:

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

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

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

        Returns:
            `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
        """
        if token_ids_1 is None:
            return self.prefix_tokens + token_ids_0 + self.suffix_tokens
        return self.prefix_tokens + token_ids_0 + token_ids_1 + self.suffix_tokens

    def create_token_type_ids_from_sequences(self, token_ids_0: List[int], token_ids_1: Optional[List[int]]=None) -> List[int]:
        """
        Create a mask from the two sequences passed to be used in a sequence-pair classification task. nllb does not
        make use of token type ids, therefore a list of zeros is returned.

        Args:
            token_ids_0 (`List[int]`):
                List of IDs.
            token_ids_1 (`List[int]`, *optional*):
                Optional second list of IDs for sequence pairs.

        Returns:
            `List[int]`: List of zeros.

        """
        sep = [self.sep_token_id]
        cls = [self.cls_token_id]
        if token_ids_1 is None:
            return len(cls + token_ids_0 + sep) * [0]
        return len(cls + token_ids_0 + sep + sep + token_ids_1 + sep) * [0]

    def _build_translation_inputs(self, raw_inputs, return_tensors: str, src_lang: Optional[str], tgt_lang: Optional[str], **extra_kwargs):
        """Used by translation pipeline, to prepare inputs for the generate function"""
        if src_lang is None or tgt_lang is None:
            raise ValueError('Translation requires a `src_lang` and a `tgt_lang` for this model')
        self.src_lang = src_lang
        inputs = self(raw_inputs, add_special_tokens=True, return_tensors=return_tensors, **extra_kwargs)
        tgt_lang_id = self.convert_tokens_to_ids(tgt_lang)
        inputs['forced_bos_token_id'] = tgt_lang_id
        return inputs

    def get_vocab(self):
        vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
        vocab.update(self.added_tokens_encoder)
        return vocab

    def _tokenize(self, text: str) -> List[str]:
        return self.sp_model.encode(text, out_type=str)

    def _convert_token_to_id(self, token):
        """Converts a token (str) in an id using the vocab."""
        spm_id = self.sp_model.PieceToId(token)
        return spm_id + self.fairseq_offset if spm_id else self.unk_token_id

    def _convert_id_to_token(self, index):
        """Converts an index (integer) in a token (str) using the vocab."""
        return self.sp_model.IdToPiece(index - self.fairseq_offset)

    def convert_tokens_to_string(self, tokens):
        """Converts a sequence of tokens (strings for sub-words) in a single string."""
        out_string = ''.join(tokens).replace(SPIECE_UNDERLINE, ' ').strip()
        return out_string

    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str]=None) -> Tuple[str]:
        if not os.path.isdir(save_directory):
            logger.error(f'Vocabulary path ({save_directory}) should be a directory')
            return
        out_vocab_file = os.path.join(save_directory, (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file'])
        if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
            copyfile(self.vocab_file, out_vocab_file)
        elif not os.path.isfile(self.vocab_file):
            with open(out_vocab_file, 'wb') as fi:
                content_spiece_model = self.sp_model.serialized_model_proto()
                fi.write(content_spiece_model)
        return (out_vocab_file,)

    def prepare_seq2seq_batch(self, src_texts: List[str], src_lang: str='eng_Latn', tgt_texts: Optional[List[str]]=None, tgt_lang: str='fra_Latn', **kwargs) -> BatchEncoding:
        self.src_lang = src_lang
        self.tgt_lang = tgt_lang
        return super().prepare_seq2seq_batch(src_texts, tgt_texts, **kwargs)

    def _switch_to_input_mode(self):
        return self.set_src_lang_special_tokens(self.src_lang)

    def _switch_to_target_mode(self):
        return self.set_tgt_lang_special_tokens(self.tgt_lang)

    def set_src_lang_special_tokens(self, src_lang) -> None:
        """Reset the special tokens to the source lang setting.
        - In legacy mode: No prefix and suffix=[eos, src_lang_code].
        - In default mode: Prefix=[src_lang_code], suffix = [eos]
        """
        self.cur_lang_code = self.convert_tokens_to_ids(src_lang)
        if self.legacy_behaviour:
            self.prefix_tokens = []
            self.suffix_tokens = [self.eos_token_id, self.cur_lang_code]
        else:
            self.prefix_tokens = [self.cur_lang_code]
            self.suffix_tokens = [self.eos_token_id]

    def set_tgt_lang_special_tokens(self, lang: str) -> None:
        """Reset the special tokens to the target lang setting.
        - In legacy mode: No prefix and suffix=[eos, tgt_lang_code].
        - In default mode: Prefix=[tgt_lang_code], suffix = [eos]
        """
        self.cur_lang_code = self.convert_tokens_to_ids(lang)
        if self.legacy_behaviour:
            self.prefix_tokens = []
            self.suffix_tokens = [self.eos_token_id, self.cur_lang_code]
        else:
            self.prefix_tokens = [self.cur_lang_code]
            self.suffix_tokens = [self.eos_token_id]

Construct an NLLB tokenizer.

Adapted from [`RobertaTokenizer`] and [`XLNetTokenizer`]. Based on
[SentencePiece](https://github.com/google/sentencepiece).

The tokenization method is `<tokens> <eos> <language code>` for source language documents, and `<language code>
<tokens> <eos>` for target language documents.

Examples:

```python
>>> from transformers import NllbTokenizer

>>> tokenizer = NllbTokenizer.from_pretrained(
...     "facebook/nllb-200-distilled-600M", src_lang="eng_Latn", tgt_lang="fra_Latn"
... )
>>> example_english_phrase = " UN Chief Says There Is No Military Solution in Syria"
>>> expected_translation_french = "Le chef de l'ONU affirme qu'il n'y a pas de solution militaire en Syrie."
>>> inputs = tokenizer(example_english_phrase, text_target=expected_translation_french, return_tensors="pt")
```

Args:
    vocab_file (`str`):
        Path to the vocabulary file.
    bos_token (`str`, *optional*, defaults to `"<s>"`):
        The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.

        <Tip>

        When building a sequence using special tokens, this is not the token that is used for the beginning of
        sequence. The token used is the `cls_token`.

        </Tip>

    eos_token (`str`, *optional*, defaults to `"</s>"`):
        The end of sequence token.

        <Tip>

        When building a sequence using special tokens, this is not the token that is used for the end of sequence.
        The token used is the `sep_token`.

        </Tip>

    sep_token (`str`, *optional*, defaults to `"</s>"`):
        The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
        sequence classification or for a text and a question for question answering. It is also used as the last
        token of a sequence built with special tokens.
    cls_token (`str`, *optional*, defaults to `"<s>"`):
        The classifier token which is used when doing sequence classification (classification of the whole sequence
        instead of per-token classification). It is the first token of the sequence when built with special tokens.
    unk_token (`str`, *optional*, defaults to `"<unk>"`):
        The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
        token instead.
    pad_token (`str`, *optional*, defaults to `"<pad>"`):
        The token used for padding, for example when batching sequences of different lengths.
    mask_token (`str`, *optional*, defaults to `"<mask>"`):
        The token used for masking values. This is the token used when training this model with masked language
        modeling. This is the token which the model will try to predict.
    tokenizer_file (`str`, *optional*):
        The path to a tokenizer file to use instead of the vocab file.
    src_lang (`str`, *optional*):
        The language to use as source language for translation.
    tgt_lang (`str`, *optional*):
        The language to use as target language for translation.
    sp_model_kwargs (`Dict[str, str]`):
        Additional keyword arguments to pass to the model initialization.