def do_test(create_module_fn, exported_names=None, show_debug_info=False):
    if exported_names is None:
        exported_names = []
    logging.set_stderrthreshold('error')
    tf.enable_v2_behavior()

    def app_main(argv):
        """Function passed to absl.app.run."""
        if len(argv) > 1:
            raise app.UsageError('Too many command-line arguments.')
        if FLAGS.save_model_path:
            save_model_path = FLAGS.save_model_path
        else:
            save_model_path = tempfile.mkdtemp(suffix='.saved_model')
        save_options = tf.saved_model.SaveOptions(save_debug_info=show_debug_info)
        tf.saved_model.save(create_module_fn(), save_model_path, options=save_options)
        logging.info('Saved model to: %s', save_model_path)
        mlir = pywrap_mlir.experimental_convert_saved_model_to_mlir(save_model_path, ','.join(exported_names), show_debug_info)
        mlir = pywrap_mlir.experimental_run_pass_pipeline(mlir, 'canonicalize', show_debug_info)
        print(mlir)
        filename = '%s/result.mlirbc' % save_model_path
        pywrap_mlir.experimental_write_bytecode(filename, mlir)
        if not file_io.file_exists(filename):
            raise app.UsageError('Failed to create bytecode output.')
    app.run(app_main)

Runs test.

1. Performs absl and tf "main"-like initialization that must run before almost
   anything else.
2. Converts `tf.Module` to SavedModel
3. Converts SavedModel to MLIR
4. Prints the textual MLIR to stdout (it is expected that the caller will have
   FileCheck checks in its file to check this output).

This is only for use by the MLIR SavedModel importer tests.

Args:
  create_module_fn: A callable taking no arguments, which returns the
    `tf.Module` to be converted and printed.
  exported_names: A set of exported names for the MLIR converter (default is
    "export all").
  show_debug_info: If true, shows debug locations in the resulting MLIR.


def SyncSleep(delay, name=None):
    return examples_sync_sleep(delay=delay, name=name)

Pause for `delay` seconds (which need not be an integer).

This is a synchronous (blocking) version of a sleep op. It's purpose is
to be contrasted with Examples>AsyncSleep.

Args:
  delay: tf.Tensor which is a scalar of type float.
  name: An optional name for the op.

Returns:
  The `delay` value.


