class DeiTConfig(PretrainedConfig):
    model_type = 'deit'

    def __init__(self, hidden_size=768, num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072, hidden_act='gelu', hidden_dropout_prob=0.0, attention_probs_dropout_prob=0.0, initializer_range=0.02, layer_norm_eps=1e-12, image_size=224, patch_size=16, num_channels=3, qkv_bias=True, encoder_stride=16, pooler_output_size=None, pooler_act='tanh', **kwargs):
        super().__init__(**kwargs)
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.intermediate_size = intermediate_size
        self.hidden_act = hidden_act
        self.hidden_dropout_prob = hidden_dropout_prob
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.initializer_range = initializer_range
        self.layer_norm_eps = layer_norm_eps
        self.image_size = image_size
        self.patch_size = patch_size
        self.num_channels = num_channels
        self.qkv_bias = qkv_bias
        self.encoder_stride = encoder_stride
        self.pooler_output_size = pooler_output_size if pooler_output_size else hidden_size
        self.pooler_act = pooler_act

This is the configuration class to store the configuration of a [`DeiTModel`]. It is used to instantiate an DeiT
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 DeiT
[facebook/deit-base-distilled-patch16-224](https://huggingface.co/facebook/deit-base-distilled-patch16-224)
architecture.

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

Args:
    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" (i.e., feed-forward) layer in the Transformer encoder.
    hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
        The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
        `"relu"`, `"selu"` and `"gelu_new"` are supported.
    hidden_dropout_prob (`float`, *optional*, defaults to 0.0):
        The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
    attention_probs_dropout_prob (`float`, *optional*, defaults to 0.0):
        The dropout ratio for the attention probabilities.
    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.
    image_size (`int`, *optional*, defaults to 224):
        The size (resolution) of each image.
    patch_size (`int`, *optional*, defaults to 16):
        The size (resolution) of each patch.
    num_channels (`int`, *optional*, defaults to 3):
        The number of input channels.
    qkv_bias (`bool`, *optional*, defaults to `True`):
        Whether to add a bias to the queries, keys and values.
    encoder_stride (`int`, *optional*, defaults to 16):
        Factor to increase the spatial resolution by in the decoder head for masked image modeling.
    pooler_output_size (`int`, *optional*):
       Dimensionality of the pooler layer. If None, defaults to `hidden_size`.
    pooler_act (`str`, *optional*, defaults to `"tanh"`):
       The activation function to be used by the pooler. Keys of ACT2FN are supported for Flax and
       Pytorch, and elements of https://www.tensorflow.org/api_docs/python/tf/keras/activations are
       supported for Tensorflow.

Example:

```python
>>> from transformers import DeiTConfig, DeiTModel

>>> # Initializing a DeiT deit-base-distilled-patch16-224 style configuration
>>> configuration = DeiTConfig()

>>> # Initializing a model (with random weights) from the deit-base-distilled-patch16-224 style configuration
>>> model = DeiTModel(configuration)

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


class KerasFileEditor:

    def __init__(self, filepath):
        self.filepath = filepath
        self.metadata = None
        self.config = None
        self.model = None
        self.console = rich.console.Console(highlight=False)
        if filepath.endswith('.keras'):
            zf = zipfile.ZipFile(filepath, 'r')
            weights_store = H5IOStore(saving_lib._VARS_FNAME + '.h5', archive=zf, mode='r')
            with zf.open(saving_lib._CONFIG_FILENAME, 'r') as f:
                config_json = f.read()
            with zf.open(saving_lib._METADATA_FILENAME, 'r') as f:
                metadata_json = f.read()
            self.config = json.loads(config_json)
            self.metadata = json.loads(metadata_json)
        elif filepath.endswith('.weights.h5'):
            weights_store = H5IOStore(filepath, mode='r')
        else:
            raise ValueError(f'Invalid filename: expected a `.keras` `.weights.h5` extension. Received: filepath={filepath}')
        weights_dict, object_metadata = self._extract_weights_from_store(weights_store.h5_file)
        weights_store.close()
        self.weights_dict = weights_dict
        self.object_metadata = object_metadata
        self.console.print(self._generate_filepath_info(rich_style=True))
        if self.metadata is not None:
            self.console.print(self._generate_metadata_info(rich_style=True))

    def summary(self):
        """Prints the weight structure of the opened file."""
        self._weights_summary_cli()

    def compare(self, reference_model):
        """Compares the opened file to a reference model.

        This method will list all mismatches between the
        currently opened file and the provided reference model.

        Args:
            reference_model: Model instance to compare to.

        Returns:
            Dict with the following keys:
            `'status'`, `'error_count'`, `'match_count'`.
            Status can be `'success'` or `'error'`.
            `'error_count'` is the number of mismatches found.
            `'match_count'` is the number of matching weights found.
        """
        self.console.print('Running comparison')
        ref_spec = {}
        get_weight_spec_of_saveable(reference_model, ref_spec)

        def _compare(target, ref_spec, inner_path, target_name, ref_name, error_count, match_count, checked_paths):
            base_inner_path = inner_path
            for ref_key, ref_val in ref_spec.items():
                inner_path = base_inner_path + '/' + ref_key
                if inner_path in checked_paths:
                    continue
                if ref_key not in target:
                    error_count += 1
                    checked_paths.add(inner_path)
                    if isinstance(ref_val, dict):
                        self.console.print(f'[color(160)]...Object [bold]{inner_path}[/] present in {ref_name}, missing from {target_name}[/]')
                        self.console.print(f'    In {ref_name}, {inner_path} contains the following keys: {list(ref_val.keys())}')
                    else:
                        self.console.print(f'[color(160)]...Weight [bold]{inner_path}[/] present in {ref_name}, missing from {target_name}[/]')
                elif isinstance(ref_val, dict):
                    _error_count, _match_count = _compare(target[ref_key], ref_spec[ref_key], inner_path, target_name, ref_name, error_count=error_count, match_count=match_count, checked_paths=checked_paths)
                    error_count += _error_count
                    match_count += _match_count
                elif target[ref_key].shape != ref_val.shape:
                    error_count += 1
                    checked_paths.add(inner_path)
                    self.console.print(f'[color(160)]...Weight shape mismatch for [bold]{inner_path}[/][/]\n    In {ref_name}: shape={ref_val.shape}\n    In {target_name}: shape={target[ref_key].shape}')
                else:
                    match_count += 1
            return (error_count, match_count)
        checked_paths = set()
        error_count, match_count = _compare(self.weights_dict, ref_spec, inner_path='', target_name='saved file', ref_name='reference model', error_count=0, match_count=0, checked_paths=checked_paths)
        _error_count, _ = _compare(ref_spec, self.weights_dict, inner_path='', target_name='reference model', ref_name='saved file', error_count=0, match_count=0, checked_paths=checked_paths)
        error_count += _error_count
        self.console.print('─────────────────────')
        if error_count == 0:
            status = 'success'
            self.console.print('[color(28)][bold]Comparison successful:[/] saved file is compatible with the reference model[/]')
            if match_count == 1:
                plural = ''
            else:
                plural = 's'
            self.console.print(f'    Found {match_count} matching weight{plural}')
        else:
            status = 'error'
            if error_count == 1:
                plural = ''
            else:
                plural = 's'
            self.console.print(f'[color(160)][bold]Found {error_count} error{plural}:[/] saved file is not compatible with the reference model[/]')
        return {'status': status, 'error_count': error_count, 'match_count': match_count}

    def _edit_object(self, edit_fn, source_name, target_name=None):
        if target_name is not None and '/' in target_name:
            raise ValueError(f"Argument `target_name` should be a leaf name, not a full path name. Received: target_name='{target_name}'")
        if '/' in source_name:
            elements = source_name.split('/')
            weights_dict = self.weights_dict
            for e in elements[:-1]:
                if e not in weights_dict:
                    raise ValueError(f"Path '{source_name}' not found in model.")
                weights_dict = weights_dict[e]
            if elements[-1] not in weights_dict:
                raise ValueError(f"Path '{source_name}' not found in model.")
            edit_fn(weights_dict, source_name=elements[-1], target_name=target_name)
        else:

            def count_occurences(d, name, count=0):
                for k in d:
                    if isinstance(d[k], dict):
                        count += count_occurences(d[k], name, count)
                if name in d:
                    count += 1
                return count
            occurrences = count_occurences(self.weights_dict, source_name)
            if occurrences > 1:
                raise ValueError(f"Name '{source_name}' occurs more than once in the model; try passing a complete path")
            if occurrences == 0:
                raise ValueError(f"Source name '{source_name}' does not appear in the model. Use `editor.weights_summary()` to list all objects.")

            def _edit(d):
                for k in d:
                    if isinstance(d[k], dict):
                        _edit(d[k])
                if source_name in d:
                    edit_fn(d, source_name=source_name, target_name=target_name)
            _edit(self.weights_dict)

    def rename_object(self, object_name, new_name):
        """Rename an object in the file (e.g. a layer).

        Args:
            object_name: String, name or path of the
                object to rename (e.g. `"dense_2"` or
                `"layers/dense_2"`).
            new_name: String, new name of the object.
        """

        def rename_fn(weights_dict, source_name, target_name):
            weights_dict[target_name] = weights_dict[source_name]
            weights_dict.pop(source_name)
        self._edit_object(rename_fn, object_name, new_name)

    def delete_object(self, object_name):
        """Removes an object from the file (e.g. a layer).

        Args:
            object_name: String, name or path of the
                object to delete (e.g. `"dense_2"` or
                `"layers/dense_2"`).
        """

        def delete_fn(weights_dict, source_name, target_name=None):
            weights_dict.pop(source_name)
        self._edit_object(delete_fn, object_name)

    def add_object(self, object_path, weights):
        """Add a new object to the file (e.g. a layer).

        Args:
            object_path: String, full path of the
                object to add (e.g. `"layers/dense_2"`).
            weights: Dict mapping weight names to weight
                values (arrays),
                e.g. `{"0": kernel_value, "1": bias_value}`.
        """
        if not isinstance(weights, dict):
            raise ValueError(f"Argument `weights` should be a dict where keys are weight names (usually '0', '1', etc.) and values are NumPy arrays. Received: type(weights)={type(weights)}")
        if '/' in object_path:
            elements = object_path.split('/')
            partial_path = '/'.join(elements[:-1])
            weights_dict = self.weights_dict
            for e in elements[:-1]:
                if e not in weights_dict:
                    raise ValueError(f"Path '{partial_path}' not found in model.")
                weights_dict = weights_dict[e]
            weights_dict[elements[-1]] = weights
        else:
            self.weights_dict[object_path] = weights

    def delete_weight(self, object_name, weight_name):
        """Removes a weight from an existing object.

        Args:
            object_name: String, name or path of the
                object from which to remove the weight
                (e.g. `"dense_2"` or `"layers/dense_2"`).
            weight_name: String, name of the weight to
                delete (e.g. `"0"`).
        """

        def delete_weight_fn(weights_dict, source_name, target_name=None):
            if weight_name not in weights_dict[source_name]:
                raise ValueError(f'Weight {weight_name} not found in object {object_name}. Weights found: {list(weights_dict[source_name].keys())}')
            weights_dict[source_name].pop(weight_name)
        self._edit_object(delete_weight_fn, object_name)

    def add_weights(self, object_name, weights):
        """Add one or more new weights to an existing object.

        Args:
            object_name: String, name or path of the
                object to add the weights to
                (e.g. `"dense_2"` or `"layers/dense_2"`).
            weights: Dict mapping weight names to weight
                values (arrays),
                e.g. `{"0": kernel_value, "1": bias_value}`.
        """
        if not isinstance(weights, dict):
            raise ValueError(f"Argument `weights` should be a dict where keys are weight names (usually '0', '1', etc.) and values are NumPy arrays. Received: type(weights)={type(weights)}")

        def add_weight_fn(weights_dict, source_name, target_name=None):
            weights_dict[source_name].update(weights)
        self._edit_object(add_weight_fn, object_name)

    def save(self, filepath):
        """Save the edited weights file.

        Args:
            filepath: Path to save the file to.
                Must be a `.weights.h5` file.
        """
        filepath = str(filepath)
        if not filepath.endswith('.weights.h5'):
            raise ValueError(f'Invalid `filepath` argument: expected a `.weights.h5` extension. Received: filepath={filepath}')
        weights_store = H5IOStore(filepath, mode='w')

        def _save(weights_dict, weights_store, inner_path):
            vars_to_create = {}
            for name, value in weights_dict.items():
                if isinstance(value, dict):
                    if value:
                        _save(weights_dict[name], weights_store, inner_path=inner_path + '/' + name)
                else:
                    vars_to_create[name] = value
            if vars_to_create:
                var_store = weights_store.make(inner_path)
                for name, value in vars_to_create.items():
                    var_store[name] = value
        _save(self.weights_dict, weights_store, inner_path='')
        weights_store.close()

    def resave_weights(self, filepath):
        self.save(filepath)

    def _extract_weights_from_store(self, data, metadata=None, inner_path=''):
        metadata = metadata or {}
        object_metadata = {}
        for k, v in data.attrs.items():
            object_metadata[k] = v
        if object_metadata:
            metadata[inner_path] = object_metadata
        result = collections.OrderedDict()
        for key in data.keys():
            inner_path = inner_path + '/' + key
            value = data[key]
            if isinstance(value, h5py.Group):
                if len(value) == 0:
                    continue
                if 'vars' in value.keys() and len(value['vars']) == 0:
                    continue
            if hasattr(value, 'keys'):
                if 'vars' in value.keys():
                    result[key], metadata = self._extract_weights_from_store(value['vars'], metadata=metadata, inner_path=inner_path)
                else:
                    result[key], metadata = self._extract_weights_from_store(value, metadata=metadata, inner_path=inner_path)
            else:
                result[key] = value[()]
        return (result, metadata)

    def _generate_filepath_info(self, rich_style=False):
        if rich_style:
            filepath = f"'{self.filepath}'"
            filepath = f'{summary_utils.highlight_symbol(filepath)}'
        else:
            filepath = f"'{self.filepath}'"
        return f'Keras model file {filepath}'

    def _generate_config_info(self, rich_style=False):
        return pprint.pformat(self.config)

    def _generate_metadata_info(self, rich_style=False):
        version = self.metadata['keras_version']
        date = self.metadata['date_saved']
        if rich_style:
            version = f'{summary_utils.highlight_symbol(version)}'
            date = f'{summary_utils.highlight_symbol(date)}'
        return f'Saved with Keras {version} - date: {date}'

    def _print_weights_structure(self, weights_dict, indent=0, is_first=True, prefix='', inner_path=''):
        for idx, (key, value) in enumerate(weights_dict.items()):
            inner_path = inner_path + '/' + key
            is_last = idx == len(weights_dict) - 1
            if is_first:
                is_first = False
                connector = '> '
            elif is_last:
                connector = '└─ '
            else:
                connector = '├─ '
            if isinstance(value, dict):
                bold_key = summary_utils.bold_text(key)
                object_label = f'{prefix}{connector}{bold_key}'
                if inner_path in self.object_metadata:
                    metadata = self.object_metadata[inner_path]
                    if 'name' in metadata:
                        name = metadata['name']
                        object_label += f" ('{name}')"
                self.console.print(object_label)
                if is_last:
                    appended = '    '
                else:
                    appended = '│   '
                new_prefix = prefix + appended
                self._print_weights_structure(value, indent + 1, is_first=is_first, prefix=new_prefix, inner_path=inner_path)
            elif hasattr(value, 'shape'):
                bold_key = summary_utils.bold_text(key)
                self.console.print(f'{prefix}{connector}{bold_key}:' + f' shape={value.shape}, dtype={value.dtype}')
            else:
                self.console.print(f'{prefix}{connector}{key}: {value}')

    def _weights_summary_cli(self):
        self.console.print('Weights structure')
        self._print_weights_structure(self.weights_dict, prefix=' ' * 2)

    def _weights_summary_interactive(self):

        def _generate_html_weights(dictionary, margin_left=0, font_size=1):
            html = ''
            for key, value in dictionary.items():
                if isinstance(value, dict) and value:
                    html += f'' + '{key}' + _generate_html_weights(value, margin_left + 20, font_size - 1) + ''
                else:
                    html += f'' + f'' + f'{key} : shape={value.shape}' + f', dtype={value.dtype}' + f'
' + f'{display_weight(value)}' + '
' + ''
            return html
        output = 'Weights structure'
        initialize_id_counter()
        output += _generate_html_weights(self.weights_dict)
        ipython.display.display(ipython.display.HTML(output))

Utility to inspect, edit, and resave Keras weights files.

You will find this class useful when adapting
an old saved weights file after having made
architecture changes to a model.

Args:
    filepath: The path to a local file to inspect and edit.

Examples:

```python
editor = KerasFileEditor("my_model.weights.h5")

# Displays current contents
editor.summary()

# Remove the weights of an existing layer
editor.delete_object("layers/dense_2")

# Add the weights of a new layer
editor.add_object("layers/einsum_dense", weights={"0": ..., "1": ...})

# Save the weights of the edited model
editor.resave_weights("edited_model.weights.h5")
```


def cardinality(dataset):
    return gen_dataset_ops.dataset_cardinality(dataset._variant_tensor)

Returns the cardinality of `dataset`, if known.

The operation returns the cardinality of `dataset`. The operation may return
`tf.data.experimental.INFINITE_CARDINALITY` if `dataset` contains an infinite
number of elements or `tf.data.experimental.UNKNOWN_CARDINALITY` if the
analysis fails to determine the number of elements in `dataset` (e.g. when the
dataset source is a file).

>>> dataset = tf.data.Dataset.range(42)
>>> print(tf.data.experimental.cardinality(dataset).numpy())
42
>>> dataset = dataset.repeat()
>>> cardinality = tf.data.experimental.cardinality(dataset)
>>> print((cardinality == tf.data.experimental.INFINITE_CARDINALITY).numpy())
True
>>> dataset = dataset.filter(lambda x: True)
>>> cardinality = tf.data.experimental.cardinality(dataset)
>>> print((cardinality == tf.data.experimental.UNKNOWN_CARDINALITY).numpy())
True

Args:
  dataset: A `tf.data.Dataset` for which to determine cardinality.

Returns:
  A scalar `tf.int64` `Tensor` representing the cardinality of `dataset`. If
  the cardinality is infinite or unknown, the operation returns the named
  constant `INFINITE_CARDINALITY` and `UNKNOWN_CARDINALITY` respectively.


def _ready(op: ops.Operation, sess: session.Session, msg) -> Tuple[bool, Optional[str]]:
    if op is None:
        return (True, None)
    else:
        try:
            ready_value = sess.run(op)
            if ready_value is None or ready_value.dtype == np.int32 or ready_value.size == 0:
                return (True, None)
            else:
                non_initialized_varnames = ', '.join([i.decode('utf-8') for i in ready_value])
                return (False, 'Variables not initialized: ' + non_initialized_varnames)
        except errors.FailedPreconditionError as e:
            if 'uninitialized' not in str(e):
                logging.warning('%s : error [%s]', msg, str(e))
                raise e
            return (False, str(e))

Checks if the model is ready or not, as determined by op.

Args:
  op: An op, either _ready_op or _ready_for_local_init_op, which defines the
    readiness of the model.
  sess: A `Session`.
  msg: A message to log to warning if not ready

Returns:
  A tuple (is_ready, msg), where is_ready is True if ready and False
  otherwise, and msg is `None` if the model is ready, a `String` with the
  reason why it is not ready otherwise.


def enter(self, layer, inputs, build_graph, training, saving=None):
    state = {'layer': layer, 'inputs': inputs, 'build_graph': build_graph, 'training': training, 'saving': saving}
    return CallContextManager(self, state)

Push a Layer and its inputs and state onto the current call context.

Args:
  layer: The `Layer` whose `call` is currently active.
  inputs: The inputs to the currently active `Layer`.
  build_graph: Whether currently inside a Graph or FuncGraph.
  training: Whether currently executing in training or inference mode.
  saving: Whether currently saving to SavedModel.

Returns:
  Context manager.


def _get_pmap_impl(f, devices, has_tpu):
    if has_tpu:
        output_is_list = [False]

        def recorder(args, kwargs, res):
            del args, kwargs
            output_is_list[0] = isinstance(res, list)
            return res
        f = _record_result_type(recorder, f)

    def tf_f(*tf_args):
        """A wrapper for `f` that takes/returns tensors."""
        np_args = _tf_to_np(tf_args)
        np_out = f(*np_args)
        return np_out
    if has_tpu:

        @polymorphic_function.function(autograph=False)
        def fn(inputs):
            res = tpu.replicate(tf_f, inputs)
            if res and isinstance(res[0], list) and (len(res[0]) == 1) and (not output_is_list[0]):
                res = [x[0] for x in res]
            return res
        return fn
    else:
        jit_tf_f = polymorphic_function.function(tf_f, autograph=False)

        @polymorphic_function.function(autograph=False)
        def fn(all_per_device_args):
            """Multi-device function with calls placed on the correct device."""
            results = []
            for per_device_args, device in zip(all_per_device_args, devices):
                with ops.device(device):
                    results.append(jit_tf_f(*per_device_args))
            return results
        return fn

This is a helper function to return the pmap impl.

Args:
  f: a function that takes ndarrays and returns ndarrays.
  devices: a list of strings; the device list.
  has_tpu: boolean; whether `devices` contains TPU devices.

Returns:
  A function that takes tensors and returns tensors.


def standardize_sample_or_class_weights(x_weight, output_names, weight_type):
    if x_weight is None or (isinstance(x_weight, (list, tuple)) and len(x_weight) == 0):
        return [None for _ in output_names]
    if len(output_names) == 1:
        if isinstance(x_weight, (list, tuple)) and len(x_weight) == 1:
            return x_weight
        if isinstance(x_weight, dict) and output_names[0] in x_weight:
            return [x_weight[output_names[0]]]
        else:
            return [x_weight]
    if isinstance(x_weight, (list, tuple)):
        if len(x_weight) != len(output_names):
            raise ValueError('Provided `' + weight_type + '` was a list of ' + str(len(x_weight)) + ' elements, but the model has ' + str(len(output_names)) + ' outputs. You should provide one `' + weight_type + '`array per model output.')
        return x_weight
    if isinstance(x_weight, collections.abc.Mapping):
        generic_utils.check_for_unexpected_keys(weight_type, x_weight, output_names)
        x_weights = []
        for name in output_names:
            x_weights.append(x_weight.get(name))
        return x_weights
    else:
        raise TypeError('The model has multiple outputs, so `' + weight_type + '` should be either a list or a dict. Provided `' + weight_type + '` type not understood: ' + str(x_weight))

Maps `sample_weight` or `class_weight` to model outputs.

Args:
    x_weight: User-provided `sample_weight` or `class_weight` argument.
    output_names: List of output names (strings) in the model.
    weight_type: A string used purely for exception printing.

Returns:
    A list of `sample_weight` or `class_weight` where there are exactly
        one element per model output.

Raises:
    ValueError: In case of invalid user-provided argument.


def get_settable_properties(cls):
    results = []
    for attr, value in vars(cls).items():
        if isinstance(value, property) and value.fset is not None:
            results.append(attr)
    return results

Gets the settable properties of a class.

Only returns the explicitly defined properties with setters.

Args:
  cls: A class in Python.


def transform_function(self, fn, user_context):
    cache_subkey = self.get_caching_key(user_context)
    if self._cache.has(fn, cache_subkey):
        factory = self._cached_factory(fn, cache_subkey)
    else:
        with self._cache_lock:
            if self._cache.has(fn, cache_subkey):
                factory = self._cached_factory(fn, cache_subkey)
            else:
                logging.log(1, '%s is not cached for subkey %s', fn, cache_subkey)
                nodes, ctx = super(PyToPy, self).transform_function(fn, user_context)
                if isinstance(nodes, gast.Lambda):
                    nodes = gast.Assign(targets=[gast.Name(ctx.info.name, ctx=gast.Store(), annotation=None, type_comment=None)], value=nodes)
                else:
                    nodes.name = ctx.info.name
                if logging.has_verbosity(2):
                    logging.log(2, 'Transformed %s:\n\n%s\n', fn, parser.unparse(nodes))
                factory = _PythonFnFactory(ctx.info.name, fn.__code__.co_freevars, self.get_extra_locals())
                factory.create(nodes, ctx.namer, future_features=ctx.info.future_features)
                self._cache[fn][cache_subkey] = factory
    transformed_fn = factory.instantiate(globals_=fn.__globals__, closure=fn.__closure__ or (), defaults=fn.__defaults__, kwdefaults=getattr(fn, '__kwdefaults__', None))
    return (transformed_fn, factory.module, factory.source_map)

Transforms a function. See GenericTranspiler.transform_function.

This overload wraps the parent's `transform_function`, adding caching and
facilities to instantiate the output as a Python object. It also
adds facilities to make new symbols available to the generated Python code,
visible as local variables - see `get_extra_locals`.

Args:
  fn: A function or lambda.
  user_context: An opaque object (may be None) that is forwarded to
    transform_ast, through the ctx.user attribute.

Returns:
  A tuple:
    * A function or lambda with the same signature and closure as `fn`
    * The temporary module into which the transformed function was loaded
    * The source map as a
        Dict[origin_info.LineLocation, origin_info.OriginInfo]


class Adafactor(Optimizer):

    def __init__(self, params, lr=None, eps=(1e-30, 0.001), clip_threshold=1.0, decay_rate=-0.8, beta1=None, weight_decay=0.0, scale_parameter=True, relative_step=True, warmup_init=False):
        if lr is not None and relative_step:
            raise ValueError('Cannot combine manual `lr` and `relative_step=True` options')
        if warmup_init and (not relative_step):
            raise ValueError('`warmup_init=True` requires `relative_step=True`')
        defaults = {'lr': lr, 'eps': eps, 'clip_threshold': clip_threshold, 'decay_rate': decay_rate, 'beta1': beta1, 'weight_decay': weight_decay, 'scale_parameter': scale_parameter, 'relative_step': relative_step, 'warmup_init': warmup_init}
        super().__init__(params, defaults)

    @staticmethod
    def _get_lr(param_group, param_state):
        rel_step_sz = param_group['lr']
        if param_group['relative_step']:
            min_step = 1e-06 * param_state['step'] if param_group['warmup_init'] else 0.01
            rel_step_sz = min(min_step, 1.0 / math.sqrt(param_state['step']))
        param_scale = 1.0
        if param_group['scale_parameter']:
            param_scale = max(param_group['eps'][1], param_state['RMS'])
        return param_scale * rel_step_sz

    @staticmethod
    def _get_options(param_group, param_shape):
        factored = len(param_shape) >= 2
        use_first_moment = param_group['beta1'] is not None
        return (factored, use_first_moment)

    @staticmethod
    def _rms(tensor):
        return tensor.norm(2) / tensor.numel() ** 0.5

    @staticmethod
    def _approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col):
        r_factor = (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1, keepdim=True)).rsqrt_().unsqueeze(-1)
        c_factor = exp_avg_sq_col.unsqueeze(-2).rsqrt()
        return torch.mul(r_factor, c_factor)

    @torch.no_grad()
    def step(self, closure=None):
        """
        Performs a single optimization step

        Arguments:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        """
        loss = None
        if closure is not None:
            loss = closure()
        for group in self.param_groups:
            for p in group['params']:
                if p.grad is None:
                    continue
                grad = p.grad
                if grad.dtype in {torch.float16, torch.bfloat16}:
                    grad = grad.float()
                if grad.is_sparse:
                    raise RuntimeError('Adafactor does not support sparse gradients.')
                state = self.state[p]
                grad_shape = grad.shape
                factored, use_first_moment = self._get_options(group, grad_shape)
                if len(state) == 0:
                    state['step'] = 0
                    if use_first_moment:
                        state['exp_avg'] = torch.zeros_like(grad)
                    if factored:
                        state['exp_avg_sq_row'] = torch.zeros(grad_shape[:-1]).to(grad)
                        state['exp_avg_sq_col'] = torch.zeros(grad_shape[:-2] + grad_shape[-1:]).to(grad)
                    else:
                        state['exp_avg_sq'] = torch.zeros_like(grad)
                    state['RMS'] = 0
                else:
                    if use_first_moment:
                        state['exp_avg'] = state['exp_avg'].to(grad)
                    if factored:
                        state['exp_avg_sq_row'] = state['exp_avg_sq_row'].to(grad)
                        state['exp_avg_sq_col'] = state['exp_avg_sq_col'].to(grad)
                    else:
                        state['exp_avg_sq'] = state['exp_avg_sq'].to(grad)
                p_data_fp32 = p
                if p.dtype in {torch.float16, torch.bfloat16}:
                    p_data_fp32 = p_data_fp32.float()
                state['step'] += 1
                state['RMS'] = self._rms(p_data_fp32)
                lr = self._get_lr(group, state)
                beta2t = 1.0 - math.pow(state['step'], group['decay_rate'])
                update = grad ** 2 + group['eps'][0]
                if factored:
                    exp_avg_sq_row = state['exp_avg_sq_row']
                    exp_avg_sq_col = state['exp_avg_sq_col']
                    exp_avg_sq_row.mul_(beta2t).add_(update.mean(dim=-1), alpha=1.0 - beta2t)
                    exp_avg_sq_col.mul_(beta2t).add_(update.mean(dim=-2), alpha=1.0 - beta2t)
                    update = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col)
                    update.mul_(grad)
                else:
                    exp_avg_sq = state['exp_avg_sq']
                    exp_avg_sq.mul_(beta2t).add_(update, alpha=1.0 - beta2t)
                    update = exp_avg_sq.rsqrt().mul_(grad)
                update.div_((self._rms(update) / group['clip_threshold']).clamp_(min=1.0))
                update.mul_(lr)
                if use_first_moment:
                    exp_avg = state['exp_avg']
                    exp_avg.mul_(group['beta1']).add_(update, alpha=1 - group['beta1'])
                    update = exp_avg
                if group['weight_decay'] != 0:
                    p_data_fp32.add_(p_data_fp32, alpha=-group['weight_decay'] * lr)
                p_data_fp32.add_(-update)
                if p.dtype in {torch.float16, torch.bfloat16}:
                    p.copy_(p_data_fp32)
        return loss

AdaFactor pytorch implementation can be used as a drop in replacement for Adam original fairseq code:
https://github.com/pytorch/fairseq/blob/master/fairseq/optim/adafactor.py

Paper: *Adafactor: Adaptive Learning Rates with Sublinear Memory Cost* https://huggingface.co/papers/1804.04235 Note that
this optimizer internally adjusts the learning rate depending on the `scale_parameter`, `relative_step` and
`warmup_init` options. To use a manual (external) learning rate schedule you should set `scale_parameter=False` and
`relative_step=False`.

Arguments:
    params (`Iterable[nn.parameter.Parameter]`):
        Iterable of parameters to optimize or dictionaries defining parameter groups.
    lr (`float`, *optional*):
        The external learning rate.
    eps (`Tuple[float, float]`, *optional*, defaults to `(1e-30, 0.001)`):
        Regularization constants for square gradient and parameter scale respectively
    clip_threshold (`float`, *optional*, defaults to 1.0):
        Threshold of root mean square of final gradient update
    decay_rate (`float`, *optional*, defaults to -0.8):
        Coefficient used to compute running averages of square
    beta1 (`float`, *optional*):
        Coefficient used for computing running averages of gradient
    weight_decay (`float`, *optional*, defaults to 0.0):
        Weight decay (L2 penalty)
    scale_parameter (`bool`, *optional*, defaults to `True`):
        If True, learning rate is scaled by root mean square
    relative_step (`bool`, *optional*, defaults to `True`):
        If True, time-dependent learning rate is computed instead of external learning rate
    warmup_init (`bool`, *optional*, defaults to `False`):
        Time-dependent learning rate computation depends on whether warm-up initialization is being used

This implementation handles low-precision (FP16, bfloat) values, but we have not thoroughly tested.

Recommended T5 finetuning settings (https://discuss.huggingface.co/t/t5-finetuning-tips/684/3):

    - Training without LR warmup or clip_threshold is not recommended.

       - use scheduled LR warm-up to fixed LR
       - use clip_threshold=1.0 (https://huggingface.co/papers/1804.04235)
    - Disable relative updates
    - Use scale_parameter=False
    - Additional optimizer operations like gradient clipping should not be used alongside Adafactor

Example:

```python
Adafactor(model.parameters(), scale_parameter=False, relative_step=False, warmup_init=False, lr=1e-3)
```

Others reported the following combination to work well:

```python
Adafactor(model.parameters(), scale_parameter=True, relative_step=True, warmup_init=True, lr=None)
```

When using `lr=None` with [`Trainer`] you will most likely need to use [`~optimization.AdafactorSchedule`]
scheduler as following:

```python
from transformers.optimization import Adafactor, AdafactorSchedule

optimizer = Adafactor(model.parameters(), scale_parameter=True, relative_step=True, warmup_init=True, lr=None)
lr_scheduler = AdafactorSchedule(optimizer)
trainer = Trainer(..., optimizers=(optimizer, lr_scheduler))
```

Usage:

```python
# replace AdamW with Adafactor
optimizer = Adafactor(
    model.parameters(),
    lr=1e-3,
    eps=(1e-30, 1e-3),
    clip_threshold=1.0,
    decay_rate=-0.8,
    beta1=None,
    weight_decay=0.0,
    relative_step=False,
    scale_parameter=False,
    warmup_init=False,
)
```


def pad(self, images: 'torch.Tensor', size: int) -> 'torch.Tensor':
    height, width = get_image_size(images, ChannelDimension.FIRST)
    pad_height = (height // size + 1) * size - height
    pad_width = (width // size + 1) * size - width
    return F.pad(images, (0, 0, pad_width, pad_height), padding_mode='symmetric')

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

Args:
    images (`torch.Tensor`):
        Images to pad.
    size (`int`):
        The size to make the height and width divisible by.

Returns:
    `torch.Tensor`: The padded images.


def register_converter(src_type: Union[Type[Any], Tuple[Type[Any], ...]], dest_type: Union[Type[Any], Tuple[Type[Any], ...]], convert_fn: Callable[[Any], Any]) -> None:
    _TYPE_CONVERTER_REGISTRY.register(src_type, dest_type, convert_fn)

Register converter from source type to destination type.

Examples::

  # Add converter from int to float.
  pg.typing.register_converter(int, float, float)

  assert pg.typing.Float().apply(1) is 1.0

  # Add converter from a dict to class A.
  def from_dict(d):
    return A(**d)

  assert isinstance(pg.typing.Object(A).apply({'x': 1, 'y': 2}), A)

Args:
    src_type: Source value type.
    dest_type: Target value type.
    convert_fn: Function that performs the conversion, in signature
      (src_type) -> dest_type.


def _infer_num_gpus_per_worker(devices):
    if _is_device_list_single_worker(devices):
        return sum((1 for d in devices if _is_gpu_device(d)))
    else:
        device_dict = _group_device_list(devices)
        num_gpus = None
        for _, devices_in_task in device_dict.items():
            for device_in_task in devices_in_task:
                if num_gpus is None:
                    num_gpus = sum((1 for d in device_in_task if _is_gpu_device(d)))
                elif num_gpus != sum((1 for d in device_in_task if _is_gpu_device(d))):
                    raise ValueError('All workers should have the same number of GPUs.')
                for d in device_in_task:
                    d_spec = tf_device.DeviceSpec.from_string(d)
                    if d_spec.device_type == 'GPU' and d_spec.device_index >= num_gpus:
                        raise ValueError('GPU `device_index` on a worker should be consecutive and start from 0.')
        return num_gpus

Infers the number of GPUs on each worker.

Currently to make multi-worker cross device ops work, we need all workers to
have the same number of GPUs.

Args:
  devices: a list of device strings, can be either local devices or remote
    devices.

Returns:
  number of GPUs per worker.

Raises:
  ValueError if workers have different number of GPUs or GPU indices are not
  consecutive and starting from 0.


def _delegate_property(keras_tensor_cls, property_name):
    property_access = property(lambda self: InstanceProperty(property_name)(self))
    setattr(keras_tensor_cls, property_name, property_access)

Register property on a KerasTensor class.

Calling this multiple times with the same arguments should be a no-op.

This method exposes a property on the KerasTensor class that will use an
`InstanceProperty` layer to access the property on the represented
intermediate values in the model.

Args:
  keras_tensor_cls: The KerasTensor subclass that should expose the property.
  property_name: The name of the property to expose and delegate to the
    represented (Composite)Tensor.


def call(self, input_ids: TFModelInputType | None=None, inputs_embeds: np.ndarray | tf.Tensor | None=None, attention_mask: np.ndarray | tf.Tensor | None=None, position_ids: np.ndarray | tf.Tensor | None=None, encoder_hidden_states: np.ndarray | tf.Tensor | None=None, encoder_attention_mask: np.ndarray | tf.Tensor | None=None, head_mask: np.ndarray | tf.Tensor | None=None, cross_attn_head_mask: np.ndarray | tf.Tensor | None=None, past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, training: Optional[bool]=False) -> Union[TFBaseModelOutputWithPastAndCrossAttentions, Tuple[tf.Tensor]]:
    if input_ids is not None and inputs_embeds is not None:
        raise ValueError('You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time')
    elif input_ids is not None:
        input_shape = shape_list(input_ids)
    elif inputs_embeds is not None:
        input_shape = shape_list(inputs_embeds)[:-1]
    else:
        raise ValueError('You have to specify either decoder_input_ids or decoder_inputs_embeds')
    past_key_values_length = shape_list(past_key_values[0][0])[2] if past_key_values is not None else 0
    if position_ids is None:
        positions = self.embed_positions(input_shape, past_key_values_length)
    else:
        positions = self.embed_positions(input_shape, position_ids=position_ids)
    if inputs_embeds is None:
        check_embeddings_within_bounds(input_ids, self.embed_tokens.input_dim)
        inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
    hidden_states = inputs_embeds
    if input_shape[-1] > 1:
        combined_attention_mask = _make_causal_mask(input_shape, past_key_values_length=past_key_values_length)
    else:
        combined_attention_mask = _expand_mask(tf.ones((input_shape[0], input_shape[1] + past_key_values_length)), tgt_len=input_shape[-1])
    if attention_mask is not None:
        combined_attention_mask = combined_attention_mask + _expand_mask(attention_mask, tgt_len=input_shape[-1])
    if encoder_hidden_states is not None and encoder_attention_mask is not None:
        encoder_attention_mask = _expand_mask(encoder_attention_mask, tgt_len=input_shape[-1])
    hidden_states = self.layernorm_embedding(hidden_states + positions)
    hidden_states = self.dropout(hidden_states, training=training)
    all_hidden_states = () if output_hidden_states else None
    all_self_attns = () if output_attentions else None
    all_cross_attns = () if output_attentions and encoder_hidden_states is not None else None
    present_key_values = () if use_cache else None
    for attn_mask_name, attn_mask in [('head_mask', head_mask), ('cross_attn_head_mask', cross_attn_head_mask)]:
        if attn_mask is not None:
            tf.debugging.assert_equal(shape_list(attn_mask)[0], len(self.layers), message=f'The {attn_mask_name} should be specified for {len(self.layers)} layers, but it is for {shape_list(attn_mask)[0]}.')
    for idx, decoder_layer in enumerate(self.layers):
        if output_hidden_states:
            all_hidden_states += (hidden_states,)
        dropout_probability = random.uniform(0, 1)
        if training and dropout_probability < self.layerdrop:
            continue
        past_key_value = past_key_values[idx] if past_key_values is not None else None
        hidden_states, layer_self_attn, layer_cross_attn, present_key_value = decoder_layer(hidden_states, attention_mask=combined_attention_mask, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, layer_head_mask=head_mask[idx] if head_mask is not None else None, cross_attn_layer_head_mask=cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None, past_key_value=past_key_value)
        if use_cache:
            present_key_values += (present_key_value,)
        if output_attentions:
            all_self_attns += (layer_self_attn,)
            if encoder_hidden_states is not None:
                all_cross_attns += (layer_cross_attn,)
    if output_hidden_states:
        all_hidden_states += (hidden_states,)
    if not return_dict:
        return (hidden_states, present_key_values, all_hidden_states, all_self_attns, all_cross_attns)
    else:
        return TFBaseModelOutputWithPastAndCrossAttentions(last_hidden_state=hidden_states, past_key_values=present_key_values, hidden_states=all_hidden_states, attentions=all_self_attns, cross_attentions=all_cross_attns)

Args:
    input_ids (`tf.Tensor` of shape `(batch_size, sequence_length)`):
        Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you
        provide it.

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

        [What are input IDs?](../glossary#input-ids)
    attention_mask (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
        Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:

        - 1 for tokens that are **not masked**,
        - 0 for tokens that are **masked**.

        [What are attention masks?](../glossary#attention-mask)
    position_ids (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
        Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the
        range `[0, config.max_position_embeddings - 1]`.
    encoder_hidden_states (`tf.Tensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):
        Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention
        of the decoder.
    encoder_attention_mask (`tf.Tensor` of shape `(batch_size, encoder_sequence_length)`, *optional*):
        Mask to avoid performing cross-attention on padding tokens indices of encoder input_ids. Mask values
        selected in `[0, 1]`:

        - 1 for tokens that are **not masked**,
        - 0 for tokens that are **masked**.

        [What are attention masks?](../glossary#attention-mask)
    head_mask (`tf.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
        Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:

        - 1 indicates the head is **not masked**,
        - 0 indicates the head is **masked**.

    cross_attn_head_mask (`tf.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
        Mask to nullify selected heads of the cross-attention modules. Mask values selected in `[0, 1]`:

        - 1 indicates the head is **not masked**,
        - 0 indicates the head is **masked**.

    past_key_values (`Tuple[Tuple[tf.Tensor]]` of length `config.n_layers` with each tuple having 2 tuples each of which has 2 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
        Contains precomputed key and value hidden-states of the attention blocks. Can be used to speed up
        decoding.

        If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those
        that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of
        all `decoder_input_ids` of shape `(batch_size, sequence_length)`.
    inputs_embeds (`tf.tTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
        Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
        This is useful if you want more control over how to convert `input_ids` indices into associated vectors
        than the model's internal embedding lookup matrix.
    output_attentions (`bool`, *optional*):
        Whether or not to return the attentions tensors of all attention layers. See `attentions` under
        returned tensors for more detail.
    output_hidden_states (`bool`, *optional*):
        Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
        for more detail.
    return_dict (`bool`, *optional*):
        Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.


def _normalize_var_keyword_hint(hint, arg_name):
    if not hint or type(hint) != dict:
        raise TypeCheckError('Unexpected VAR_KEYWORD value: %s' % hint)
    keys = list(hint.keys())
    values = list(hint.values())
    if len(values) == 1 and keys[0] == arg_name and isinstance(values[0], typehints.DictConstraint):
        return values[0]
    else:
        return typehints.Dict[str, typehints.Union[values]]

Converts a var_keyword hint into Dict[<key type>, <value type>] form.

Args:
  hint: (dict) Should either contain a pair (arg_name,
    Dict[<key type>, <value type>]), or one or more possible types for the
    value.
  arg_name: (str) The keyword receiving this hint.

Raises:
  TypeCheckError if hint does not have the right form.


def resize(self, image: torch.Tensor, size: SizeDict, interpolation: 'F.InterpolationMode'=None, **kwargs) -> torch.Tensor:
    interpolation = interpolation if interpolation is not None else F.InterpolationMode.BICUBIC
    if size.shortest_edge:
        shortest_edge = int(256 / 224 * size['shortest_edge'])
        new_size = get_resize_output_image_size(image, size=shortest_edge, default_to_square=False, input_data_format=ChannelDimension.FIRST)
    elif size.height and size.width:
        new_size = (size.height, size.width)
    else:
        raise ValueError(f"Size dict must have keys 'height' and 'width' or 'shortest_edge'. Got {size.keys()} {size.keys()}.")
    return F.resize(image, size=new_size, interpolation=interpolation, **kwargs)

Resize an image.

If size is a dict with keys "width" and "height", the image will be resized to `(size["height"],
size["width"])`.

If size is a dict with key "shortest_edge", the shortest edge value `c` is rescaled to `int(c * (256/224))`.
The smaller edge of the image will be matched to this value i.e, if height > width, then image will be rescaled
to `(size["shortest_egde"] * height / width, size["shortest_egde"])`.

Args:
    image (`torch.Tensor`):
        Image to resize.
    size (`SizeDict`):
        Size of the output image after resizing. If size is a dict with keys "width" and "height", the image
        will be resized to (height, width). If size is a dict with key "shortest_edge", the shortest edge value
        `c` is rescaled to int(`c` * (256/224)). The smaller edge of the image will be matched to this value
        i.e, if height > width, then image will be rescaled to (size * height / width, size).
    interpolation (`InterpolationMode`, *optional*, defaults to `InterpolationMode.BICUBIC`):
        Resampling filter to use when resiizing the image.


def actual_360(*, start_date, end_date, schedule_info=None, dtype=None, name=None):
    del schedule_info
    with tf.name_scope(name or 'actual_360'):
        end_date = dt.convert_to_date_tensor(end_date)
        start_date = dt.convert_to_date_tensor(start_date)
        dtype = dtype or tf.constant(0.0).dtype
        actual_days = tf.cast(start_date.days_until(end_date), dtype=dtype)
        return actual_days / 360

Computes the year fraction between the specified dates.

The actual/360 convention specifies the year fraction between the start and
end date as the actual number of days between the two dates divided by 360.

Note that the schedule info is not needed for this convention and is ignored
if supplied.

For more details see:
https://en.wikipedia.org/wiki/Day_count_convention#Actual/360

Args:
  start_date: A `DateTensor` object of any shape.
  end_date: A `DateTensor` object of compatible shape with `start_date`.
  schedule_info: The schedule info. Ignored for this convention.
  dtype: The dtype of the result. Either `tf.float32` or `tf.float64`. If not
    supplied, `tf.float32` is returned.
  name: Python `str` name prefixed to ops created by this function. If not
    supplied, `actual_360` is used.

Returns:
  A real `Tensor` of supplied `dtype` and shape of `start_date`. The year
  fraction between the start and end date as computed by Actual/360
  convention.


def __init__(self, project_name=None, bq_table=None, bq_dataset=None, publish_to_bq=False, influxdb_options: Optional['InfluxDBMetricsPublisherOptions']=None, namespace=None, filters=None):
    self._namespace = namespace
    self.publishers: List[MetricsPublisher] = []
    self.publishers.append(ConsoleMetricsPublisher())
    bq_check = project_name and bq_table and bq_dataset and publish_to_bq
    if bq_check:
        bq_publisher = BigQueryMetricsPublisher(project_name, bq_table, bq_dataset)
        self.publishers.append(bq_publisher)
    if influxdb_options and influxdb_options.validate():
        self.publishers.append(InfluxDBMetricsPublisher(influxdb_options))
    else:
        _LOGGER.info('Missing InfluxDB options. Metrics will not be published to InfluxDB')
    self.filters = filters

Initializes :class:`MetricsReader` .

Args:
  project_name (str): project with BigQuery where metrics will be saved
  bq_table (str): BigQuery table where metrics will be saved
  bq_dataset (str): BigQuery dataset where metrics will be saved
  namespace (str): Namespace of the metrics
  filters: MetricFilter to query only filtered metrics


def get_parameters(model: nn.Module) -> Iterable[torch.Tensor]:
    for name, module in model._modules.items():
        for attr_name, attr in module.__dict__.items():
            if isinstance(attr, torch.Tensor) and attr.requires_grad:
                yield attr
        for param in get_parameters(module):
            yield param

Get all parameters from a model by iterating over its modules.
This is an alternative to model.parameters() that works with DTensor models.

Args:
    model (nn.Module): The model to get parameters from

Returns:
    Iterable[torch.Tensor]: An iterator over all parameters in the model


def add(self, arg, tag=None, name=None, aggregate=None, index_override=None):
    if tag is None:
        if aggregate is not None:
            raise ValueError('You must specify `tag` if using aggregate.')
        global_index = self._get_new_global_index(index_override)
        sort_index = None
    else:
        if aggregate is None:
            raise ValueError('You must specify `aggregate` if using tag.')
        if tag not in self._tag_to_global_index:
            self._tag_to_global_index[tag] = self._get_new_global_index(index_override)
            self._tag_to_next_sort_index[tag] = 0
        elif index_override and index_override != self._tag_to_global_index[tag]:
            raise ValueError('Tag %r was called with two indices %r and %r' % (tag, index_override, self._tag_to_global_index[tag]))
        global_index = self._tag_to_global_index[tag]
        sort_index = self._tag_to_next_sort_index[tag]
        self._tag_to_next_sort_index[tag] += 1
    uuid = self._unique_function_id
    name = '%s-%s-%s-%r-%r-%s' % (self._node_name_prefix, self._function_name, uuid, global_index, sort_index, name)
    identity_op = _array_ops.identity(arg, name=name)
    identity_op.op._set_attr(OpHint.FUNCTION_NAME_ATTR, _attr_value_pb2.AttrValue(s=_compat.as_bytes(self._function_name)))
    identity_op.op._set_attr(OpHint.FUNCTION_UUID_ATTR, _attr_value_pb2.AttrValue(s=_compat.as_bytes(self._unique_function_id)))
    identity_op.op._set_attr(self._attr_name, _attr_value_pb2.AttrValue(i=global_index))
    identity_op.op._set_attr(OpHint.FUNCTION_LEVEL_ATTR, _attr_value_pb2.AttrValue(i=self._level))
    if self._children_inputs_mappings:
        identity_op.op._set_attr(OpHint.CHILDREN_INPUTS_MAPPINGS, _attr_value_pb2.AttrValue(s=_compat.as_bytes(_json.dumps(self._children_inputs_mappings))))
    if sort_index is not None:
        identity_op.op._set_attr(OpHint.FUNCTION_SORT_INDEX_ATTR, _attr_value_pb2.AttrValue(i=sort_index))
    if aggregate is not None:
        identity_op.op._set_attr(OpHint.FUNCTION_AGGREGATE_ATTR, _attr_value_pb2.AttrValue(s=_compat.as_bytes(aggregate)))
    return identity_op

Return a wrapped tensor of an input tensor as an argument.

Args:
  arg: A TensorFlow tensor that should be considered an argument.
  tag: String tag to identify arguments that should be packed.
  name: Name of argument. This is included in the Identity hint op names.
  aggregate: Strategy to aggregate.
  Acceptable values are OpHint.AGGREGATE_FIRST, OpHint.AGGREGATE_LAST,
    and OpHint.AGGREGATE_STACK.
    Note, aggregate is only valid if tag is specified.
  index_override: Specify what input/output index should this be in the
    final stub. i.e. add(arg0, index=1); add(arg1, index=0) will make the
    final stub be as stub_func(inputs[arg1, arg0], outputs=[]) rather than
    the default call order based ordering.

Returns:
  A tensor representing the wrapped argument.

Raises:
  ValueError: When indices are not consistent.


def from_concrete_functions(cls, funcs, trackable_obj=None):
    TFLiteConverterBase._set_original_model_type(conversion_metadata_fb.ModelType.TF_CONCRETE_FUNCTIONS)
    if trackable_obj is None:
        logging.warning('Please consider providing the trackable_obj argument in the from_concrete_functions. Providing without the trackable_obj argument is deprecated and it will use the deprecated conversion path.')
    for func in funcs:
        if not isinstance(func, _function.ConcreteFunction):
            message = 'This function takes in a list of ConcreteFunction.'
            if isinstance(func, _def_function.Function):
                message += ' To get the ConcreteFunction from a Function, call get_concrete_function.'
            raise ValueError(message)
    return cls(funcs, trackable_obj)

Creates a TFLiteConverter object from ConcreteFunctions.

Args:
  funcs: List of TensorFlow ConcreteFunctions. The list should not contain
    duplicate elements. Currently converter can only convert a single
    ConcreteFunction. Converting multiple functions is under development.
  trackable_obj:   An `AutoTrackable` object (typically `tf.module`)
    associated with `funcs`. A reference to this object needs to be
    maintained so that Variables do not get garbage collected since
    functions have a weak reference to Variables.

Returns:
  TFLiteConverter object.

Raises:
  Invalid input type.


def create_model_card(self, output_dir, model_name: str, language: Optional[str]=None, license: Optional[str]=None, tags: Optional[str]=None, finetuned_from: Optional[str]=None, tasks: Optional[str]=None, dataset_tags: Optional[Union[str, List[str]]]=None, dataset: Optional[Union[str, List[str]]]=None, dataset_args: Optional[Union[str, List[str]]]=None):
    from .modelcard import TrainingSummary
    training_summary = TrainingSummary.from_keras(self, keras_history=self.history, language=language, license=license, tags=tags, model_name=model_name, finetuned_from=finetuned_from, tasks=tasks, dataset_tags=dataset_tags, dataset=dataset, dataset_args=dataset_args)
    model_card = training_summary.to_model_card()
    with open(os.path.join(output_dir, 'README.md'), 'w') as f:
        f.write(model_card)

Creates a draft of a model card using the information available to the `Trainer`.

Args:
    output_dir (`str` or `os.PathLike`):
        The folder in which to create the model card.
    model_name (`str`, *optional*):
        The name of the model.
    language (`str`, *optional*):
        The language of the model (if applicable)
    license (`str`, *optional*):
        The license of the model. Will default to the license of the pretrained model used, if the original
        model given to the `Trainer` comes from a repo on the Hub.
    tags (`str` or `List[str]`, *optional*):
        Some tags to be included in the metadata of the model card.
    finetuned_from (`str`, *optional*):
        The name of the model used to fine-tune this one (if applicable). Will default to the name of the repo
        of the original model given to the `Trainer` (if it comes from the Hub).
    tasks (`str` or `List[str]`, *optional*):
        One or several task identifiers, to be included in the metadata of the model card.
    dataset_tags (`str` or `List[str]`, *optional*):
        One or several dataset tags, to be included in the metadata of the model card.
    dataset (`str` or `List[str]`, *optional*):
        One or several dataset identifiers, to be included in the metadata of the model card.
    dataset_args (`str` or `List[str]`, *optional*):
       One or several dataset arguments, to be included in the metadata of the model card.


def get_tensor_sharding(tensor):
    if isinstance(tensor, resource_variable_ops.BaseResourceVariable) and context.xla_sharding_for_resource_variables_enabled():
        sharding = tensor._get_xla_sharding()
        if sharding is None:
            return None
        else:
            return sharding.SerializeToString()
    try:
        return get_op_sharding(tensor.op)
    except AttributeError:
        return None

Returns sharding attribute of a Tensor.

Args:
  tensor: a Tensor.

Returns:
  The attribute representing XLA sharding on tensor's op.


def log10(x):
    if any_symbolic_tensors((x,)):
        return Log10().symbolic_call(x)
    return backend.numpy.log10(x)

Return the base 10 logarithm of the input tensor, element-wise.

Args:
    x: Input tensor.

Returns:
    Output tensor, element-wise base 10 logarithm of `x`.


def find_executable_batch_size(function: Optional[callable]=None, starting_batch_size: int=128, auto_find_batch_size: bool=False):
    if function is None:
        return functools.partial(find_executable_batch_size, starting_batch_size=starting_batch_size, auto_find_batch_size=auto_find_batch_size)
    if auto_find_batch_size:
        requires_backends(find_executable_batch_size, 'accelerate')
        from accelerate.utils import find_executable_batch_size as accelerate_find_executable_batch_size
        return accelerate_find_executable_batch_size(function=function, starting_batch_size=starting_batch_size)
    return functools.partial(function, batch_size=starting_batch_size)

Args:
A basic decorator that will try to execute `function`. If it fails from exceptions related to out-of-memory or
CUDNN, the batch size is cut in half and passed to `function`. `function` must take in a `batch_size` parameter as
its first argument.
    function (`callable`, *optional*)
        A function to wrap
    starting_batch_size (`int`, *optional*)
        The batch size to try and fit into memory
    auto_find_batch_size (`bool`, *optional*)
        If False, will just execute `function`


def trace(service_addr, logdir, duration_ms, worker_list='', num_tracing_attempts=3, options=None):
    if duration_ms <= 0:
        raise errors.InvalidArgumentError(None, None, 'duration_ms must be greater than zero.')
    opts = dict(options._asdict()) if options is not None else {}
    _pywrap_profiler_plugin.trace(_strip_addresses(service_addr, _GRPC_PREFIX), logdir, worker_list, True, duration_ms, num_tracing_attempts, opts)

Sends gRPC requests to one or more profiler servers to perform on-demand profiling.

This method will block the calling thread until it receives responses from all
servers or until deadline expiration. Both single host and multiple host
profiling are supported on CPU, GPU, and TPU.
The profiled results will be saved by each server to the specified TensorBoard
log directory (i.e. the directory you save your model checkpoints). Use the
TensorBoard profile plugin to view the visualization and analysis results.

Args:
  service_addr: A comma delimited string of gRPC addresses of the workers to
    profile.
    e.g. service_addr='grpc://localhost:6009'
         service_addr='grpc://10.0.0.2:8466,grpc://10.0.0.3:8466'
         service_addr='grpc://localhost:12345,grpc://localhost:23456'
  logdir: Path to save profile data to, typically a TensorBoard log directory.
    This path must be accessible to both the client and server.
    e.g. logdir='gs://your_tb_dir'
  duration_ms: Duration of tracing or monitoring in milliseconds. Must be
    greater than zero.
  worker_list: An optional TPU only configuration. The list of workers to
    profile in the current session.
  num_tracing_attempts: Optional. Automatically retry N times when no trace
    event is collected (default 3).
  options: profiler.experimental.ProfilerOptions namedtuple for miscellaneous
    profiler options.

Raises:
  InvalidArgumentError: For when arguments fail validation checks.
  UnavailableError: If no trace event was collected.

Example usage (CPU/GPU):

```python
  # Start a profiler server before your model runs.
  tf.profiler.experimental.server.start(6009)
  # (Model code goes here).
  # Send gRPC request to the profiler server to collect a trace of your model.
  tf.profiler.experimental.client.trace('grpc://localhost:6009',
                                        '/nfs/tb_log', 2000)
```

Example usage (Multiple GPUs):

```python
  # E.g. your worker IP addresses are 10.0.0.2, 10.0.0.3, 10.0.0.4, and you
  # would like to schedule start of profiling 1 second from now, for a
  # duration of 2 seconds.
  options['delay_ms'] = 1000
  tf.profiler.experimental.client.trace(
      'grpc://10.0.0.2:8466,grpc://10.0.0.3:8466,grpc://10.0.0.4:8466',
      'gs://your_tb_dir',
      2000,
      options=options)
```

Example usage (TPU):

```python
  # Send gRPC request to a TPU worker to collect a trace of your model. A
  # profiler service has been started in the TPU worker at port 8466.
  # E.g. your TPU IP address is 10.0.0.2 and you want to profile for 2 seconds
  # .
  tf.profiler.experimental.client.trace('grpc://10.0.0.2:8466',
                                        'gs://your_tb_dir', 2000)
```

Example usage (Multiple TPUs):

```python
  # Send gRPC request to a TPU pod to collect a trace of your model on
  # multiple TPUs. A profiler service has been started in all the TPU workers
  # at the port 8466.
  # E.g. your TPU IP addresses are 10.0.0.2, 10.0.0.3, 10.0.0.4, and you want
  # to profile for 2 seconds.
  tf.profiler.experimental.client.trace(
      'grpc://10.0.0.2:8466',
      'gs://your_tb_dir',
      2000,
      '10.0.0.2:8466,10.0.0.3:8466,10.0.0.4:8466')
```

Launch TensorBoard and point it to the same logdir you provided to this API.

```shell
  # logdir can be gs://your_tb_dir as in the above examples.
  $ tensorboard --logdir=/tmp/tb_log
```

Open your browser and go to localhost:6006/#profile to view profiling results.


def initialize_multi_client_cluster(job_name: str, dtensor_jobs: List[str], client_id: int, collective_leader: str, port: Optional[int]=None, gpu_use_nccl_communication: bool=False, enable_coordination_service: bool=True):
    assert context.executing_eagerly()
    if not collective_leader.startswith('/job:'):
        collective_leader = '/job:' + collective_leader
    context.context().configure_collective_ops(use_nccl_communication=gpu_use_nccl_communication, collective_leader=collective_leader)
    if enable_coordination_service:
        context.context().configure_coordination_service(service_type='standalone', service_leader=collective_leader)
    config_proto = context.get_config()
    cluster_def = cluster_pb2.ClusterDef()
    cluster_def.job.add(name=job_name, tasks=dict(enumerate(dtensor_jobs)))
    server_def = tensorflow_server_pb2.ServerDef(cluster=cluster_def, default_session_config=config_proto, job_name=job_name, task_index=client_id, protocol=remote_utils.get_default_communication_protocol(), port=port)
    server_def.default_session_config.rpc_options.num_channels_per_target = 4
    server_def.default_session_config.experimental.recv_buf_max_chunk = -1
    logging.info('Enabling collectives with server_def: %s', server_def)
    context.context().enable_collective_ops(server_def)
    context.ensure_initialized()

Initialize GRPC servers and collectives for multi-client DTensor setup.

This function can be used to initialize a multi-client cluster and enable
collective ops. GRPC servers are necessary in the multi-client mode, even
when the number of clientis is 1.

NOTE: this function must be called in an eager context.

Args:
  job_name: The job name used by all clients in the DTensor cluster.
  dtensor_jobs: A list of the DTensor client jobs participating in the
    cluster. Must be strings of the form "hostname:port".
  client_id: The ID of the DTensor client this function is being called in.
  collective_leader: The job/task that will be used to run collectives.
  port: The port this client's GRPC server will run on. If omitted, use the
    port from dtensor_jobs for this client.
  gpu_use_nccl_communication: if True, configure TensorFlow to use NCCL by
    default.
  enable_coordination_service: If true, enable distributed coordination
    service to make sure that workers know the devices on each other, a
    prerequisite for data transfer through cross-worker rendezvous.

Raises:
  RuntimeError: If running inside a tf.function.


def forward(self, data: torch.Tensor, observed_indicator: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    ts_sum = (data * observed_indicator).abs().sum(self.dim, keepdim=True)
    num_observed = observed_indicator.sum(self.dim, keepdim=True)
    scale = ts_sum / torch.clamp(num_observed, min=1)
    if self.default_scale is None:
        batch_sum = ts_sum.sum(dim=0)
        batch_observations = torch.clamp(num_observed.sum(0), min=1)
        default_scale = torch.squeeze(batch_sum / batch_observations)
    else:
        default_scale = self.default_scale * torch.ones_like(scale)
    scale = torch.where(num_observed > 0, scale, default_scale)
    scale = torch.clamp(scale, min=self.minimum_scale)
    scaled_data = data / scale
    if not self.keepdim:
        scale = scale.squeeze(dim=self.dim)
    return (scaled_data, torch.zeros_like(scale), scale)

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


def get_metrics_namespace(self) -> str:
    return 'BeamML_Sklearn'

Returns:
   A namespace for metrics collected by the RunInference transform.


def __init__(self, experimental_io_device=None, experimental_enable_async_checkpoint=False, experimental_write_callbacks=None, enable_async=False, experimental_skip_slot_variables=False, experimental_sharding_callback=None):
    self.experimental_io_device = experimental_io_device
    self.enable_async = experimental_enable_async_checkpoint or enable_async
    self.experimental_enable_async_checkpoint = self.enable_async
    if experimental_write_callbacks is not None:
        for callback in experimental_write_callbacks:
            assert len(inspect.signature(callback).parameters) <= 1
    self.experimental_write_callbacks = experimental_write_callbacks
    if experimental_sharding_callback is not None:
        if not isinstance(experimental_sharding_callback, sharding_util.ShardingCallback):
            raise ValueError(f'The experimental_sharding_callback checkpoint optionmust be of type ShardingCallback. The option providedwas of type {type(experimental_sharding_callback)}.')
    self.experimental_sharding_callback = experimental_sharding_callback
    self.experimental_skip_slot_variables = experimental_skip_slot_variables

Creates an object that stores options for a Checkpoint.

Args:
  experimental_io_device: string. Applies in a distributed setting.
    Tensorflow device to use to access the filesystem. If `None` (default)
    then for each variable the filesystem is accessed from the CPU:0 device
    of the host where that variable is assigned. If specified, the
    filesystem is instead accessed from that device for all variables.  This
    is for example useful if you want to save to a local directory, such as
    "/tmp" when running in a distributed setting. In that case pass a device
    for the host where the "/tmp" directory is accessible.
  experimental_enable_async_checkpoint: bool Type. Deprecated, please use
    the enable_async option.
  experimental_write_callbacks: List[Callable]. A list of callback functions
    that will be executed after each saving event finishes (i.e. after
    `save()` or `write()`). For async checkpoint, the callbacks will be
    executed only after the async thread finishes saving.  The return values
    of the callback(s) will be ignored. The callback(s) can optionally take
    the `save_path` (the result of `save()` or `write()`) as an argument.
    The callbacks will be executed in the same order of this list after the
    checkpoint has been written.
  enable_async: bool Type. Indicates whether async checkpointing is enabled.
    Default is False, i.e., no async checkpoint.  Async checkpoint moves the
    checkpoint file writing off the main thread, so that the model can
    continue to train while the checkpoint file writing runs in the
    background. Async checkpoint reduces TPU device idle cycles and speeds
    up model training process, while memory consumption may increase.
  experimental_skip_slot_variables: bool Type. If true, ignores slot
    variables during restore. Context: TPU Embedding layers for Serving do
    not properly restore slot variables. This option is a way to omit
    restoring slot variables which are not required for Serving usecase
    anyways.(b/315912101)
  experimental_sharding_callback: `tf.train.experimental.ShardingCallback`.
    A pre-made or custom callback that determines how checkpoints are
    sharded on disk. Pre-made callback options are
    `tf.train.experimental.ShardByDevicePolicy` and
    `tf.train.experimental.MaxShardSizePolicy`. You may also write a custom
    callback, see `tf.train.experimental.ShardingCallback`.


def amplitude_to_db(spectrogram: np.ndarray, reference: float=1.0, min_value: float=1e-05, db_range: Optional[float]=None) -> np.ndarray:
    if reference <= 0.0:
        raise ValueError('reference must be greater than zero')
    if min_value <= 0.0:
        raise ValueError('min_value must be greater than zero')
    reference = max(min_value, reference)
    spectrogram = np.clip(spectrogram, a_min=min_value, a_max=None)
    spectrogram = 20.0 * (np.log10(spectrogram) - np.log10(reference))
    if db_range is not None:
        if db_range <= 0.0:
            raise ValueError('db_range must be greater than zero')
        spectrogram = np.clip(spectrogram, a_min=spectrogram.max() - db_range, a_max=None)
    return spectrogram

Converts an amplitude spectrogram to the decibel scale. This computes `20 * log10(spectrogram / reference)`, using
basic logarithm properties for numerical stability.

The motivation behind applying the log function on the (mel) spectrogram is that humans do not hear loudness on a
linear scale. Generally to double the perceived volume of a sound we need to put 8 times as much energy into it.
This means that large variations in energy may not sound all that different if the sound is loud to begin with.
This compression operation makes the (mel) spectrogram features match more closely what humans actually hear.

Args:
    spectrogram (`np.ndarray`):
        The input amplitude (mel) spectrogram.
    reference (`float`, *optional*, defaults to 1.0):
        Sets the input spectrogram value that corresponds to 0 dB. For example, use `np.max(spectrogram)` to set
        the loudest part to 0 dB. Must be greater than zero.
    min_value (`float`, *optional*, defaults to `1e-5`):
        The spectrogram will be clipped to this minimum value before conversion to decibels, to avoid taking
        `log(0)`. The default of `1e-5` corresponds to a minimum of -100 dB. Must be greater than zero.
    db_range (`float`, *optional*):
        Sets the maximum dynamic range in decibels. For example, if `db_range = 80`, the difference between the
        peak value and the smallest value will never be more than 80 dB. Must be greater than zero.

Returns:
    `np.ndarray`: the spectrogram in decibels


def dropout_add(x: torch.Tensor, residual: torch.Tensor, prob: float, training: bool) -> torch.Tensor:
    out = F.dropout(x, p=prob, training=training)
    out = residual + out
    return out

Dropout add function

Args:
    x (`torch.tensor`):
        input tensor
    residual (`torch.tensor`):
        residual tensor
    prob (`float`):
        dropout probability
    training (`bool`):
        training mode


def __init__(self, live_api_processor: live_model.LiveProcessor, chattiness: float=1.0, unsafe_string_list: list[str] | None=None):
    self._processor = live_api_processor
    self._chattiness = chattiness
    self._commentator = CommentatorStateMachine()
    self.ttfts = collections.deque(maxlen=50)
    self._unsafe_string_list = unsafe_string_list
    if unsafe_string_list is not None:
        pattern = '|'.join((re.escape(s) for s in unsafe_string_list))
        self._processor += text.MatchProcessor(pattern=pattern, substream_input='output_transcription', substream_output='unsafe_regex', remove_from_input_stream=False, flush_fn=lambda x: x.get_metadata('generation_complete') or x.get_metadata('interrupted') or x.get_metadata('interrupt_request') or x.get_metadata('turn_complete') or x.get_metadata('go_away'))

Initializes the processor.

Args:
  live_api_processor: The live API processor to use.
  chattiness: Probability of triggering a comment when the model has
    finished talking or every 3 seconds. Set to 0 to disable commenting.
  unsafe_string_list: The strings to use for unsafe content. If None, the
    processor will not block unsafe content. If set, the processor will
    interrupt itself when it sees the string in the output.


class PegasusXDecoder(PegasusXPreTrainedModel):

    def __init__(self, config: PegasusXConfig, embed_tokens: Optional[nn.Embedding]=None):
        super().__init__(config)
        self.dropout = config.dropout
        self.layerdrop = config.decoder_layerdrop
        self.max_target_positions = config.max_position_embeddings
        embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
        padding_idx = config.pad_token_id
        if embed_tokens is not None:
            self.embed_tokens = embed_tokens
        else:
            self.embed_tokens = PegasusXScaledWordEmbedding(config.vocab_size, config.d_model, padding_idx=padding_idx, embed_scale=embed_scale)
        self.embed_positions = PegasusXSinusoidalPositionalEmbedding(config.d_model)
        self.layers = nn.ModuleList([PegasusXDecoderLayer(config, layer_idx=i) for i in range(config.decoder_layers)])
        self.layer_norm = nn.LayerNorm(config.d_model)
        self.gradient_checkpointing = False
        self.post_init()

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, value):
        self.embed_tokens = value

    def forward(self, input_ids=None, attention_mask=None, encoder_hidden_states=None, encoder_attention_mask=None, past_key_values=None, inputs_embeds=None, use_cache=None, output_attentions=None, output_hidden_states=None, return_dict=None, cache_position=None):
        """
        Args:
            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
                Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you
                provide it.

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

                [What are input IDs?](../glossary#input-ids)
            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:

                - 1 for tokens that are **not masked**,
                - 0 for tokens that are **masked**.

                [What are attention masks?](../glossary#attention-mask)
            encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):
                Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention
                of the decoder.
            encoder_attention_mask (`torch.LongTensor` of shape `(batch_size, encoder_sequence_length)`, *optional*):
                Mask to avoid performing cross-attention on padding tokens indices of encoder input_ids. Mask values
                selected in `[0, 1]`:

                - 1 for tokens that are **not masked**,
                - 0 for tokens that are **masked**.

                [What are attention masks?](../glossary#attention-mask)

            past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
                Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
                shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of
                shape `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.

                Contains pre-computed hidden-states (key and values in the self-attention blocks and in the
                cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.

                If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those
                that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of
                all `decoder_input_ids` of shape `(batch_size, sequence_length)`.
            inputs_embeds (`torch.FloatTensor` of
                shape `(batch_size, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing
                `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more
                control over how to convert `input_ids` indices into associated vectors than the model's internal
                embedding lookup matrix.
            output_attentions (`bool`, *optional*):
                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
                returned tensors for more detail.
            output_hidden_states (`bool`, *optional*):
                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
                for more detail.
            return_dict (`bool`, *optional*):
                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
            cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
                Indices depicting the position of the input sequence tokens in the sequence. It is used to update the
                cache in the correct position and to infer the complete sequence length.
        """
        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
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError('You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time')
        elif input_ids is not None:
            input = input_ids
            input_shape = input.shape
            input_ids = input_ids.view(-1, input_shape[-1])
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
            input = inputs_embeds[:, :, -1]
        else:
            raise ValueError('You have to specify either decoder_input_ids or decoder_inputs_embeds')
        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input)
        if self.gradient_checkpointing and self.training:
            if use_cache:
                logger.warning_once('`use_cache=True` is incompatible with gradient checkpointing`. Setting `use_cache=False`...')
                use_cache = False
        return_legacy_cache = False
        if use_cache and (not isinstance(past_key_values, Cache)):
            return_legacy_cache = True
            logger.warning_once('Passing a tuple of `past_key_values` is deprecated and will be removed in Transformers v4.58.0. You should pass an instance of `EncoderDecoderCache` instead, e.g. `past_key_values=EncoderDecoderCache.from_legacy_cache(past_key_values)`.')
            past_key_values = EncoderDecoderCache.from_legacy_cache(past_key_values)
        batch_size, seq_length = inputs_embeds.size()[:-1]
        past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0
        if cache_position is None:
            cache_position = torch.arange(past_key_values_length, past_key_values_length + seq_length, device=inputs_embeds.device)
        if attention_mask is None and (not is_torchdynamo_compiling()):
            mask_seq_length = past_key_values_length + seq_length
            attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device)
        self_attn_cache = past_key_values.self_attention_cache if isinstance(past_key_values, EncoderDecoderCache) else past_key_values
        causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position, self_attn_cache)
        encoder_attention_mask = self._update_cross_attn_mask(encoder_hidden_states, encoder_attention_mask, input_shape, inputs_embeds)
        position_ids = cache_position.unsqueeze(1)
        position_ids = self.embed_positions(inputs_embeds, past_key_values_length, position_ids)
        position_ids = position_ids.to(inputs_embeds.device)
        hidden_states = inputs_embeds + position_ids
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None
        all_cross_attentions = () if output_attentions and encoder_hidden_states is not None else None
        next_decoder_cache = None
        for idx, decoder_layer in enumerate(self.layers):
            if output_hidden_states:
                all_hidden_states += (hidden_states,)
            if self.training:
                dropout_probability = torch.rand([])
                if dropout_probability < self.layerdrop:
                    continue
            if self.gradient_checkpointing and self.training:
                layer_outputs = self._gradient_checkpointing_func(decoder_layer.__call__, hidden_states, causal_mask, encoder_hidden_states, encoder_attention_mask, None, output_attentions, use_cache, cache_position)
            else:
                layer_outputs = decoder_layer(hidden_states, attention_mask=causal_mask, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, past_key_value=past_key_values, output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position)
            hidden_states = layer_outputs[0]
            if use_cache:
                next_decoder_cache = layer_outputs[3 if output_attentions else 1]
            if output_attentions:
                all_self_attns += (layer_outputs[1],)
                if encoder_hidden_states is not None:
                    all_cross_attentions += (layer_outputs[2],)
        hidden_states = self.layer_norm(hidden_states)
        if output_hidden_states:
            all_hidden_states += (hidden_states,)
        next_cache = next_decoder_cache if use_cache else None
        if return_legacy_cache:
            next_cache = past_key_values.to_legacy_cache()
        if not return_dict:
            return tuple((v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_cross_attentions] if v is not None))
        return BaseModelOutputWithPastAndCrossAttentions(last_hidden_state=hidden_states, past_key_values=next_cache, hidden_states=all_hidden_states, attentions=all_self_attns, cross_attentions=all_cross_attentions)

Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`PegasusDecoderLayer`]

Args:
    config: PegasusXConfig
    embed_tokens (nn.Embedding): output embedding


def gather_nd(params: ragged_tensor.RaggedOrDense, indices: ragged_tensor.RaggedOrDense, batch_dims=0, name=None, bad_indices_policy=''):
    if not isinstance(batch_dims, int) or batch_dims != 0:
        raise ValueError('batch_dims != 0 is not supported for ragged gather yet.')
    if not (ragged_tensor.is_ragged(params) or ragged_tensor.is_ragged(indices)):
        return array_ops.gather_nd(params, indices, name=name, bad_indices_policy=bad_indices_policy)
    if bad_indices_policy not in ('', 'DEFAULT'):
        raise ValueError('non-default bad_indices_policy not supported for ragged gather')
    with ops.name_scope(name, 'RaggedGatherNd', [params, indices]):
        params = ragged_tensor.convert_to_tensor_or_ragged_tensor(params, name='params')
        indices = ragged_tensor.convert_to_tensor_or_ragged_tensor(indices, name='indices')
        params, indices = ragged_tensor.match_row_splits_dtypes(params, indices)
        indices_shape = indices.shape
        indices_ndims = indices_shape.ndims
        if indices_ndims is None:
            raise ValueError('indices.rank be statically known.')
        if indices_ndims == 0:
            raise ValueError('indices.rank must be at least 1.')
        if ragged_tensor.is_ragged(indices) and indices_ndims == indices.ragged_rank + 1:
            raise ValueError('The innermost dimension of indices may not be ragged')
        index_size = tensor_shape.dimension_value(indices_shape[-1])
        if index_size is None:
            raise ValueError('indices.shape[-1] must be statically known.')
        if indices_ndims > 2:
            indices_is_dense = not ragged_tensor.is_ragged(indices)
            if indices_is_dense:
                indices = ragged_tensor.RaggedTensor.from_tensor(indices, ragged_rank=indices_ndims - 2, row_splits_dtype=params.row_splits.dtype)
            result = indices.with_flat_values(gather_nd(params, indices.flat_values))
            if indices_is_dense and ragged_tensor.is_ragged(result) and (result.ragged_rank == indices_ndims - 2):
                result = ragged_tensor.RaggedTensor.to_tensor(result)
            return result
        assert not ragged_tensor.is_ragged(indices)
        assert ragged_tensor.is_ragged(params)
        if index_size == 0:
            params_ndims = params.ragged_rank + array_ops.rank(params.flat_values)
            for dim in range(indices_ndims - 1):
                params = ragged_array_ops.expand_dims(params, axis=0)
            multiples = array_ops.concat([array_ops.shape(indices)[:-1], array_ops.ones([params_ndims], dtypes.int32)], axis=0)
            return ragged_array_ops.tile(params, multiples)
        elif index_size == 1:
            flattened_index_tuples = array_ops.reshape(indices, [-1])
            return gather(params, flattened_index_tuples)
        else:
            indices = math_ops.cast(indices, params.row_splits.dtype)
            flattened_index_tuples = array_ops.gather(params.row_splits, indices[..., 0])
            flattened_index_tuples += indices[..., 1]
            flattened_params = params.values
            for dim in range(2, index_size):
                if not ragged_tensor.is_ragged(flattened_params):
                    flattened_index_tuples = array_ops.expand_dims(flattened_index_tuples, axis=1)
                    flattened_index_tuples = array_ops.concat([flattened_index_tuples, indices[..., dim:]], axis=1)
                    return array_ops.gather_nd(flattened_params, flattened_index_tuples)
                flattened_index_tuples = array_ops.gather(flattened_params.row_starts(), flattened_index_tuples)
                flattened_index_tuples += indices[..., dim]
                flattened_params = flattened_params.values
            return gather(flattened_params, flattened_index_tuples)

Gather slices from `params` using `n`-dimensional indices.

This operation is similar to `gather`, but it uses the innermost dimension
of `indices` to define a slice into `params`.  In particular, if:

* `indices` has shape `[A1...AN, I]`
* `params` has shape `[B1...BM]`

Then:

* `result` has shape `[A1...AN, B_{I+1}...BM]`.
* `result[a1...aN] = params[indices[a1...aN, :]]`

Args:
  params: A potentially ragged tensor with shape `[A1...AN, I]`.
  indices: A potentially ragged tensor with shape `[B1...BM]`.
  batch_dims: Must be zero.
  name: A name for the operation (optional).
  bad_indices_policy: A string. If `""` or `"DEFAULT"`, the default behavior
    is used (error on CPU and ignore on GPU). If `"IGNORE"`, the bad indices
    are ignored and 0 is stored in the

Returns:
  A potentially ragged tensor with shape `[A1...AN, B_{I+1}...BM]`.

#### Examples:

>>> params = tf.ragged.constant(
...     [ [ ['000', '001'], ['010'              ]          ],
...       [ ['100'       ], ['110', '111', '112'], ['120'] ],
...       [ [            ], ['210'              ]          ] ])

>>> # Gather 2D slices from a 3D tensor
>>> tf.gather_nd(params, [[2], [0]])
<tf.RaggedTensor [[[], [b'210']], [[b'000', b'001'], [b'010']]]>

>>> # Gather 1D slices from a 3D tensor
>>> tf.gather_nd(params, [[2, 1], [0, 0]])
<tf.RaggedTensor [[b'210'], [b'000', b'001']]>

>>> # Gather scalars from a 3D tensor
>>> tf.gather_nd(params, [[0, 0, 1], [1, 1, 2]]).numpy()
array([b'001', b'112'], dtype=object)


def rescale(self, image: 'torch.Tensor', scale: float, **kwargs) -> 'torch.Tensor':
    return image * scale

Rescale an image by a scale factor. image = image * scale.

Args:
    image (`torch.Tensor`):
        Image to rescale.
    scale (`float`):
        The scaling factor to rescale pixel values by.

Returns:
    `torch.Tensor`: The rescaled image.


def wait_for_other_workers(self):
    if not self._worker_barrier:
        return
    self._worker_barrier.wait()

Waits for other workers to reach the same call to this method.

Raises:
  ValueError: if `worker_barrier` is not passed to the __init__ method.


def variance(self, name='variance'):
    with self._name_scope(name):
        try:
            return self._variance()
        except NotImplementedError as original_exception:
            try:
                return math_ops.square(self._stddev())
            except NotImplementedError:
                raise original_exception

Variance.

Variance is defined as,

```none
Var = E[(X - E[X])**2]
```

where `X` is the random variable associated with this distribution, `E`
denotes expectation, and `Var.shape = batch_shape + event_shape`.

Args:
  name: Python `str` prepended to names of ops created by this function.

Returns:
  variance: Floating-point `Tensor` with shape identical to
    `batch_shape + event_shape`, i.e., the same shape as `self.mean()`.


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

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

Modified Bessel function of order 0.

>>> tf.math.special.bessel_y0([0.5, 1., 2., 4.]).numpy()
array([-0.44451873,  0.08825696,  0.51037567, -0.01694074], dtype=float32)

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

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

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