def get_audio_features(self, input_features: Optional[torch.Tensor]=None, is_longer: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> torch.FloatTensor:
    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
    audio_outputs = self.audio_model(input_features=input_features, is_longer=is_longer, return_dict=return_dict)
    pooled_output = audio_outputs[1] if not return_dict else audio_outputs.pooler_output
    audio_features = self.audio_projection(pooled_output)
    audio_features = F.normalize(audio_features, dim=-1)
    return audio_features

input_features (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
    Input audio features. This should be returned by the [`ClapFeatureExtractor`] class that you can also
    retrieve from [`AutoFeatureExtractor`]. See [`ClapFeatureExtractor.__call__`] for details.
is_longer (`torch.FloatTensor`, of shape `(batch_size, 1)`, *optional*):
    Whether the audio clip is longer than `max_length`. If `True`, a feature fusion will be enabled to enhance
    the features.

Returns:
    audio_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The audio embeddings obtained by
    applying the projection layer to the pooled output of [`ClapAudioModel`].

Examples:

```python
>>> from transformers import AutoFeatureExtractor, ClapModel
>>> import torch

>>> model = ClapModel.from_pretrained("laion/clap-htsat-unfused")
>>> feature_extractor = AutoFeatureExtractor.from_pretrained("laion/clap-htsat-unfused")
>>> random_audio = torch.rand((16_000))
>>> inputs = feature_extractor(random_audio, return_tensors="pt")
>>> audio_features = model.get_audio_features(**inputs)
```


def on_test_batch_end(self, batch, logs=None):

Called at the end of a batch in `evaluate` methods.

Also called at the end of a validation batch in the `fit`
methods, if validation data is provided.

Subclasses should override for any actions to run.

Note that if the `steps_per_execution` argument to `compile` in
`tf.keras.Model` is set to `N`, this method will only be called every `N`
batches.

Args:
    batch: Integer, index of batch within the current epoch.
    logs: Dict. Aggregated metric results up until this batch.


def distribute(processing_mode, service, job_name=None, consumer_index=None, num_consumers=None, max_outstanding_requests=None, data_transfer_protocol=None, compression='AUTO', cross_trainer_cache=None, target_workers='AUTO') -> Callable[dataset_ops.Dataset, dataset_ops.Dataset]:
    _validate_job_name(job_name)
    return _distribute(processing_mode=processing_mode, service=service, job_name=job_name, consumer_index=consumer_index, num_consumers=num_consumers, max_outstanding_requests=max_outstanding_requests, data_transfer_protocol=data_transfer_protocol, compression=compression, cross_trainer_cache=cross_trainer_cache, target_workers=target_workers)

A transformation that moves dataset processing to the tf.data service.

When you iterate over a dataset containing the `distribute` transformation,
the tf.data service creates a "job" which produces data for the dataset
iteration.

The tf.data service uses a cluster of workers to prepare data for training
your model.
The `processing_mode` argument to `tf.data.experimental.service.distribute`
describes how to leverage multiple workers to process the input dataset.
Currently, there are two processing modes to choose from: "distributed_epoch"
and "parallel_epochs".

"distributed_epoch" means that the dataset will be split across all tf.data
service workers.
The dispatcher produces "splits" for the dataset and sends them to workers for
further processing. For example, if a dataset begins with a list of filenames,
the dispatcher will iterate through the filenames and send the filenames to
tf.data workers, which will perform the rest of the dataset transformations on
those files. "distributed_epoch" is useful when your model needs to see each
element of the dataset exactly once, or if it needs to see the data in a
generally-sequential order. "distributed_epoch" only works for datasets with
splittable sources, such as `Dataset.from_tensor_slices`,
`Dataset.list_files`, or `Dataset.range`.

"parallel_epochs" means that the entire input dataset will be processed
independently by each of the tf.data service workers.
For this reason, it is important to shuffle data (e.g. filenames)
non-deterministically, so that each worker will process the elements of the
dataset in a different order. "parallel_epochs" can be used to distribute
datasets that aren't splittable.

With two workers, "parallel_epochs" will produce every element of the dataset
twice:

>>> dispatcher = tf.data.experimental.service.DispatchServer()
>>> dispatcher_address = dispatcher.target.split("://")[1]
>>> # Start two workers
>>> workers = [
...     tf.data.experimental.service.WorkerServer(
...         tf.data.experimental.service.WorkerConfig(
...             dispatcher_address=dispatcher_address)) for _ in range(2)
... ]
>>> dataset = tf.data.Dataset.range(10)
>>> dataset = dataset.apply(tf.data.experimental.service.distribute(
...     processing_mode="parallel_epochs", service=dispatcher.target))
>>> sorted([a.item() for a in dataset.as_numpy_iterator()])
[0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9]

"distributed_epoch", on the other hand, will still produce each element once:

>>> dispatcher = tf.data.experimental.service.DispatchServer()
>>> dispatcher_address = dispatcher.target.split("://")[1]
>>> workers = [
...     tf.data.experimental.service.WorkerServer(
...         tf.data.experimental.service.WorkerConfig(
...             dispatcher_address=dispatcher_address)) for _ in range(2)
... ]
>>> dataset = tf.data.Dataset.range(10)
>>> dataset = dataset.apply(tf.data.experimental.service.distribute(
...     processing_mode="distributed_epoch", service=dispatcher.target))
>>> sorted([a.item() for a in dataset.as_numpy_iterator()])
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

When using `apply(tf.data.experimental.service.distribute(...))`, the dataset
before the `apply` transformation executes within the tf.data service, while
the operations after `apply` happen within the local process.

>>> dispatcher = tf.data.experimental.service.DispatchServer()
>>> dispatcher_address = dispatcher.target.split("://")[1]
>>> workers = [
...     tf.data.experimental.service.WorkerServer(
...         tf.data.experimental.service.WorkerConfig(
...             dispatcher_address=dispatcher_address)) for _ in range(2)
... ]
>>> dataset = tf.data.Dataset.range(5)
>>> dataset = dataset.map(lambda x: x*x)
>>> dataset = dataset.apply(
...    tf.data.experimental.service.distribute("parallel_epochs",
...                                            dispatcher.target))
>>> dataset = dataset.map(lambda x: x+1)
>>> sorted([a.item() for a in dataset.as_numpy_iterator()])
[1, 1, 2, 2, 5, 5, 10, 10, 17, 17]

In the above example, the dataset operations (before applying the `distribute`
function on the elements) will be executed on the tf.data workers,
and the elements are provided over RPC. The remaining transformations
(after the call to `distribute`) will be executed locally. The dispatcher
and the workers will bind to usused free ports (which are chosen at random),
in order to communicate with each other. However, to bind them to specific
ports, the `port` parameter can be passed.

The `job_name` argument allows jobs to be shared across multiple
datasets. Instead of each dataset creating its own job, all
datasets with the same `job_name` will consume from the same job. A new job
will be created for each iteration of the dataset (with each repetition of
`Dataset.repeat` counting as a new iteration). Suppose the `DispatchServer`
is serving on `localhost:5000` and two training workers (in either a single
client or multi-client setup) iterate over the below dataset, and there is a
single tf.data worker:

```
range5_dataset = tf.data.Dataset.range(5)
dataset = range5_dataset.apply(tf.data.experimental.service.distribute(
    "parallel_epochs", "localhost:5000", job_name="my_job_name"))
for iteration in range(3):
  print(list(dataset))
```

The elements of each job will be split between the two processes, with
elements being consumed by the processes on a first-come first-served basis.
One possible result is that process 1 prints

```
[0, 2, 4]
[0, 1, 3]
[1]
```

and process 2 prints

```
[1, 3]
[2, 4]
[0, 2, 3, 4]
```

Job names must not be re-used across different training jobs within the
lifetime of the tf.data service. In general, the tf.data service is expected
to live for the duration of a single training job.
To use the tf.data service with multiple training jobs, make sure to use
different job names to avoid conflicts. For example, suppose a training job
calls `distribute` with `job_name="job"` and reads until end of input. If
another independent job connects to the same tf.data service and tries to read
from `job_name="job"`, it will immediately receive end of input, without
getting any data.

**Coordinated data read**

By default, when multiple consumers read from the same job, they receive data
on a first-come first-served basis. In some use cases, it is advantageous to
coordinate the consumers. At each step, consumers read data from the same
worker.

For example, the tf.data service can be used to coordinate example sizes
across a cluster during synchronous training, so that during each step all
replicas train on similar-sized elements. To achieve this, define a dataset
which generates rounds of `num_consumers` consecutive similar-sized batches,
then enable coordinated reads by setting `consumer_index` and `num_consumers`.

NOTE: To keep consumers in sync, round robin data consumption requires that
the dataset have infinite cardinality. You can get this by adding `.repeat()`
at the end of the dataset definition.

**Keras and Distribution Strategies**

The dataset produced by the `distribute` transformation can be passed to
Keras' `Model.fit` or Distribution Strategy's
`tf.distribute.Strategy.experimental_distribute_dataset` like any other
`tf.data.Dataset`. We recommend setting a `job_name` on the call to
`distribute` so that if there are multiple workers, they read data from the
same job. Note that the autosharding normally performed by
`experimental_distribute_dataset` will be disabled when setting a `job_name`,
since sharing the job already results in splitting data across the workers.
When using a shared job, data will be dynamically balanced across workers, so
that they reach end of input about the same time. This results in better
worker utilization than with autosharding, where each worker processes an
independent set of files, and some workers may run out of data earlier than
others.

Args:
  processing_mode: A `tf.data.experimental.service.ShardingPolicy` specifying
    how to shard the dataset among tf.data workers. See
    `tf.data.experimental.service.ShardingPolicy` for details. For backwards
    compatibility, `processing_mode` may also be set to the strings
    `"parallel_epochs"` or `"distributed_epoch"`, which are respectively
    equivalent to `ShardingPolicy.OFF` and `ShardingPolicy.DYNAMIC`.
  service: A string or a tuple indicating how to connect to the tf.data
    service. If it's a string, it should be in the format
    `[<protocol>://]<address>`, where `<address>` identifies the dispatcher
    address and `<protocol>` can optionally be used to override the default
    protocol to use. If it's a tuple, it should be (protocol, address).
  job_name: (Optional.) The name of the job. If provided, it must be a
    non-empty string. This argument makes it possible for multiple datasets to
    share the same job. The default behavior is that the dataset creates
    anonymous, exclusively owned jobs.
  consumer_index: (Optional.) The index of the consumer in the range from `0`
    to `num_consumers`. Must be specified alongside `num_consumers`. When
    specified, consumers will read from the job in a strict round-robin order,
    instead of the default first-come-first-served order.
  num_consumers: (Optional.) The number of consumers which will consume from
    the job. Must be specified alongside `consumer_index`. When specified,
    consumers will read from the job in a strict round-robin order, instead of
    the default first-come-first-served order. When `num_consumers` is
    specified, the dataset must have infinite cardinality to prevent a
    producer from running out of data early and causing consumers to go out of
    sync.
  max_outstanding_requests: (Optional.) A limit on how many elements may be
    requested at the same time. You can use this option to control the amount
    of memory used, since `distribute` won't use more than `element_size` *
    `max_outstanding_requests` of memory.
  data_transfer_protocol: (Optional.) The protocol to use for transferring
    data with the tf.data service. If not provided, a protocol is determined
    at runtime.
  compression: How to compress the dataset's elements before transferring them
    over the network. "AUTO" leaves the decision of how to compress up to the
    tf.data service runtime. `None` indicates not to compress. "SNAPPY" forces
    the use of snappy compression.
  cross_trainer_cache: (Optional.) If a `CrossTrainerCache` object is
    provided, dataset iteration will be shared across concurrently running
    trainers. See
    https://www.tensorflow.org/api_docs/python/tf/data/experimental/service#sharing_tfdata_service_with_concurrent_trainers
      for details.
  target_workers: (Optional.) Which workers to read from. If `"AUTO"`, tf.data
    runtime decides which workers to read from. If `"ANY"`, reads from any
    tf.data service workers. If `"LOCAL"`, only reads from local in-process
    tf.data service workers. `"AUTO"` works well for most cases, while users
    can specify other targets. For example, `"LOCAL"` helps avoid RPCs and
    data copy if every TF worker colocates with a tf.data service worker.
    Consumers of a shared job must use the same `target_workers`. Defaults to
    `"AUTO"`.

Returns:
  Dataset: A `Dataset` of the elements produced by the data service.


def run_op_benchmark(self, op, iters=1, warmup=True, session_config=None):
    if context.executing_eagerly():
        return self._run_eager_benchmark(iterable=op, iters=iters, warmup=warmup)
    return self._run_graph_benchmark(iterable=op, iters=iters, warmup=warmup, session_config=session_config)

Benchmarks the op.

Runs the op `iters` times. In each iteration, the benchmark measures
the time it takes to go execute the op.

Args:
  op: The tf op to benchmark.
  iters: Number of times to repeat the timing.
  warmup: If true, warms up the session caches by running an untimed run.
  session_config: A ConfigProto protocol buffer with configuration options
    for the session. Applicable only for benchmarking in graph mode.

Returns:
  A float, representing the per-execution wall time of the op in seconds.
  This is the median time (with respect to `iters`) it takes for the op
  to be executed `iters` num of times.


def __getitem__(self, index: Any) -> Rigid:
    if type(index) is not tuple:
        index = (index,)
    return Rigid(self._rots[index], self._trans[index + (slice(None),)])

Indexes the affine transformation with PyTorch-style indices. The index is applied to the shared dimensions of
both the rotation and the translation.

E.g.::

    r = Rotation(rot_mats=torch.rand(10, 10, 3, 3), quats=None) t = Rigid(r, torch.rand(10, 10, 3)) indexed =
    t[3, 4:6] assert(indexed.shape == (2,)) assert(indexed.get_rots().shape == (2,))
    assert(indexed.get_trans().shape == (2, 3))

Args:
    index: A standard torch tensor index. E.g. 8, (10, None, 3),
    or (3, slice(0, 1, None))
Returns:
    The indexed tensor


def get_resource(self, uri: str) -> Optional[message.Message]:
    for collection in (self.structure_definitions, self.search_parameters, self.code_systems, self.value_sets):
        resource = collection.get(uri)
        if resource is not None:
            return resource
    return None

Retrieves a protocol buffer representation of the given resource.

Args:
  uri: The URI of the resource to retrieve.

Returns:
  Protocol buffer for the resource or `None` if the `uri` can not be found.


def __init__(self, num_workers, dispatcher_port=0, work_dir=TMP_WORK_DIR, fault_tolerant_mode=True, job_gc_check_interval_ms=TEST_JOB_GC_CHECK_INTERNAL_MS, job_gc_timeout_ms=None, worker_timeout_ms=TEST_WORKER_TIMEOUT_MS, worker_shutdown_quiet_period_ms=0, snapshot_max_chunk_size_bytes=TEST_SNAPSHOT_MAX_CHUNK_SIZE_BYTES, worker_max_concurrent_snapshots=0, start=True, protocol=PROTOCOL, data_transfer_protocol=None):
    if work_dir == TMP_WORK_DIR:
        work_dir = tempfile.mkdtemp(dir=googletest.GetTempDir())
    self._worker_shutdown_quiet_period_ms = worker_shutdown_quiet_period_ms
    self._snapshot_max_chunk_size_bytes = snapshot_max_chunk_size_bytes
    self._protocol = protocol
    self._data_transfer_protocol = data_transfer_protocol
    self._job_gc_check_interval_ms = job_gc_check_interval_ms
    self._job_gc_timeout_ms = job_gc_timeout_ms
    self._worker_timeout_ms = worker_timeout_ms
    self._worker_max_concurrent_snapshots = worker_max_concurrent_snapshots
    self.dispatcher = server_lib.DispatchServer(server_lib.DispatcherConfig(port=dispatcher_port, work_dir=work_dir, protocol=protocol, fault_tolerant_mode=fault_tolerant_mode, job_gc_check_interval_ms=job_gc_check_interval_ms, job_gc_timeout_ms=job_gc_timeout_ms, worker_timeout_ms=worker_timeout_ms, worker_max_concurrent_snapshots=worker_max_concurrent_snapshots), start=start)
    self.workers = []
    for _ in range(num_workers):
        self.add_worker(start=start)

Creates a tf.data service test cluster.

Args:
  num_workers: The number of workers to initially add to the cluster.
  dispatcher_port: The port to use for the dispatcher.
  work_dir: The work directory to use for the dispatcher. If set to
    `TMP_WORK_DIR`, the cluster will create a new temporary directory to use
    as the work directory. If set to `NO_WORK_DIR`, no work directory will
    be used.
  fault_tolerant_mode: Whether the dispatcher should write its state to a
    journal so that it can recover from restarts.
  job_gc_check_interval_ms: How often the dispatcher should scan through to
    delete old and unused jobs, in milliseconds.
  job_gc_timeout_ms: How long a job needs to be unused before it becomes a
    candidate for garbage collection, in milliseconds.
  worker_timeout_ms: How long to wait for a worker to heartbeat before
    considering it missing, in milliseconds.
  worker_shutdown_quiet_period_ms: When shutting down a worker, how long to
    wait for the gRPC server to process the final requests.
  snapshot_max_chunk_size_bytes: The maximum size of a distributed snapshot
    chunk file.
  worker_max_concurrent_snapshots: The maximum number of snapshots a worker
    can concurrently process.
  start: Whether to immediately start the servers in the cluster. If
    `False`, the servers can be started later by calling
    `start_dispatcher()` and `start_workers()`.
  protocol: The protocol to use for communicating with the tf.data service,
    e.g. "grpc".
  data_transfer_protocol: (Optional.) The protocol to use for transferring
    data with the tf.data service.


def open(self, url, mime_type='application/octet-stream', compression_type=CompressionTypes.AUTO) -> BinaryIO:
    _, path = self._parse_url(url)
    return self._open(path, mime_type, compression_type)


def resize_tensor_input(self, input_index, tensor_size, strict=False):
    self._ensure_safe()
    tensor_size = np.array(tensor_size, dtype=np.int32)
    self._interpreter.ResizeInputTensor(input_index, tensor_size, strict)

Resizes an input tensor.

Args:
  input_index: Tensor index of input to set. This value can be gotten from
    the 'index' field in get_input_details.
  tensor_size: The tensor_shape to resize the input to.
  strict: Only unknown dimensions can be resized when `strict` is True.
    Unknown dimensions are indicated as `-1` in the `shape_signature`
    attribute of a given tensor. (default False)

Raises:
  ValueError: If the interpreter could not resize the input tensor.

Usage:
```
interpreter = Interpreter(model_content=tflite_model)
interpreter.resize_tensor_input(0, [num_test_images, 224, 224, 3])
interpreter.allocate_tensors()
interpreter.set_tensor(0, test_images)
interpreter.invoke()
```


def ldu(load_v, name):
    try:
        return load_v()
    except (KeyError, AttributeError, NameError):
        return Undefined(name)

Load variable operator that returns Undefined when failing to evaluate.

Note: the name ("load or return undefined") is abbreviated to minimize
the amount of clutter in generated code.

This variant of `ld` is useful when loading symbols that may be undefined at
runtime, such as composite symbols, and whether they are defined or not cannot
be determined statically. For example `d['a']` is undefined when `d` is an
empty dict.

Args:
  load_v: Lambda that executes the actual read.
  name: Human-readable name of the symbol being read.
Returns:
  Either the value of the symbol, or Undefined, if the symbol is not fully
  defined.


def _parse_test_option_args(self, argv):
    parser = argparse.ArgumentParser()
    parser.add_argument('--test-pipeline-options', type=str, action='store', help='only run tests providing service options')
    parser.add_argument('--not-use-test-runner-api', action='store_true', default=False, help='whether not to use test-runner-api')
    known, unused_argv = parser.parse_known_args(argv)
    test_pipeline_options = known.test_pipeline_options or TestPipeline.pytest_test_pipeline_options
    if self.is_integration_test and (not test_pipeline_options):
        raise SkipTest('IT is skipped because --test-pipeline-options is not specified')
    self.not_use_test_runner_api = known.not_use_test_runner_api
    return shlex.split(test_pipeline_options) if test_pipeline_options else []

Parse value of command line argument: --test-pipeline-options to get
pipeline options.

Args:
  argv: An iterable of command line arguments to be used. If not specified
    then sys.argv will be used as input for parsing arguments.

Returns:
  An argument list of options that can be parsed by argparser or directly
  build a pipeline option.


def preprocess(self, images: ImageInput, do_resize: Optional[bool]=None, size: Optional[Dict[str, int]]=None, resample: PILImageResampling=None, do_center_crop: Optional[bool]=None, crop_size: Optional[int]=None, do_rescale: Optional[bool]=None, rescale_factor: Optional[float]=None, do_normalize: Optional[bool]=None, image_mean: Optional[Union[float, List[float]]]=None, image_std: Optional[Union[float, List[float]]]=None, do_convert_rgb: Optional[bool]=None, return_tensors: Optional[Union[str, TensorType]]=None, data_format: Optional[ChannelDimension]=ChannelDimension.FIRST, input_data_format: Optional[Union[str, ChannelDimension]]=None) -> PIL.Image.Image:
    do_resize = do_resize if do_resize is not None else self.do_resize
    size = size if size is not None else self.size
    size = get_size_dict(size, default_to_square=False)
    resample = resample if resample is not None else self.resample
    do_center_crop = do_center_crop if do_center_crop is not None else self.do_center_crop
    crop_size = crop_size if crop_size is not None else self.crop_size
    crop_size = get_size_dict(crop_size)
    do_rescale = do_rescale if do_rescale is not None else self.do_rescale
    rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
    do_normalize = do_normalize if do_normalize is not None else self.do_normalize
    image_mean = image_mean if image_mean is not None else self.image_mean
    image_std = image_std if image_std is not None else self.image_std
    do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
    images = make_list_of_images(images)
    if not valid_images(images):
        raise ValueError('Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, torch.Tensor, tf.Tensor or jax.ndarray.')
    validate_preprocess_arguments(do_rescale=do_rescale, rescale_factor=rescale_factor, do_normalize=do_normalize, image_mean=image_mean, image_std=image_std, do_center_crop=do_center_crop, crop_size=crop_size, do_resize=do_resize, size=size, resample=resample)
    if do_convert_rgb:
        images = [convert_to_rgb(image) for image in images]
    images = [to_numpy_array(image) for image in images]
    if do_rescale and is_scaled_image(images[0]):
        logger.warning_once('It looks like you are trying to rescale already rescaled images. If the input images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again.')
    if input_data_format is None:
        input_data_format = infer_channel_dimension_format(images[0])
    all_images = []
    for image in images:
        if do_resize:
            image = self.resize(image=image, size=size, resample=resample, input_data_format=input_data_format)
        if do_center_crop:
            image = self.center_crop(image=image, size=crop_size, input_data_format=input_data_format)
        if do_rescale:
            image = self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format)
        if do_normalize:
            image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
        all_images.append(image)
    images = [to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) for image in all_images]
    data = {'pixel_values': images}
    return BatchFeature(data=data, tensor_type=return_tensors)

Preprocess an image or batch of images.

Args:
    images (`ImageInput`):
        Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
        passing in images with pixel values between 0 and 1, set `do_rescale=False`.
    do_resize (`bool`, *optional*, defaults to `self.do_resize`):
        Whether to resize the image.
    size (`Dict[str, int]`, *optional*, defaults to `self.size`):
        Size of the image after resizing. Shortest edge of the image is resized to size["shortest_edge"], with
        the longest edge resized to keep the input aspect ratio.
    resample (`int`, *optional*, defaults to `self.resample`):
        Resampling filter to use if resizing the image. This can be one of the enum `PILImageResampling`. Only
        has an effect if `do_resize` is set to `True`.
    do_center_crop (`bool`, *optional*, defaults to `self.do_center_crop`):
        Whether to center crop the image.
    crop_size (`Dict[str, int]`, *optional*, defaults to `self.crop_size`):
        Size of the center crop. Only has an effect if `do_center_crop` is set to `True`.
    do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
        Whether to rescale the image.
    rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
        Rescale factor to rescale the image by if `do_rescale` is set to `True`.
    do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
        Whether to normalize the image.
    image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
        Image mean to use for normalization. Only has an effect if `do_normalize` is set to `True`.
    image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
        Image standard deviation to use for normalization. Only has an effect if `do_normalize` is set to
        `True`.
    do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):
        Whether to convert the image to RGB.
    return_tensors (`str` or `TensorType`, *optional*):
        The type of tensors to return. Can be one of:
        - Unset: Return a list of `np.ndarray`.
        - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
        - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
        - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
        - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
    data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
        The channel dimension format for the output 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.
        - Unset: Use the channel dimension format of the input image.
    input_data_format (`ChannelDimension` or `str`, *optional*):
        The channel dimension format for the input image. If unset, the channel dimension format 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.
        - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.


def distance(a, b):
    a = a.lower()
    b = b.lower()
    if len(a) == 0:
        return len(b)
    if len(b) == 0:
        return len(a)
    dist = [[0 for _ in range(len(b) + 1)] for _ in range(len(a) + 1)]
    for i in range(len(a) + 1):
        dist[i][0] = i
    for j in range(len(b) + 1):
        dist[0][j] = j
    for i in range(1, len(a) + 1):
        for j in range(1, len(b) + 1):
            cost = 0 if a[i - 1] == b[j - 1] else 1
            dist[i][j] = min(dist[i - 1][j] + 1, dist[i][j - 1] + 1, dist[i - 1][j - 1] + cost)
    return dist[-1][-1]

Returns the case-insensitive Levenshtein edit distance between |a| and |b|.

The Levenshtein distance is a metric for measuring the difference between
two strings. If |a| == |b|, the distance is 0. It is roughly the number
of insertions, deletions, and substitutions needed to convert |a| -> |b|.

This distance is at most the length of the longer string.
This distance is 0 iff the strings are equal.

Examples:
  levenshtein_distance("cow", "bow") == 1
  levenshtein_distance("cow", "bowl") == 2
  levenshtein_distance("cow", "blrp") == 4

See https://en.wikipedia.org/wiki/Levenshtein_distance for more background.

Args:
  a: A string
  b: A string

Returns:
  The Levenshtein distance between the inputs.


def _pad(self, encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding], max_length: Optional[int]=None, max_entity_length: Optional[int]=None, padding_strategy: PaddingStrategy=PaddingStrategy.DO_NOT_PAD, pad_to_multiple_of: Optional[int]=None, padding_side: Optional[str]=None, return_attention_mask: Optional[bool]=None) -> dict:
    entities_provided = bool('entity_ids' in encoded_inputs)
    if return_attention_mask is None:
        return_attention_mask = 'attention_mask' in self.model_input_names
    if padding_strategy == PaddingStrategy.LONGEST:
        max_length = len(encoded_inputs['input_ids'])
        if entities_provided:
            max_entity_length = len(encoded_inputs['entity_ids'])
    if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
        max_length = (max_length // pad_to_multiple_of + 1) * pad_to_multiple_of
    if entities_provided and max_entity_length is not None and (pad_to_multiple_of is not None) and (max_entity_length % pad_to_multiple_of != 0):
        max_entity_length = (max_entity_length // pad_to_multiple_of + 1) * pad_to_multiple_of
    needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and (len(encoded_inputs['input_ids']) != max_length or (entities_provided and len(encoded_inputs['entity_ids']) != max_entity_length))
    if return_attention_mask and 'attention_mask' not in encoded_inputs:
        encoded_inputs['attention_mask'] = [1] * len(encoded_inputs['input_ids'])
    if entities_provided and return_attention_mask and ('entity_attention_mask' not in encoded_inputs):
        encoded_inputs['entity_attention_mask'] = [1] * len(encoded_inputs['entity_ids'])
    if needs_to_be_padded:
        difference = max_length - len(encoded_inputs['input_ids'])
        padding_side = padding_side if padding_side is not None else self.padding_side
        if entities_provided:
            entity_difference = max_entity_length - len(encoded_inputs['entity_ids'])
        if padding_side == 'right':
            if return_attention_mask:
                encoded_inputs['attention_mask'] = encoded_inputs['attention_mask'] + [0] * difference
                if entities_provided:
                    encoded_inputs['entity_attention_mask'] = encoded_inputs['entity_attention_mask'] + [0] * entity_difference
            if 'token_type_ids' in encoded_inputs:
                encoded_inputs['token_type_ids'] = encoded_inputs['token_type_ids'] + [0] * difference
                if entities_provided:
                    encoded_inputs['entity_token_type_ids'] = encoded_inputs['entity_token_type_ids'] + [0] * entity_difference
            if 'special_tokens_mask' in encoded_inputs:
                encoded_inputs['special_tokens_mask'] = encoded_inputs['special_tokens_mask'] + [1] * difference
            encoded_inputs['input_ids'] = encoded_inputs['input_ids'] + [self.pad_token_id] * difference
            if entities_provided:
                encoded_inputs['entity_ids'] = encoded_inputs['entity_ids'] + [self.entity_pad_token_id] * entity_difference
                encoded_inputs['entity_position_ids'] = encoded_inputs['entity_position_ids'] + [[-1] * self.max_mention_length] * entity_difference
                if self.task == 'entity_span_classification':
                    encoded_inputs['entity_start_positions'] = encoded_inputs['entity_start_positions'] + [0] * entity_difference
                    encoded_inputs['entity_end_positions'] = encoded_inputs['entity_end_positions'] + [0] * entity_difference
        elif padding_side == 'left':
            if return_attention_mask:
                encoded_inputs['attention_mask'] = [0] * difference + encoded_inputs['attention_mask']
                if entities_provided:
                    encoded_inputs['entity_attention_mask'] = [0] * entity_difference + encoded_inputs['entity_attention_mask']
            if 'token_type_ids' in encoded_inputs:
                encoded_inputs['token_type_ids'] = [0] * difference + encoded_inputs['token_type_ids']
                if entities_provided:
                    encoded_inputs['entity_token_type_ids'] = [0] * entity_difference + encoded_inputs['entity_token_type_ids']
            if 'special_tokens_mask' in encoded_inputs:
                encoded_inputs['special_tokens_mask'] = [1] * difference + encoded_inputs['special_tokens_mask']
            encoded_inputs['input_ids'] = [self.pad_token_id] * difference + encoded_inputs['input_ids']
            if entities_provided:
                encoded_inputs['entity_ids'] = [self.entity_pad_token_id] * entity_difference + encoded_inputs['entity_ids']
                encoded_inputs['entity_position_ids'] = [[-1] * self.max_mention_length] * entity_difference + encoded_inputs['entity_position_ids']
                if self.task == 'entity_span_classification':
                    encoded_inputs['entity_start_positions'] = [0] * entity_difference + encoded_inputs['entity_start_positions']
                    encoded_inputs['entity_end_positions'] = [0] * entity_difference + encoded_inputs['entity_end_positions']
        else:
            raise ValueError('Invalid padding strategy:' + str(padding_side))
    return encoded_inputs

Pad encoded inputs (on left/right and up to predefined length or max length in the batch)

Args:
    encoded_inputs:
        Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`).
    max_length: maximum length of the returned list and optionally padding length (see below).
        Will truncate by taking into account the special tokens.
    max_entity_length: The maximum length of the entity sequence.
    padding_strategy: PaddingStrategy to use for padding.

        - PaddingStrategy.LONGEST Pad to the longest sequence in the batch
        - PaddingStrategy.MAX_LENGTH: Pad to the max length (default)
        - PaddingStrategy.DO_NOT_PAD: Do not pad
        The tokenizer padding sides are defined in self.padding_side:

            - 'left': pads on the left of the sequences
            - 'right': pads on the right of the sequences
    pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value.
        This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability
        `>= 7.5` (Volta).
    padding_side:
        The side on which the model should have padding applied. Should be selected between ['right', 'left'].
        Default value is picked from the class attribute of the same name.
    return_attention_mask:
        (optional) Set to False to avoid returning attention mask (default: set to model specifics)


def get_module(dir_path: str, relative_to_dir: str) -> str:
    dir_path = dir_path[len(relative_to_dir):]
    dir_path = dir_path.replace(os.sep, '/')
    return dir_path.replace('/', '.').strip('.')

Get module that corresponds to path relative to relative_to_dir.

Args:
  dir_path: Path to directory.
  relative_to_dir: Get module relative to this directory.

Returns:
  Name of module that corresponds to the given directory.


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

Applies Rotary Position Embedding to the query and key tensors.

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


def move_test_classes_into_scope(wrapped_test_module):
    for name, obj in wrapped_test_module.__dict__.items():
        if _is_test_class(obj):
            module_variables['tpu_test_imported_%s' % name] = obj

Add all test classes defined in wrapped module to our module.

The test runner works by inspecting the main module for TestCase classes, so
by adding a module-level reference to the TestCase we cause it to execute the
wrapped TestCase.

Args:
  wrapped_test_module: The user-provided test code to run.


def reinterpret(value, new_type):
    if not isinstance(value, ExtensionType):
        raise ValueError(f'reinterpret expects `value` to be a tf.ExtensionType instance; got {value!r}')
    if not (isinstance(new_type, type) and issubclass(new_type, ExtensionType)):
        raise ValueError(f'reinterpret expects `new_type` to be a subclass of tf.ExtensionType; got {new_type!r}')
    fields = [item for item in value.__dict__.items() if not extension_type_field.ExtensionTypeField.is_reserved_name(item[0])]
    new_value = _create_object_from_type_and_dict(new_type, fields)
    new_value._tf_extension_type_convert_fields()
    new_value.__validate__()
    return new_value

Converts a given `ExtensionType` to a new type with compatible fields.

In particular, this can be used to convert a concrete subclass of
`ExtensionType` to an `AnonymousExtensionType`, or vice versa.  When
converting to a non-anonymous ExtensionType, field values are type-checked to
ensure they are consistent with `new_type`'s type annotations, and validated
with `new_type.__validate__`.

Args:
  value: An instance of a subclass of `tf.ExtensionType`
  new_type: A subclass of `tf.ExtensionType`

Returns:
  An instance of `new_type`, whose fields are copied from `value`.


class TvltDecoderOutput(ModelOutput):
    logits: Optional[torch.FloatTensor] = None
    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None

Class for TvltDecoder's outputs, with potential hidden states and attentions.

Args:
    logits (`torch.FloatTensor` of shape `(batch_size, patch_size ** 2 * num_channels)`):
        Pixel reconstruction logits.
    hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `torch.FloatTensor` (one for the output of the embeddings and 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(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
        Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
        sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in
        the self-attention heads.


class TFXLNetForMultipleChoiceOutput(ModelOutput):
    loss: tf.Tensor | None = None
    logits: Optional[tf.Tensor] = None
    mems: List[tf.Tensor] | None = None
    hidden_states: Tuple[tf.Tensor, ...] | None = None
    attentions: Tuple[tf.Tensor, ...] | None = None

Output type of [`TFXLNetForMultipleChoice`].

Args:
    loss (`tf.Tensor` of shape *(1,)*, *optional*, returned when `labels` is provided):
        Classification loss.
    logits (`tf.Tensor` of shape `(batch_size, num_choices)`):
        *num_choices* is the second dimension of the input tensors. (see *input_ids* above).

        Classification scores (before SoftMax).
    mems (`List[tf.Tensor]` of length `config.n_layers`):
        Contains pre-computed hidden-states. Can be used (see `mems` input) to speed up sequential decoding. The
        token ids which have their past given to this model should not be passed as `input_ids` as they have
        already been computed.
    hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape
        `(batch_size, sequence_length, hidden_size)`.

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

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


def _patch_compile(source, filename, mode, flags=0, dont_inherit=False, optimize=-1):
    del filename
    del mode
    source_ast = ast.parse(source)
    final = source_ast.body[-1]
    if isinstance(final, ast.Expr):
        print_it = ast.Expr(lineno=-1, col_offset=-1, value=ast.Call(func=ast.Name(id='_print_if_not_none', ctx=ast.Load(), lineno=-1, col_offset=-1), lineno=-1, col_offset=-1, args=[final], keywords=[]))
        source_ast.body[-1] = print_it
        source = astor.to_source(source_ast)
    return compile(source, filename='dummy.py', mode='exec', flags=flags, dont_inherit=dont_inherit, optimize=optimize)

Patch `doctest.compile` to make doctest to behave like a notebook.

Default settings for doctest are configured to run like a repl: one statement
at a time. The doctest source uses `compile(..., mode="single")`

So to let doctest act like a notebook:

1. We need `mode="exec"` (easy)
2. We need the last expression to be printed (harder).

To print the last expression, just wrap the last expression in
`_print_if_not_none(expr)`. To detect the last expression use `AST`.
If the last node is an expression modify the ast to call
`_print_if_not_none` on it, convert the ast back to source and compile that.

https://docs.python.org/3/library/functions.html#compile

Args:
  source: Can either be a normal string, a byte string, or an AST object.
  filename: Argument should give the file from which the code was read; pass
    some recognizable value if it wasn’t read from a file ('<string>' is
    commonly used).
  mode: [Ignored] always use exec.
  flags: Compiler options.
  dont_inherit: Compiler options.
  optimize: Compiler options.

Returns:
  The resulting code object.


def make_transform_from_reference(n_xyz: torch.Tensor, ca_xyz: torch.Tensor, c_xyz: torch.Tensor, eps: float=1e-20) -> Rigid:
    translation = -1 * ca_xyz
    n_xyz = n_xyz + translation
    c_xyz = c_xyz + translation
    c_x, c_y, c_z = [c_xyz[..., i] for i in range(3)]
    norm = torch.sqrt(eps + c_x ** 2 + c_y ** 2)
    sin_c1 = -c_y / norm
    cos_c1 = c_x / norm
    c1_rots = sin_c1.new_zeros((*sin_c1.shape, 3, 3))
    c1_rots[..., 0, 0] = cos_c1
    c1_rots[..., 0, 1] = -1 * sin_c1
    c1_rots[..., 1, 0] = sin_c1
    c1_rots[..., 1, 1] = cos_c1
    c1_rots[..., 2, 2] = 1
    norm = torch.sqrt(eps + c_x ** 2 + c_y ** 2 + c_z ** 2)
    sin_c2 = c_z / norm
    cos_c2 = torch.sqrt(c_x ** 2 + c_y ** 2) / norm
    c2_rots = sin_c2.new_zeros((*sin_c2.shape, 3, 3))
    c2_rots[..., 0, 0] = cos_c2
    c2_rots[..., 0, 2] = sin_c2
    c2_rots[..., 1, 1] = 1
    c2_rots[..., 2, 0] = -1 * sin_c2
    c2_rots[..., 2, 2] = cos_c2
    c_rots = rot_matmul(c2_rots, c1_rots)
    n_xyz = rot_vec_mul(c_rots, n_xyz)
    _, n_y, n_z = [n_xyz[..., i] for i in range(3)]
    norm = torch.sqrt(eps + n_y ** 2 + n_z ** 2)
    sin_n = -n_z / norm
    cos_n = n_y / norm
    n_rots = sin_c2.new_zeros((*sin_c2.shape, 3, 3))
    n_rots[..., 0, 0] = 1
    n_rots[..., 1, 1] = cos_n
    n_rots[..., 1, 2] = -1 * sin_n
    n_rots[..., 2, 1] = sin_n
    n_rots[..., 2, 2] = cos_n
    rots = rot_matmul(n_rots, c_rots)
    rots = rots.transpose(-1, -2)
    translation = -1 * translation
    rot_obj = Rotation(rot_mats=rots, quats=None)
    return Rigid(rot_obj, translation)

Returns a transformation object from reference coordinates.

Note that this method does not take care of symmetries. If you provide the atom positions in the non-standard
way, the N atom will end up not at [-0.527250, 1.359329, 0.0] but instead at [-0.527250, -1.359329, 0.0]. You
need to take care of such cases in your code.

Args:
    n_xyz: A [*, 3] tensor of nitrogen xyz coordinates.
    ca_xyz: A [*, 3] tensor of carbon alpha xyz coordinates.
    c_xyz: A [*, 3] tensor of carbon xyz coordinates.
Returns:
    A transformation object. After applying the translation and rotation to the reference backbone, the
    coordinates will approximately equal to the input coordinates.


def separable_conv1d(x, depthwise_kernel, pointwise_kernel, strides=1, padding='valid', data_format=None, dilation_rate=1):
    if data_format is None:
        data_format = image_data_format()
    if data_format not in {'channels_first', 'channels_last'}:
        raise ValueError('Unknown data_format: ' + str(data_format))
    if isinstance(strides, int):
        strides = (strides,)
    if isinstance(dilation_rate, int):
        dilation_rate = (dilation_rate,)
    x, tf_data_format = _preprocess_conv1d_input(x, data_format)
    padding = _preprocess_padding(padding)
    if not isinstance(strides, tuple):
        strides = tuple(strides)
    if tf_data_format == 'NWC':
        spatial_start_dim = 1
        strides = (1,) + strides * 2 + (1,)
    else:
        spatial_start_dim = 2
        strides = (1, 1) + strides * 2
    x = array_ops.expand_dims(x, spatial_start_dim)
    depthwise_kernel = array_ops.expand_dims(depthwise_kernel, 0)
    pointwise_kernel = array_ops.expand_dims(pointwise_kernel, 0)
    dilation_rate = (1,) + dilation_rate
    x = nn.separable_conv2d(x, depthwise_kernel, pointwise_kernel, strides=strides, padding=padding, rate=dilation_rate, data_format=tf_data_format)
    x = array_ops.squeeze(x, [spatial_start_dim])
    if data_format == 'channels_first' and tf_data_format == 'NWC':
        x = array_ops.transpose(x, (0, 2, 1))
    return x

1D convolution with separable filters.

Args:
    x: input tensor
    depthwise_kernel: convolution kernel for the depthwise convolution.
    pointwise_kernel: kernel for the 1x1 convolution.
    strides: stride integer.
    padding: string, `"same"` or `"valid"`.
    data_format: string, `"channels_last"` or `"channels_first"`.
    dilation_rate: integer dilation rate.

Returns:
    Output tensor.

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


def __init__(self, name: str, coder: Optional[Coder]=None, combine_fn: Any=None) -> None:
    from apache_beam.transforms.core import CombineFn
    if combine_fn is None:
        if coder is None:
            raise ValueError('combine_fn must be provided')
        else:
            coder, combine_fn = (None, coder)
    self.combine_fn = CombineFn.maybe_from_callable(combine_fn)
    if coder is None:
        coder = self.combine_fn.get_accumulator_coder()
    super().__init__(name, coder)

Initialize the specification for CombiningValue state.

CombiningValueStateSpec(name, combine_fn) -> Coder-inferred combining value
  state spec.
CombiningValueStateSpec(name, coder, combine_fn) -> Combining value state
  spec with coder and combine_fn specified.

Args:
  name (str): The name by which the state is identified.
  coder (Coder): Coder specifying how to encode the values to be combined.
    May be inferred.
  combine_fn (``CombineFn`` or ``callable``): Function specifying how to
    combine the values passed to state.


def _print_primitive_field(self, field_name: str, field: descriptor.FieldDescriptor, value: Any) -> None:
    if proto_utils.field_is_repeated(field):
        string_values = []
        elements = []
        extensions_found = False
        nonnull_values_found = False
        for primitive in value:
            wrapper = self.primitive_handler.primitive_wrapper_from_primitive(primitive)
            string_values.append(wrapper.json_value())
            elements.append(wrapper.get_element())
            nonnull_values_found = nonnull_values_found or wrapper.has_value()
            extensions_found = extensions_found or wrapper.has_element()
        if nonnull_values_found:
            self.generator.add_field(field_name)
            self._print_list(string_values, self.generator.push)
        if extensions_found:
            if nonnull_values_found:
                self.generator.push(',')
                self.generator.add_newline()
            self.generator.add_field(f'_{field_name}')
            self._print_list(elements, self._print)
    elif self.json_format == _FhirJsonFormat.ANALYTIC and field.message_type.name == 'ReferenceId':
        str_value = proto_utils.get_value_at_field(value, 'value')
        self.generator.add_field(field_name, f'"{str_value}"')
    else:
        wrapper = self.primitive_handler.primitive_wrapper_from_primitive(value)
        if wrapper.has_value():
            self.generator.add_field(field_name, wrapper.json_value())
        if wrapper.has_element() and self.json_format == _FhirJsonFormat.PURE:
            if wrapper.has_value():
                self.generator.push(',')
                self.generator.add_newline()
            self.generator.add_field(f'_{field_name}')
            self._print(wrapper.get_element())

Prints the primitive field.

Args:
  field_name: The name of the field.
  field: The FielDescriptor whose contents to print.
  value: The value present at field to print.


def get_enum_value_original_code(enum_value_descriptor: descriptor.EnumValueDescriptor) -> Optional[str]:
    return get_value_for_annotation_extension(enum_value_descriptor, annotations_pb2.fhir_original_code)

Returns the original name if the provided enum value had to be renamed.

Args:
  enum_value_descriptor: The EnumValueDescriptor to examine.

Returns:
  If the code had to be renamed to make a valid enum identifier, this function
  returns the original name. Otherwise returns None.

Raises:
  ValueError: Unable to retrieve options for type: <type>.


def pad(self, image: np.ndarray, size: int, data_format: Optional[Union[str, ChannelDimension]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None):
    old_height, old_width = get_image_size(image, input_data_format)
    pad_height = (old_height // size + 1) * size - old_height
    pad_width = (old_width // size + 1) * size - old_width
    return pad(image, ((0, pad_height), (0, pad_width)), mode='symmetric', data_format=data_format, input_data_format=input_data_format)

Pad an image to make the height and width divisible by `size`.

Args:
    image (`np.ndarray`):
        Image to pad.
    size (`int`):
        The size to make the height and width divisible by.
    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, the channel dimension format 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.

Returns:
    `np.ndarray`: The padded image.


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

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

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

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

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

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

Args:
  func: Test case function to be decorated.

Returns:
  Decorated test case function.


class TFBaseModelOutputWithCLSToken(ModelOutput):
    last_hidden_state: Optional[tf.Tensor] = None
    cls_token_value: Optional[tf.Tensor] = None
    hidden_states: Tuple[tf.Tensor, ...] | None = None

Base class for model's outputs.

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.
    cls_token_value (`tf.Tensor` of shape `(batch_size, 1, hidden_size)`):
        Classification token at the output of the last layer of the model.
    hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape
        `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus
        the initial embedding outputs.


def get_imports(filename: Union[str, os.PathLike]) -> list[str]:
    with open(filename, encoding='utf-8') as f:
        content = f.read()
    imported_modules = set()
    import transformers.utils

    def recursive_look_for_imports(node):
        if isinstance(node, ast.Try):
            return
        elif isinstance(node, ast.If):
            test = node.test
            for condition_node in ast.walk(test):
                if isinstance(condition_node, ast.Call):
                    check_function = getattr(condition_node.func, 'id', '')
                    if check_function.endswith('available') and check_function.startswith('is_flash_attn') or hasattr(transformers.utils.import_utils, check_function):
                        return
        elif isinstance(node, ast.Import):
            for alias in node.names:
                top_module = alias.name.split('.')[0]
                if top_module:
                    imported_modules.add(top_module)
        elif isinstance(node, ast.ImportFrom):
            if node.level == 0 and node.module:
                top_module = node.module.split('.')[0]
                if top_module:
                    imported_modules.add(top_module)
        for child in ast.iter_child_nodes(node):
            recursive_look_for_imports(child)
    tree = ast.parse(content)
    recursive_look_for_imports(tree)
    return sorted(imported_modules)

Extracts all the libraries (not relative imports this time) that are imported in a file.

Args:
    filename (`str` or `os.PathLike`): The module file to inspect.

Returns:
    `list[str]`: The list of all packages required to use the input module.


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. Longformer 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 _create_extra_packages(extra_packages, temp_dir) -> List[beam_runner_api_pb2.ArtifactInformation]:
    resources: List[beam_runner_api_pb2.ArtifactInformation] = []
    staging_temp_dir = tempfile.mkdtemp(dir=temp_dir)
    local_packages: List[str] = []
    for package in extra_packages:
        if not (os.path.basename(package).endswith('.tar') or os.path.basename(package).endswith('.tar.gz') or os.path.basename(package).endswith('.whl') or os.path.basename(package).endswith('.zip')):
            raise RuntimeError('The --extra_package option expects a full path ending with ".tar", ".tar.gz", ".whl" or ".zip" instead of %s' % package)
        if os.path.basename(package).endswith('.whl'):
            _LOGGER.warning('The .whl package "%s" provided in --extra_package must be binary-compatible with the worker runtime environment.' % package)
        if not os.path.isfile(package):
            if Stager._is_remote_path(package):
                _LOGGER.info('Downloading extra package: %s locally before staging', package)
                _, last_component = FileSystems.split(package)
                local_file_path = FileSystems.join(staging_temp_dir, last_component)
                Stager._download_file(package, local_file_path)
            else:
                raise RuntimeError('The file %s cannot be found. It was specified in the --extra_packages command line option.' % package)
        else:
            local_packages.append(package)
    local_packages.extend([FileSystems.join(staging_temp_dir, f) for f in os.listdir(staging_temp_dir)])
    for package in local_packages:
        basename = os.path.basename(package)
        resources.append(Stager._create_file_stage_to_artifact(package, basename))
    with open(os.path.join(temp_dir, EXTRA_PACKAGES_FILE), 'wt') as f:
        for package in local_packages:
            f.write('%s\n' % os.path.basename(package))
    resources.append(Stager._create_file_stage_to_artifact(os.path.join(temp_dir, EXTRA_PACKAGES_FILE), EXTRA_PACKAGES_FILE))
    return resources

Creates a list of local extra packages.

Args:
  extra_packages: Ordered list of local paths to extra packages to be
    staged. Only packages on localfile system and GCS are supported.
  temp_dir: Temporary folder where the resource building can happen.
    Caller is responsible for cleaning up this folder after this function
    returns.

Returns:
  A list of ArtifactInformation of local file paths and file names
  (no paths) for the resources staged. All the files are assumed to be
  staged in staging_location.

Raises:
  RuntimeError: If files specified are not found or do not have expected
    name patterns.


class DTypePolicyMap(DTypePolicy, MutableMapping):

    def __init__(self, default_policy=None, policy_map=None):
        if isinstance(default_policy, DTypePolicyMap):
            raise ValueError('`default_policy` cannot be a `DTypePolicyMap`.')
        if policy_map is not None and (not isinstance(policy_map, dict)):
            raise TypeError(f'If specified, `policy_map` must be a dict. Received: policy_map={policy_map} of type {type(policy_map)}')
        self._default_policy_arg = default_policy
        self._default_policy = dtype_policies.get(default_policy)
        self._policy_map = policy_map or dict()

    @property
    def name(self):
        return 'map_' + self.default_policy._name

    @property
    def default_policy(self):
        """The default dtype policy.

        If `default_policy` is not specified in the constructor, this property
        will be `keras.config.dtype_policy()`.
        """
        return dtype_policies.get(self._default_policy)

    @property
    def variable_dtype(self):
        return self.default_policy.variable_dtype

    @property
    def compute_dtype(self):
        return self.default_policy.compute_dtype

    @property
    def quantization_mode(self):
        return self.default_policy.quantization_mode

    def __getitem__(self, key):
        """Retrieves the corresponding `DTypePolicy` by the string key.

        When there isn't an exact match, all the existing keys in the map
        will be treated as a regex and map against the input key again. When
        there are multiple matches for the regex, an `ValueError` will be
        raised. Returns `self.default_policy` if there isn't any match found.

        Args:
            key: String key to query a `DTypePolicy`.

        Returns:
            Corresponding `DTypePolicy` based on the query.
        """
        if key in self._policy_map:
            return self._policy_map[key]
        matching_keys = []
        for k in self._policy_map:
            if re.search(k, key):
                matching_keys.append(k)
        if len(matching_keys) > 1:
            raise ValueError(f"Path '{key}' matches multiple dtype policy specification keys: {matching_keys}. Please make sure each path only matches at most one dtype policy specification key in the DTypePolicyMap.")
        elif len(matching_keys) == 1:
            return self._policy_map[matching_keys[0]]
        return self.default_policy

    def __setitem__(self, key, policy):
        """Insert `DTypePolicy` to the `DTypePolicyMap`.

        Args:
            key: String key for the `DTypePolicy`.
            policy: The `DTypePolicy`.
        """
        if key in self._policy_map:
            raise ValueError(f'{key} already exist in the DTypePolicyMap with value {self._policy_map[key]}. Please make sure to not use duplicated keys.')
        try:
            policy = dtype_policies.get(policy)
        except Exception:
            raise ValueError(f'Cannot interpret the assigned value by `keras.dtype_policies.get`. Received: {policy} of type {type(policy)}')
        self._policy_map[key] = policy

    def __delitem__(self, key):
        return self._policy_map.pop(key)

    def __contains__(self, key):
        return key in self._policy_map

    def get_config(self):
        from keras.src.saving import serialization_lib
        policy_map = self._policy_map
        if self._default_policy_arg is None:
            for policy in policy_map.values():
                if isinstance(policy, dtype_policies.QuantizedDTypePolicy):
                    policy._name = None
                    policy._source_name = None
                elif isinstance(policy, dtype_policies.DTypePolicy):
                    policy._name = None
        return {'default_policy': self._default_policy_arg, 'policy_map': serialization_lib.serialize_keras_object(policy_map)}

    @classmethod
    def from_config(cls, config, custom_objects=None):
        from keras.src.saving import serialization_lib
        config = config.copy()
        config['policy_map'] = serialization_lib.deserialize_keras_object(config['policy_map'], custom_objects=custom_objects)
        return cls(**config)

    def __len__(self):
        return len(self._policy_map)

    def __iter__(self):
        return iter(self._policy_map)

    def __repr__(self):
        default_policy = self._default_policy.name if self._default_policy is not None else None
        mapping = []
        for k, v in self._policy_map.items():
            mapping.append((k, v.name))
        return f'<DTypePolicyMap at {hex(id(self))} default_policy={default_policy}, mapping={mapping}>'

Dict-like object mapping layer paths to `DTypePolicy` instances.

`DTypePolicyMap` can be used in `get_config` in layers and subclasses to
support a complex configurations of dtype policies.

For example, we can modify `get_config` in `layers.MultiHeadAttention` as
follows to support the mixing of dtype policies, such as quantization.

```python
@keras.saving.register_keras_serializable("MyPackage")
class MyMultiHeadAttention(keras.layers.MultiHeadAttention):
    def get_config(self):
        config = super().get_config()
        dtype_policy_map = dtype_policies.DTypePolicyMap()
        for layer in self._flatten_layers():
            if layer.dtype_policy.quantization_mode is not None:
                dtype_policy_map[layer.path] = layer.dtype_policy
        if len(dtype_policy_map) > 0:
            config.update({"dtype": dtype_policy_map})
        return config
```

Internally, `DTypePolicyMap` uses a string as a key and a `DTypePolicy`
as the value. Typically, the key used for querying is the `Layer.path`.
However, it is also possible to set a regex as the key. See the docstring of
`get` for more details.

See below for a usage example. You can define the naming schema
of the `DTypePolicy`, and then retrieve the corresponding `DTypePolicy`
instance.

```python
dtype_policy_map = DTypePolicyMap()
dtype_policy_map["layer/dense_0"] = DTypePolicy("bfloat16")
dtype_policy_map["layer/dense_1"] = QuantizedDTypePolicy("int8", "bfloat16")

policy_0 = dtype_policy_map["layer/dense_0"]
policy_1 = dtype_policy_map["layer/dense_1"]
policy_2 = dtype_policy_map["layer/dense_2"]  # No hit
assert policy_0 == DTypePolicy("bfloat16")
assert policy_1 == QuantizedDTypePolicy("int8", "bfloat16")
assert policy_2 == keras.config.dtype_policy()
```

Args:
    default_policy: An optional `DTypePolicy` instance specifying the
        default dtype policy. If not specified, the value will default to
        `keras.config.dtype_policy()`.
    policy_map: An optional dict that maps string to `DTypePolicy`
        instances. Defaults to `None`


def copy_assets_to_destination_dir(asset_filename_map, destination_dir, saved_files=None):
    if saved_files is None:
        saved_files = set()
    assets_destination_dir = path_helpers.get_or_create_assets_dir(destination_dir)
    for asset_basename, asset_source_filepath in asset_filename_map.items():
        asset_destination_filepath = file_io.join(compat.as_bytes(assets_destination_dir), compat.as_bytes(asset_basename))
        if file_io.file_exists(asset_source_filepath) and asset_source_filepath != asset_destination_filepath and (asset_destination_filepath not in saved_files):
            file_io.copy(asset_source_filepath, asset_destination_filepath, overwrite=True)
            saved_files.add(asset_destination_filepath)
    tf_logging.info('Assets written to: %s', compat.as_text(assets_destination_dir))

Copy all assets from source path to destination path.

Args:
  asset_filename_map: a dict of filenames used for saving the asset in
    the SavedModel to full paths from which the filenames were derived.
  destination_dir: the destination directory that assets are stored in.
  saved_files: a set of destination filepaths that have already been copied
    and will be skipped


def conv_block(x, growth_rate, name):
    bn_axis = 3 if backend.image_data_format() == 'channels_last' else 1
    x1 = layers.BatchNormalization(axis=bn_axis, epsilon=1.001e-05, name=name + '_0_bn')(x)
    x1 = layers.Activation('relu', name=name + '_0_relu')(x1)
    x1 = layers.Conv2D(4 * growth_rate, 1, use_bias=False, name=name + '_1_conv')(x1)
    x1 = layers.BatchNormalization(axis=bn_axis, epsilon=1.001e-05, name=name + '_1_bn')(x1)
    x1 = layers.Activation('relu', name=name + '_1_relu')(x1)
    x1 = layers.Conv2D(growth_rate, 3, padding='same', use_bias=False, name=name + '_2_conv')(x1)
    x = layers.Concatenate(axis=bn_axis, name=name + '_concat')([x, x1])
    return x

A building block for a dense block.

Args:
    x: input tensor.
    growth_rate: float, growth rate at dense layers.
    name: string, block label.

Returns:
    Output tensor for the block.


def sample_paths(self, initial_forward: types.RealTensor, initial_volatility: types.RealTensor, times: types.RealTensor, time_step: types.RealTensor, num_samples: types.RealTensor=1, random_type: Optional[random.RandomType]=None, seed: Optional[types.IntTensor]=None, skip: types.IntTensor=0, validate_args: bool=False, precompute_normal_draws: bool=True, name: Optional[str]=None):
    name = name or self._name + '_sample_path'
    with tf.name_scope(name):
        initial_forward = tf.convert_to_tensor(initial_forward, self._dtype, name='initial_forward')
        initial_volatility = tf.convert_to_tensor(initial_volatility, self._dtype, name='initial_volatility')
        times = tf.convert_to_tensor(times, self._dtype, name='times')
        time_step = tf.convert_to_tensor(time_step, self._dtype, name='time_step')
        if validate_args:
            self.control_dependencies.append(tf.compat.v1.debugging.Assert(tf.compat.v1.debugging.is_strictly_increasing(times), [times]))
        with tf.compat.v1.control_dependencies(self.control_dependencies):
            initial_forward += self._shift
            if self._enable_unbiased_sampling and (not (_is_callable(self._beta) or _is_callable(self._volvol) or _is_callable(self._rho) or (self._beta == 1))):
                paths = self._sabr_sample_paths(initial_forward, initial_volatility, times, time_step, num_samples, random_type, seed, precompute_normal_draws, skip=skip)
            else:
                paths = super(SabrModel, self).sample_paths(times, num_samples, [initial_forward, initial_volatility], random_type, seed, skip=skip, time_step=time_step, precompute_normal_draws=precompute_normal_draws)
            forwards = tf.expand_dims(paths[:, :, 0] - self._shift, axis=-1)
            volatilities = tf.expand_dims(paths[:, :, 1], axis=-1)
            return tf.concat([forwards, volatilities], axis=-1)

Returns a sample of paths from the process.

Generates samples of paths from the process at the specified time points.

Currently only supports absorbing boundary conditions.

Args:
  initial_forward: Initial value of the forward. A scalar real tensor.
  initial_volatility: Initial value of the volatilities. A scalar real
    tensor.
  times: The times at which the path points are to be evaluated. Rank 1
    `Tensor` of positive real values. This `Tensor` should be sorted in
    ascending order.
  time_step: Positive Python float or a scalar `Tensor `to denote time
    discretization parameter.
  num_samples: Positive scalar `int`. The number of paths to draw.
  random_type: Enum value of `RandomType`. The type of (quasi)-random number
    generator to use to generate the paths.
    Default value: None which maps to the standard pseudo-random numbers.
  seed: Seed for the random number generator. The seed is
    only relevant if `random_type` is one of
    `[STATELESS, PSEUDO, HALTON_RANDOMIZED, PSEUDO_ANTITHETIC,
      STATELESS_ANTITHETIC]`. For `PSEUDO`, `PSEUDO_ANTITHETIC` and
    `HALTON_RANDOMIZED` the seed should be an Python integer. For
    `STATELESS` and  `STATELESS_ANTITHETIC` must be supplied as an integer
    `Tensor` of shape `[2]`.
    Default value: `None` which means no seed is set.
  skip: `int32` 0-d `Tensor`. The number of initial points of the Sobol or
    Halton sequence to skip. Used only when `random_type` is 'SOBOL',
    'HALTON', or 'HALTON_RANDOMIZED', otherwise ignored.
    Default value: `0`.
  validate_args: Python `bool`. When `True`, input `Tensor's` are checked
    for validity despite possibly degrading runtime performance. The checks
    verify that `times` is strictly increasing. When `False` invalid inputs
    may silently render incorrect outputs. Default value: False.
  precompute_normal_draws: Python bool. Indicates whether the noise
    increments are precomputed upfront (see `models.euler_sampling.sample`).
    For `HALTON` and `SOBOL` random types the increments are always
    precomputed. While the resulting graph consumes more memory, the
    performance gains might be significant. Default value: `True`.
  name: str. The name to give this op. If not supplied, default name of
    `sample_paths` is used.

Returns:
  A `Tensor`s of shape [num_samples, k, 2] where `k` is the size of the
  `times`.  The first values in `Tensor` are the simulated forward `F(t)`,
  whereas the second values in `Tensor` are the simulated volatility
  trajectories `V(t)`.


class XLMRobertaConfig(PretrainedConfig):
    model_type = 'xlm-roberta'

    def __init__(self, vocab_size=30522, hidden_size=768, num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072, hidden_act='gelu', hidden_dropout_prob=0.1, attention_probs_dropout_prob=0.1, max_position_embeddings=512, type_vocab_size=2, initializer_range=0.02, layer_norm_eps=1e-12, pad_token_id=1, bos_token_id=0, eos_token_id=2, position_embedding_type='absolute', use_cache=True, classifier_dropout=None, **kwargs):
        super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.hidden_act = hidden_act
        self.intermediate_size = intermediate_size
        self.hidden_dropout_prob = hidden_dropout_prob
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.max_position_embeddings = max_position_embeddings
        self.type_vocab_size = type_vocab_size
        self.initializer_range = initializer_range
        self.layer_norm_eps = layer_norm_eps
        self.position_embedding_type = position_embedding_type
        self.use_cache = use_cache
        self.classifier_dropout = classifier_dropout

This is the configuration class to store the configuration of a [`XLMRobertaModel`] or a [`TFXLMRobertaModel`]. It
is used to instantiate a XLM-RoBERTa 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 XLMRoBERTa
[FacebookAI/xlm-roberta-base](https://huggingface.co/FacebookAI/xlm-roberta-base) architecture.

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 30522):
        Vocabulary size of the XLM-RoBERTa model. Defines the number of different tokens that can be represented by
        the `inputs_ids` passed when calling [`XLMRobertaModel`] or [`TFXLMRobertaModel`].
    hidden_size (`int`, *optional*, defaults to 768):
        Dimensionality of the encoder layers and the pooler layer.
    num_hidden_layers (`int`, *optional*, defaults to 12):
        Number of hidden layers in the Transformer encoder.
    num_attention_heads (`int`, *optional*, defaults to 12):
        Number of attention heads for each attention layer in the Transformer encoder.
    intermediate_size (`int`, *optional*, defaults to 3072):
        Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
    hidden_act (`str` or `Callable`, *optional*, defaults to `"gelu"`):
        The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
        `"relu"`, `"silu"` and `"gelu_new"` are supported.
    hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
        The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
    attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
        The dropout ratio for the attention probabilities.
    max_position_embeddings (`int`, *optional*, defaults to 512):
        The maximum sequence length that this model might ever be used with. Typically set this to something large
        just in case (e.g., 512 or 1024 or 2048).
    type_vocab_size (`int`, *optional*, defaults to 2):
        The vocabulary size of the `token_type_ids` passed when calling [`XLMRobertaModel`] or
        [`TFXLMRobertaModel`].
    initializer_range (`float`, *optional*, defaults to 0.02):
        The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
    layer_norm_eps (`float`, *optional*, defaults to 1e-12):
        The epsilon used by the layer normalization layers.
    position_embedding_type (`str`, *optional*, defaults to `"absolute"`):
        Type of position embedding. Choose one of `"absolute"`, `"relative_key"`, `"relative_key_query"`. For
        positional embeddings use `"absolute"`. For more information on `"relative_key"`, please refer to
        [Self-Attention with Relative Position Representations (Shaw et al.)](https://huggingface.co/papers/1803.02155).
        For more information on `"relative_key_query"`, please refer to *Method 4* in [Improve Transformer Models
        with Better Relative Position Embeddings (Huang et al.)](https://huggingface.co/papers/2009.13658).
    is_decoder (`bool`, *optional*, defaults to `False`):
        Whether the model is used as a decoder or not. If `False`, the model is used as an encoder.
    use_cache (`bool`, *optional*, defaults to `True`):
        Whether or not the model should return the last key/values attentions (not used by all models). Only
        relevant if `config.is_decoder=True`.
    classifier_dropout (`float`, *optional*):
        The dropout ratio for the classification head.

Examples:

```python
>>> from transformers import XLMRobertaConfig, XLMRobertaModel

>>> # Initializing a XLM-RoBERTa FacebookAI/xlm-roberta-base style configuration
>>> configuration = XLMRobertaConfig()

>>> # Initializing a model (with random weights) from the FacebookAI/xlm-roberta-base style configuration
>>> model = XLMRobertaModel(configuration)

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


def _compare_traceback_strings(left, right):
    if left == right:
        return 0
    left = left[len(TRACEBACK_MARKER):] if left else ''
    right = right[len(TRACEBACK_MARKER):] if right else ''
    if left.endswith(right):
        return 1
    elif right.endswith(left):
        return -1
    else:
        return None

Try to compare two traceback strings.

Two traceback strings are comparable if they are equal, or if one ends with
the other. For example, these two tracebacks are comparable:
  Traceback:
    line 1, in <module>
    line 2, in foo
  Traceback:
    line 2, in foo
and the first is greater than the second.

Args:
  left: A string or None.
  right: A string or None.

Returns:
  None if the inputs aren't comparable, else an integer.


def _sort_dump_data_by(self, data, sort_by, reverse):
    if sort_by == SORT_TENSORS_BY_TIMESTAMP:
        return sorted(data, reverse=reverse, key=lambda x: x.timestamp)
    elif sort_by == SORT_TENSORS_BY_DUMP_SIZE:
        return sorted(data, reverse=reverse, key=lambda x: x.dump_size_bytes)
    elif sort_by == SORT_TENSORS_BY_OP_TYPE:
        return sorted(data, reverse=reverse, key=lambda x: self._debug_dump.node_op_type(x.node_name))
    elif sort_by == SORT_TENSORS_BY_TENSOR_NAME:
        return sorted(data, reverse=reverse, key=lambda x: '%s:%d' % (x.node_name, x.output_slot))
    else:
        raise ValueError('Unsupported key to sort tensors by: %s' % sort_by)

Sort a list of DebugTensorDatum in specified order.

Args:
  data: (list of DebugTensorDatum) the data to be sorted.
  sort_by: The field to sort data by.
  reverse: (bool) Whether to use reversed (descending) order.

Returns:
  (list of DebugTensorDatum) in sorted order.

Raises:
  ValueError: given an invalid value of sort_by.


class Module:
    path: str
    target: str
    name: str
    kind: str = 'Local'

    @property
    def full_path(self):
        return path_utils.join(self.path, self.target)

Inferred information about a module.

Attributes:
  path: The path to the module, e.g., foo/.
  target: The filename relative to the path, e.g., bar/baz.py.
  name: The module name, e.g., bar.baz.
  kind: The module kind: Builtin, Direct, Local, or System. See
    https://github.com/google/importlab/blob/main/importlab/resolve.py.
  full_path: The full path to the module (path + target).