class FlavaConfig(PretrainedConfig):
    model_type = 'flava'
    sub_configs = {'text_config': FlavaTextConfig, 'image_config': FlavaImageConfig, 'multimodal_config': FlavaMultimodalConfig, 'image_codebook_config': FlavaImageCodebookConfig}

    def __init__(self, image_config: Optional[Dict[str, Any]]=None, text_config: Optional[Dict[str, Any]]=None, multimodal_config: Optional[Dict[str, Any]]=None, image_codebook_config: Optional[Dict[str, Any]]=None, hidden_size: int=768, layer_norm_eps: float=1e-12, projection_dim: int=768, init_codebook: bool=True, logit_scale_init_value: float=2.6592, initializer_range: float=0.02, ce_ignore_index: int=-100, mim_weight: float=1.0, mlm_weight: float=1.0, global_contrastive_weight: float=1.0, itm_weight: float=1.0, mmm_image_weight: float=1.0, mmm_text_weight: float=1.0, global_backprop_contrastive: bool=True, skip_unmasked_multimodal_encoder: bool=True, return_loss: bool=True, **kwargs):
        text_config_dict = kwargs.pop('text_config_dict', None)
        image_config_dict = kwargs.pop('image_config_dict', None)
        multimodal_config_dict = kwargs.pop('multimodal_config_dict', None)
        image_codebook_config_dict = kwargs.pop('image_codebook_config_dict', None)
        super().__init__(**kwargs)
        if text_config_dict is not None:
            if text_config is None:
                text_config = {}
            _text_config_dict = FlavaTextConfig(**text_config_dict).to_dict()
            for key, value in _text_config_dict.items():
                if key in text_config and value != text_config[key] and (key not in ['transformers_version']):
                    if key in text_config_dict:
                        message = f'`{key}` is found in both `text_config_dict` and `text_config` but with different values. The value `text_config_dict["{key}"]` will be used instead.'
                    else:
                        message = f'`text_config_dict` is provided which will be used to initialize `FlavaTextConfig`. The value `text_config["{key}"]` will be overridden.'
                    logger.info(message)
            text_config.update(_text_config_dict)
        if image_config_dict is not None:
            if image_config is None:
                image_config = {}
            _image_config_dict = FlavaImageConfig(**image_config_dict).to_dict()
            if 'id2label' in _image_config_dict:
                _image_config_dict['id2label'] = {str(key): value for key, value in _image_config_dict['id2label'].items()}
            for key, value in _image_config_dict.items():
                if key in image_config and value != image_config[key] and (key not in ['transformers_version']):
                    if key in image_config_dict:
                        message = f'`{key}` is found in both `image_config_dict` and `image_config` but with different values. The value `image_config_dict["{key}"]` will be used instead.'
                    else:
                        message = f'`image_config_dict` is provided which will be used to initialize `FlavaImageConfig`. The value `image_config["{key}"]` will be overridden.'
                    logger.info(message)
            image_config.update(_image_config_dict)
        if multimodal_config_dict is not None:
            if multimodal_config is None:
                multimodal_config = {}
            _multimodal_config_dict = FlavaMultimodalConfig(**multimodal_config_dict).to_dict()
            for key, value in _multimodal_config_dict.items():
                if key in multimodal_config and value != multimodal_config[key] and (key not in ['transformers_version']):
                    if key in multimodal_config_dict:
                        message = f'`{key}` is found in both `multimodal_config_dict` and `multimodal_config` but with different values. The value `multimodal_config_dict["{key}"]` will be used instead.'
                    else:
                        message = f'`multimodal_config_dict` is provided which will be used to initialize `FlavaMultimodalConfig`. The value `multimodal_config["{key}"]` will be overridden.'
                    logger.info(message)
            multimodal_config.update(_multimodal_config_dict)
        if image_codebook_config_dict is not None:
            if image_codebook_config is None:
                image_codebook_config = {}
            _image_codebook_config_dict = FlavaImageCodebookConfig(**image_codebook_config_dict).to_dict()
            for key, value in _image_codebook_config_dict.items():
                if key in image_codebook_config and value != image_codebook_config[key] and (key not in ['transformers_version']):
                    if key in image_codebook_config_dict:
                        message = f'`{key}` is found in both `image_codebook_config_dict` and `image_codebook_config` but with different values. The value `image_codebook_config_dict["{key}"]` will be used instead.'
                    else:
                        message = f'`image_codebook_config_dict` is provided which will be used to initialize `FlavaImageCodebookConfig`. The value `image_codebook_config["{key}"]` will be overridden.'
                    logger.info(message)
            image_codebook_config.update(_image_codebook_config_dict)
        if image_config is None:
            image_config = {}
            logger.info('`image_config` is `None`. initializing the `FlavaImageConfig` with default values.')
        if text_config is None:
            text_config = {}
            logger.info('`text_config` is `None`. Initializing the `FlavaTextConfig` with default values.')
        if multimodal_config is None:
            multimodal_config = {}
            logger.info('`multimodal_config` is `None`. initializing the `FlavaMultimodalConfig` with default values.')
        if image_codebook_config is None:
            image_codebook_config = {}
            logger.info('`image_codebook_config` is `None`. initializing the `FlavaImageCodebookConfig` with default values.')
        self.image_config = FlavaImageConfig(**image_config)
        self.text_config = FlavaTextConfig(**text_config)
        self.multimodal_config = FlavaMultimodalConfig(**multimodal_config)
        self.image_codebook_config = FlavaImageCodebookConfig(**image_codebook_config)
        self.projection_dim = projection_dim
        self.init_codebook = init_codebook
        self.hidden_size = hidden_size
        self.layer_norm_eps = layer_norm_eps
        self.initializer_range = initializer_range
        self.logit_scale_init_value = logit_scale_init_value
        self.initializer_factor = 1.0
        self.ce_ignore_index = ce_ignore_index
        self.mim_weight = mim_weight
        self.mlm_weight = mlm_weight
        self.global_contrastive_weight = global_contrastive_weight
        self.itm_weight = itm_weight
        self.mmm_image_weight = mmm_image_weight
        self.mmm_text_weight = mmm_text_weight
        self.global_backprop_contrastive = global_backprop_contrastive
        self.skip_unmasked_multimodal_encoder = skip_unmasked_multimodal_encoder
        self.return_loss = return_loss

    @classmethod
    def from_configs(cls, image_config: FlavaImageConfig, text_config: FlavaTextConfig, multimodal_config: FlavaMultimodalConfig, image_codebook_config: FlavaImageCodebookConfig, **kwargs):
        """
        Instantiate a [`FlavaConfig`] (or a derived class) from flava text model configuration, flava image model
        configuration, flava multimodal model and flava codebook model configuration.

        Returns:
            [`FlavaConfig`]: An instance of a configuration object
        """
        return cls(image_config=image_config.to_dict(), text_config=text_config.to_dict(), multimodal_config=multimodal_config.to_dict(), image_codebook_config=image_codebook_config.to_dict(), **kwargs)

[`FlavaConfig`] is the configuration class to store the configuration of a [`FlavaModel`]. It is used to
instantiate FLAVA model according to the specified arguments, defining the text model, image model, image codebook
and multimodal model configs. Instantiating a configuration with the defaults will yield a similar configuration to
that of the FLAVA [facebook/flava-full](https://huggingface.co/facebook/flava-full) architecture.

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

Args:
    text_config (`dict`, *optional*):
        Dictionary of configuration options used to initialize [`FlavaTextConfig`].
    image_config (`dict`, *optional*):
        Dictionary of configuration options used to initialize [`FlavaImageConfig`].
    multimodal_config (`dict`, *optional*):
        Dictionary of configuration options used to initialize [`FlavaMultimodalConfig`].
    hidden_size (`int`, *optional*, defaults to 768):
        Dimensionality of the encoder layers and the pooler layer.
    layer_norm_eps (`float`, *optional*, defaults to 1e-12):
        The epsilon used by the layer normalization layers.
    projection_dim (`int`, *optional*, defaults to 512):
        Dimensionality of text and image projection layers.
    logit_scale_init_value (`float`, *optional*, defaults to 2.6592):
        The initial value of the *logit_scale* parameter. Default is used as per the original FLAVA/CLIP
        implementation.
    initializer_range (`float`, *optional*, defaults to 0.02):
        The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
    ce_ignore_index (`int`, *optional*, defaults to -100):
        Cross entropy index to ignore.
    mim_weight (`float`, *optional*, defaults to 1.0):
        Weight to be assigned to MIM (Masked Image Modeling) unimodal loss
    mlm_weight (`float`, *optional*, defaults to 1.0):
        Weight to be assigned to MLM (Masked Language Modeling) unimodal loss
    global_contrastive_weight (`float`, *optional*, defaults to 1.0):
        Weight to be assigned to global contrastive cross-alignment loss.
    itm_weight (`float`, *optional*, defaults to 1.0):
        Weight to be assigned to image-text matching multimodal loss.
    mmm_image_weight (`float`, *optional*, defaults to 1.0):
        Weight to be assigned to MMM loss's image part.
    mmm_text_weight (`float`, *optional*, defaults to 1.0):
        Weight to be assigned to MMM loss's text part.
    global_backprop_contrastive (`bool`, *optional*, defaults to `True`):
        Whether to use global backpropgation through all workers in contrastive loss.
    skip_unmasked_multimodal_encoder (`bool`, *optional*, defaults to `True`):
        Whether to skip running unmasked multimodal encoder whose outputs are not used by FLAVA losses.
    return_loss (`bool`, *optional*, defaults to `True`):
        Whether to return loss or not

    kwargs (*optional*):
        Dictionary of keyword arguments.

Example:

```python
>>> from transformers import FlavaConfig, FlavaModel, FlavaForPreTraining

>>> # Initializing a FlavaConfig with style configuration
>>> configuration = FlavaConfig()

>>> # Initializing a FlavaModel and FlavaForPreTraining model (with random weights) from the style configuration
>>> model = FlavaModel(configuration)
>>> model_pre = FlavaForPreTraining(configuration)

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


def assign(self, variable, value):
    variable.assign(value)

Assign a value to a variable.

This should be used in optimizers instead of `variable.assign(value)` to
support backend specific optimizations.
Note that the variable can be a model variable or an optimizer variable;
it can be a backend native variable or a Keras variable.

Args:
    variable: The variable to update.
    value: The value to add to the variable.


class FlaxMaskedLMOutput(ModelOutput):
    logits: Optional[jnp.ndarray] = None
    hidden_states: Optional[Tuple[jnp.ndarray]] = None
    attentions: Optional[Tuple[jnp.ndarray]] = None

Base class for masked language models outputs.

Args:
    logits (`jnp.ndarray` of shape `(batch_size, sequence_length, config.vocab_size)`):
        Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
    hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape
        `(batch_size, sequence_length, hidden_size)`.

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

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


class Glm4Config(PretrainedConfig):
    model_type = 'glm4'
    keys_to_ignore_at_inference = ['past_key_values']
    base_model_tp_plan = {'layers.*.self_attn.q_proj': 'colwise', 'layers.*.self_attn.k_proj': 'colwise', 'layers.*.self_attn.v_proj': 'colwise', 'layers.*.self_attn.o_proj': 'rowwise', 'layers.*.mlp.gate_up_proj': 'colwise_rep', 'layers.*.mlp.down_proj': 'rowwise_rep'}
    base_model_pp_plan = {'embed_tokens': (['input_ids'], ['inputs_embeds']), 'layers': (['hidden_states', 'attention_mask'], ['hidden_states']), 'norm': (['hidden_states'], ['hidden_states'])}

    def __init__(self, vocab_size=151552, hidden_size=4096, intermediate_size=13696, num_hidden_layers=40, num_attention_heads=32, num_key_value_heads=2, partial_rotary_factor=0.5, head_dim=128, hidden_act='silu', attention_dropout=0.0, max_position_embeddings=131072, initializer_range=0.02, rms_norm_eps=1.5625e-07, use_cache=True, tie_word_embeddings=False, rope_theta=10000.0, pad_token_id=151329, eos_token_id=[151329, 151336, 151338], bos_token_id=None, attention_bias=True, **kwargs):
        self.vocab_size = vocab_size
        self.max_position_embeddings = max_position_embeddings
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.partial_rotary_factor = partial_rotary_factor
        self.head_dim = head_dim
        self.num_key_value_heads = num_key_value_heads
        self.hidden_act = hidden_act
        self.initializer_range = initializer_range
        self.rms_norm_eps = rms_norm_eps
        self.use_cache = use_cache
        self.rope_theta = rope_theta
        self.attention_bias = attention_bias
        self.attention_dropout = attention_dropout
        super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs)

This is the configuration class to store the configuration of a [`Glm4Model`]. It is used to instantiate an Glm4
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 Glm4-4-9b-chat.
e.g. [THUDM/GLM-4-9B-0414](https://huggingface.co/THUDM/GLM-4-9B-0414)
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
documentation from [`PretrainedConfig`] for more information.
Args:
    vocab_size (`int`, *optional*, defaults to 151552):
        Vocabulary size of the Glm4 model. Defines the number of different tokens that can be represented by the
        `inputs_ids` passed when calling [`Glm4Model`]
    hidden_size (`int`, *optional*, defaults to 4096):
        Dimension of the hidden representations.
    intermediate_size (`int`, *optional*, defaults to 13696):
        Dimension of the MLP representations.
    num_hidden_layers (`int`, *optional*, defaults to 40):
        Number of hidden layers in the Transformer decoder.
    num_attention_heads (`int`, *optional*, defaults to 32):
        Number of attention heads for each attention layer in the Transformer decoder.
    num_key_value_heads (`int`, *optional*, defaults to 2):
        This is the number of key_value heads that should be used to implement Grouped Query Attention. If
        `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
        `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
        converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
        by meanpooling all the original heads within that group. For more details, check out [this
        paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to
        `num_attention_heads`.
    partial_rotary_factor (`float`, *optional*, defaults to 0.5): The factor of the partial rotary position.
    head_dim (`int`, *optional*, defaults to 128):
        The attention head dimension.
    hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
        The legacy activation function. It is overwritten by the `hidden_activation`.
    attention_dropout (`float`, *optional*, defaults to 0.0):
        The dropout ratio for the attention probabilities.
    max_position_embeddings (`int`, *optional*, defaults to 131072):
        The maximum sequence length that this model might ever be used with.
    initializer_range (`float`, *optional*, defaults to 0.02):
        The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
    rms_norm_eps (`float`, *optional*, defaults to 1.5625e-07):
        The epsilon used by the rms normalization layers.
    use_cache (`bool`, *optional*, defaults to `True`):
        Whether or not the model should return the last key/values attentions (not used by all models). Only
        relevant if `config.is_decoder=True`.
    tie_word_embeddings (`bool`, *optional*, defaults to `False`):
        Whether to tie weight embeddings
    rope_theta (`float`, *optional*, defaults to 10000.0):
        The base period of the RoPE embeddings.
    pad_token_id (`int`, *optional*, defaults to 151329):
        Padding token id.
    eos_token_id (`int` | `list`, *optional*, defaults to `[151329, 151336, 151338]`):
        End of stream token id.
    bos_token_id (`int`, *optional*):
        Beginning of stream token id.
    attention_bias (`bool`, defaults to `False`, *optional*, defaults to `True`):
        Whether to use a bias in the query, key, value and output projection layers during self-attention.
```python
>>> from transformers import Glm4Model, Glm4Config
>>> # Initializing a Glm4 glm4-4-9b-chat style configuration
>>> configuration = Glm4Config()
>>> # Initializing a model from the glm4-4-9b-chat style configuration
>>> model = Glm4Model(configuration)
>>> # Accessing the model configuration
>>> configuration = model.config
```


def predict_on_batch(self, x):
    raise NotImplementedError

Returns predictions for a single batch of samples.

Args:
    x: Input data. It must be array-like.

Returns:
    NumPy array(s) of predictions.


def ensure_graph_is_valid(graph_def: graph_pb2.GraphDef) -> None:
    node_map = {}
    for node in graph_def.node:
        if node.name not in node_map:
            node_map[node.name] = node
        else:
            raise ValueError('Duplicate node names detected for ', node.name)
    for node in graph_def.node:
        for input_name in node.input:
            input_node_name = node_name_from_input(input_name)
            if input_node_name not in node_map:
                raise ValueError('Input for ', node.name, ' not found: ', input_name)

Makes sure that the graph is internally consistent.

Checks basic properties of the graph def and raises an exception if there are
input references to missing nodes, duplicated names, or other logic errors.

Args:
  graph_def: Definition of a graph to be checked.

Raises:
  ValueError: If the graph is incorrectly constructed.


def _check_trt_version_compatibility():
    if not _pywrap_py_utils.is_tensorrt_enabled():
        logging.error('Tensorflow needs to be built with TensorRT support enabled to allow TF-TRT to operate.')
        raise RuntimeError('Tensorflow has not been built with TensorRT support.')
    if platform.system() == 'Windows':
        logging.warn('Windows support is provided experimentally. No guarantee is made regarding functionality or engineering support. Use at your own risk.')
    linked_version = _pywrap_py_utils.get_linked_tensorrt_version()
    loaded_version = _pywrap_py_utils.get_loaded_tensorrt_version()
    logging.info('Linked TensorRT version: %s', str(linked_version))
    logging.info('Loaded TensorRT version: %s', str(loaded_version))

    def raise_trt_version_deprecated(version_type, trt_version):
        assert version_type in ['linked', 'loaded'], "Incorrect value received for version_type: %s. Accepted: ['linked', 'loaded']" % version_type
        logging.error('The {version_type} version of TensorRT: `{trt_version}` has now been removed. Please upgrade to TensorRT 7 or more recent.'.format(version_type=version_type, trt_version=trt_utils.version_tuple_to_string(trt_version)))
        raise RuntimeError('Incompatible %s TensorRT versions' % version_type)
    if not trt_utils.is_linked_tensorrt_version_greater_equal(7, 0, 0):
        raise_trt_version_deprecated('linked', linked_version)
    if not trt_utils.is_loaded_tensorrt_version_greater_equal(7, 0, 0):
        raise_trt_version_deprecated('loaded', loaded_version)
    if loaded_version[0] != linked_version[0] or not trt_utils.is_loaded_tensorrt_version_greater_equal(*linked_version):
        logging.error('Loaded TensorRT %s but linked TensorFlow against TensorRT %s. A few requirements must be met:\n\t-It is required to use the same major version of TensorRT during compilation and runtime.\n\t-TensorRT does not support forward compatibility. The loaded version has to be equal or more recent than the linked version.', trt_utils.version_tuple_to_string(loaded_version), trt_utils.version_tuple_to_string(linked_version))
        raise RuntimeError('Incompatible TensorRT major version')
    elif loaded_version != linked_version:
        logging.info('Loaded TensorRT %s and linked TensorFlow against TensorRT %s. This is supported because TensorRT minor/patch upgrades are backward compatible.', trt_utils.version_tuple_to_string(loaded_version), trt_utils.version_tuple_to_string(linked_version))

Check compatibility of TensorRT version.

Raises:
  RuntimeError: if the TensorRT library version is incompatible.


def moments(x, axes, shift=None, name=None, keep_dims=None, keepdims=None):
    keep_dims = deprecated_argument_lookup('keepdims', keepdims, 'keep_dims', keep_dims)
    if keep_dims is None:
        keep_dims = False
    with ops.name_scope(name, 'moments', [x, axes]):
        y = math_ops.cast(x, dtypes.float32) if x.dtype == dtypes.float16 else x
        mean = math_ops.reduce_mean(y, axes, keepdims=True, name='mean')
        variance = math_ops.reduce_mean(math_ops.squared_difference(y, array_ops.stop_gradient(mean)), axes, keepdims=True, name='variance')
        if not keep_dims:
            mean = array_ops.squeeze(mean, axes)
            variance = array_ops.squeeze(variance, axes)
        if x.dtype == dtypes.float16:
            return (math_ops.cast(mean, dtypes.float16), math_ops.cast(variance, dtypes.float16))
        else:
            return (mean, variance)

Calculate the mean and variance of `x`.

The mean and variance are calculated by aggregating the contents of `x`
across `axes`.  If `x` is 1-D and `axes = [0]` this is just the mean
and variance of a vector.

Note: shift is currently not used; the true mean is computed and used.

When using these moments for batch normalization (see
`tf.nn.batch_normalization`):

 * for so-called "global normalization", used with convolutional filters with
   shape `[batch, height, width, depth]`, pass `axes=[0, 1, 2]`.
 * for simple batch normalization pass `axes=[0]` (batch only).

Args:
  x: A `Tensor`.
  axes: Array of ints.  Axes along which to compute mean and
    variance.
  shift: Not used in the current implementation
  name: Name used to scope the operations that compute the moments.
  keep_dims: produce moments with the same dimensionality as the input.
  keepdims: Alias to keep_dims.

Returns:
  Two `Tensor` objects: `mean` and `variance`.


def get_modified_python_files(diff_with_last_commit: bool=False) -> List[str]:
    repo = Repo(PATH_TO_REPO)
    if not diff_with_last_commit:
        print(f'main is at {repo.refs.main.commit}')
        print(f'Current head is at {repo.head.commit}')
        branching_commits = repo.merge_base(repo.refs.main, repo.head)
        for commit in branching_commits:
            print(f'Branching commit: {commit}')
        return get_diff(repo, repo.head.commit, branching_commits)
    else:
        print(f'main is at {repo.head.commit}')
        parent_commits = repo.head.commit.parents
        for commit in parent_commits:
            print(f'Parent commit: {commit}')
        return get_diff(repo, repo.head.commit, parent_commits)

Return a list of python files that have been modified between:

- the current head and the main branch if `diff_with_last_commit=False` (default)
- the current head and its parent commit otherwise.

Returns:
    `List[str]`: The list of Python files with a diff (files added, renamed or deleted are always returned, files
    modified are returned if the diff in the file is not only in docstrings or comments, see
    `diff_is_docstring_only`).


def get_config(self):
    config = {'name': self._name}
    if self.clipnorm is not None:
        config['clipnorm'] = self.clipnorm
    if self.clipvalue is not None:
        config['clipvalue'] = self.clipvalue
    if self.global_clipnorm is not None:
        config['global_clipnorm'] = self.global_clipnorm
    return config

Returns the config of the optimizer.

An optimizer config is a Python dictionary (serializable)
containing the configuration of an optimizer.
The same optimizer can be reinstantiated later
(without any saved state) from this configuration.

Returns:
    Python dictionary.


def SetModifyTimestamp(self, value):
    if value is None or isinstance(value, int):
        self._last_modification_timestamp = value
    else:
        raise TypeError('timestamp can only be int or None, not %r' % value)

Set the last modify timestamp of this map.

Args:
  value: An integer containing the number of seconds since epoch, or None.

Raises:
  TypeError: The argument is not an int or None.


def delete(self, request):
    try:
        self.client.delete_object(Bucket=request.bucket, Key=request.object)
    except Exception as e:
        raise messages.S3ClientError(str(e), get_http_error_code(e))

Deletes given object from bucket
Args:
    request: (DeleteRequest) input message
  Returns:
    (void) Void, otherwise will raise if an error occurs


def join(self, timeout=_DEFAULT_TIMEOUT_SEC):
    if timeout and (not isinstance(timeout, int)):
        raise ValueError('`timeout` must be an integer or `None`.')
    with self._process_lock:
        if self._joined:
            raise ValueError("MultiProcessRunner can't be joined twice.")
        self._joined = True
    self._watchdog_thread.join(timeout)
    if self._watchdog_thread.is_alive():
        with self._process_lock:
            self._auto_restart = False
        logging.error('Timeout when joining for child processes. Terminating...')
        self.terminate_all(sig=signal.SIGTERM)
        self._watchdog_thread.join(_FORCE_KILL_WAIT_SEC)
        if self._watchdog_thread.is_alive():
            logging.error('Timeout when waiting for child processes to print stacktrace. Sending SIGKILL...')
            self.terminate_all()
            self._watchdog_thread.join()
        process_statuses = self._get_process_statuses()
        self._reraise_if_subprocess_error(process_statuses)
        raise SubprocessTimeoutError('One or more subprocesses timed out, where timeout was set to {}s. Please change the `timeout` argument for `MultiProcessRunner.join()` or `multi_process_runner.run()` if it should be adjusted.'.format(timeout), self._get_mpr_result(process_statuses))
    for (task_type, task_id), p in self._processes.items():
        logging.info('%s-%d exit code: %s', task_type, task_id, p.exitcode)
    process_statuses = self._get_process_statuses()
    self._reraise_if_subprocess_error(process_statuses)
    for (task_type, task_id), p in self._processes.items():
        assert p.exitcode is not None
        if p.exitcode > 0 and (task_type, task_id) not in self._terminated:
            raise UnexpectedSubprocessExitError('Subprocess %s-%d exited with exit code %s. See logs for details.' % (task_type, task_id, p.exitcode), self._get_mpr_result(process_statuses))
    logging.info('Joining log reading threads.')
    for thread in self._reading_threads:
        thread.join()
    logging.info('Joined log reading threads.')
    signal.alarm(0)
    return self._get_mpr_result(process_statuses)

Joins all the processes with timeout.

If any of the subprocesses does not exit approximately after `timeout`
seconds has passed after `join` call, this raises a
`SubprocessTimeoutError`.

Note: At timeout, it uses SIGTERM to terminate the subprocesses, in order to
log the stack traces of the subprocesses when they exit. However, this
results in timeout when the test runs with tsan (thread sanitizer); if tsan
is being run on the test targets that rely on timeout to assert information,
`MultiProcessRunner.terminate_all()` must be called after `join()`, before
the test exits, so the subprocesses are terminated with SIGKILL, and data
race is removed.

Args:
  timeout: optional integer or `None`. If provided as an integer, and not
  all processes report status within roughly `timeout` seconds, a
  `SubprocessTimeoutError` exception will be raised. If `None`, `join` never
  times out.

Returns:
  A `MultiProcessRunnerResult` object, which has two attributes,
  `return_value` and `stdout`. `return_value` always contains a list of
  return values from the subprocesses, although the order is not meaningful.
  If `return_output` argument is True at `__init__`, `stdout` is available
  that contains a list of all messages from subprocesses' stdout and stderr.

Raises:
  SubprocessTimeoutError: if not all processes report status approximately
    within `timeout` seconds. When this is raised, a
    `MultiProcessRunnerResult` object can be retrieved by
    `SubprocessTimeoutError`'s mpr_result attribute, which has the same
    structure as above 'Returns' section describes.
  UnexpectedSubprocessExitError: If any of the subprocesses did not exit
    properly (for example, they exit on SIGTERM or SIGKILL signal). When
    this is raised, a `MultiProcessRunnerResult` object can be retrieved by
    `UnexpectedSubprocessExitError`'s mpr_result attribute, which has the
    same structure as above 'Returns' section describes. If `max_run_time`
    is not `None`, it is expected that some subprocesses may be
    force-killed when `max_run_time` is up, and this is raised in those
    cases.
  Exception: if there is an Exception propagated from any subprocess. When
    this is raised, a `MultiProcessRunnerResult` object can be retrieved by
    `UnexpectedSubprocessExitError`'s mpr_result attribute, which has the
    same structure as above 'Returns' section describes.


def _create_variable(self, *args, **kwargs):
    with ops.name_scope('random_generator'):
        kwargs['name'] = 'StateVar'
        v = variables.Variable(*args, **kwargs)
    if isinstance(v, sharded_variable.ShardedVariable):
        raise ValueError("tf.random.Generator state is sharded, which is not allowed. When creating a tf.distribute.experimental.ParameterServerStrategy, please make sure that the `variable_partitioner` argument won't shard a small variable of shape [2] or [3]. Ways to avoid sharding small variables include setting `variable_partitioner` to None or to tf.distribute.experimental.partitioners.MinSizePartitioner with a large enough `min_shard_bytes`.")
    return v

Creates a variable.

Args:
  *args: positional arguments passed along to `variables.Variable.
  **kwargs: keyword arguments passed along to `variables.Variable.

Returns:
  The created variable.


def split_compile_and_replicate(computation: Callable[..., Any], inputs: Optional[List[List[core_types.Tensor]]]=None, infeed_queue: Optional[tpu_feed.InfeedQueue]=None, device_assignment: Optional[device_assignment_lib.DeviceAssignment]=None, name: Optional[Text]=None, use_tpu: bool=True, maximum_shapes: Optional[Any]=None, padding_spec: Optional[PaddingSpec]=None, xla_options: Optional[XLAOptions]=None) -> List[List[core_types.Tensor]]:
    del name
    inputs = [[]] if inputs is None else inputs
    xla_options = xla_options or XLAOptions()
    metadata_kwargs = {}
    if device_assignment is not None:
        metadata_kwargs = {'topology': device_assignment.topology.serialized(), 'device_assignment': device_assignment.core_assignment.flatten().tolist()}
        metadata_kwargs['num_cores_per_replica'] = device_assignment.num_cores_per_replica
    metadata_kwargs['allow_soft_placement'] = config.get_soft_device_placement()
    if config.get_soft_device_placement():
        logging.info('Automatic outside compilation is enabled. Ops without XLA kernels will be automatically placed on CPU.')
    if not isinstance(inputs, list):
        raise TypeError(f'tpu.replicate() inputs must be a list of lists/tuples, received {type(inputs)}')
    if any((not isinstance(inp, (list, tuple)) for inp in inputs)):
        raise TypeError(f'tpu.replicate() inputs must be a list of lists/tuples, received types: {[type(inp) for inp in inputs]}')
    num_replicas = len(inputs)
    if num_replicas == 0:
        return []
    for i in range(1, num_replicas):
        nest.assert_same_structure(inputs[0], inputs[i])
    inputs = variable_utils.convert_variables_to_tensors(inputs)
    flat_inputs_with_nones = [nest.flatten(per_replica_input, expand_composites=True) for per_replica_input in inputs]
    is_composite = nest.flatten(nest.map_structure(lambda x: _flatten_and_filter_composite(x, False, True), inputs[0]))
    flat_inputs = []
    for inp in flat_inputs_with_nones:
        flat_inputs.append([constant_op.constant(0) if x is None else ops.convert_to_tensor(x) for x in inp])
    flat_input_types = [x.dtype for x in flat_inputs[0]]
    input_arity = len(inputs[0])
    flat_input_arity = len(flat_input_types)
    for i in range(num_replicas):
        if len(inputs[i]) != input_arity:
            raise ValueError('Replicas must have the same number of inputs. Replica 0 had {} inputs, replica {} had {} inputs.'.format(input_arity, i, len(inputs[i])))
        types = [x.dtype for x in flat_inputs[i]]
        if types != flat_input_types:
            raise ValueError('Replicas must have matching input types. Replica 0 had input types {}, replica {} had input types {}'.format(flat_input_types, i, types))
    arg_error = xla.check_function_argument_count(computation, input_arity, infeed_queue)
    if arg_error is not None:
        if infeed_queue is None:
            raise TypeError(f'Supplied computation cannot be called with the specified inputs. You specified {input_arity} inputs: {[i.name for i in inputs[0]]}, but the computation needs {arg_error}')
        else:
            raise TypeError(f'Supplied computation cannot be called with the specified inputs. You specified {input_arity} inputs: {[i.name for i in inputs[0]]} ', f'and {infeed_queue.number_of_tuple_elements} additional inputs from infeed, but the computation needs {arg_error}')
    dynamic_shape_inputs = False
    if maximum_shapes:
        if infeed_queue:
            raise ValueError('Dynamic input shapes are not supported with infeed queues')
        nest.assert_same_structure(inputs[0], maximum_shapes, check_types=False)
        flat_maximum_shapes = nest.flatten([_flatten_and_filter_composite(x, y) for x, y in zip(nest.flatten(inputs[0]), nest.flatten(maximum_shapes))])
        flat_maximum_shapes = [tensor_shape.TensorShape(s) if s is not None else None for s in flat_maximum_shapes]
        nest.assert_same_structure(flat_inputs[0], flat_maximum_shapes, check_types=False)
        unpadded_inputs = flat_inputs
        flat_inputs, padding_maps = _pad_all_input(unpadded_inputs, flat_maximum_shapes, padding_spec)
        if padding_maps:
            dynamic_shape_inputs = True
            logging.info('TPU has inputs with dynamic shapes: %s', inputs[0])
    metadata_kwargs['step_marker_location'] = getattr(computation, 'step_marker_location', 'STEP_MARK_AT_ENTRY')
    metadata_kwargs['use_spmd_for_xla_partitioning'] = xla_options.use_spmd_for_xla_partitioning
    graph = ops.get_default_graph()
    flat_replicated_inputs = []
    for i in range(0, len(flat_inputs[0])):
        replicas = [flat_inputs[replica][i] for replica in range(num_replicas)]
        flat_replicated_inputs.append(tpu_ops.tpu_replicated_input(replicas, name='input{}'.format(i)))
    if isinstance(graph, func_graph.FuncGraph):
        cluster_name = graph.unique_name('cluster_' + graph.name)
    else:
        cluster_name = graph.unique_name('cluster')
    pivot = control_flow_ops.no_op(name=cluster_name + '/pivot')
    pivot._set_attr(_PIVOT_FOR_CLUSTER, attr_value_pb2.AttrValue(s=compat.as_bytes(cluster_name)))
    context = tpu_replication.TPUReplicateContext(name=cluster_name, num_replicas=num_replicas, pivot=pivot)
    try:
        context.Enter()
        metadata = tpu_ops.tpu_replicate_metadata(num_replicas=num_replicas, use_tpu=use_tpu, **metadata_kwargs)
        with tpu_function.tpu_shard_context(num_replicas), ops.control_dependencies([metadata]):
            if dynamic_shape_inputs and xla_options.enable_xla_dynamic_padder:
                for padding_map in padding_maps:
                    input_shape = flat_replicated_inputs[padding_map.arg_index].shape
                    flat_replicated_inputs[padding_map.arg_index] = tf2xla.set_dynamic_dimension_size(flat_replicated_inputs[padding_map.arg_index], padding_map.shape_index, flat_replicated_inputs[padding_map.padding_arg_index])
                    flat_replicated_inputs[padding_map.arg_index].set_shape(input_shape)
            flat_replicated_inputs = [array_ops.identity(x, name='replicated_input_{}'.format(i)) for i, x in enumerate(flat_replicated_inputs)]
            for i, composite in zip(flat_replicated_inputs, is_composite):
                if not dynamic_shape_inputs or composite:
                    i.op._set_attr('_tpu_input_identity', attr_value_pb2.AttrValue(b=True))
            computation_inputs = [None if inp is None else replicated for replicated, inp in zip(flat_replicated_inputs, flat_inputs_with_nones[0])]
            computation_inputs = nest.pack_sequence_as(structure=inputs[0], flat_sequence=computation_inputs[:flat_input_arity], expand_composites=True)
            if infeed_queue is not None:
                infeed_queue.set_number_of_shards(num_replicas)
                for t in infeed_queue.generate_dequeue_op():
                    computation_inputs.append(t)
            vscope = variable_scope.get_variable_scope()
            saved_use_resource = vscope.use_resource
            saved_custom_getter = vscope.custom_getter

            def custom_getter(getter, name, *args, **kwargs):
                """Variables on TPU have a few restrictions."""
                partitioner = kwargs.get('partitioner', None)
                if partitioner is not None:
                    kwargs['partitioner'] = None
                    logging.warning('Partitioned variables are not supported on TPU. Got `partitioner` that is %s for variable %s. Setting `partitioner` to `None`.', partitioner, name)
                if saved_custom_getter is None:
                    return getter(name, *args, **kwargs)
                else:
                    return saved_custom_getter(getter, name, *args, **kwargs)
            vscope.set_use_resource(True)
            vscope.set_custom_getter(custom_getter)
            outputs = computation(*computation_inputs)
            vscope.set_use_resource(saved_use_resource)
            vscope.set_custom_getter(saved_custom_getter)
            outputs = variable_utils.convert_variables_to_tensors(outputs)
        need_spmd_partitioning = xla_options.use_spmd_for_xla_partitioning and device_assignment is not None and (device_assignment.num_cores_per_replica > 1)
        outputs_is_flat = xla.is_flat(outputs)
        if outputs_is_flat:
            output_tensors, control_deps, pack_template = _postprocess_flat_outputs(outputs, need_spmd_partitioning)
        else:
            output_tensors, control_deps, pack_template = _postprocess_non_flat_outputs(outputs, need_spmd_partitioning)
        if tensor_tracer.TensorTracer.is_enabled():
            if tf2.enabled():
                logging.warn('TF API ver >= 2.0 detected. Tensor Tracer v1 is not enabled.')
            else:
                tt = tensor_tracer.TensorTracer()
                output_tensors = tt.trace_tpu(ops.get_default_graph(), output_tensors, control_deps, num_replicas)
        context.ExitResult(output_tensors)
    finally:
        context.report_unsupported_operations()
        context.Exit()
        host_compute_core = context.HostComputeCore()
    if host_compute_core:
        attr_value = attr_value_pb2.AttrValue()
        attr_value.list.s.extend((compat.as_bytes(x) for x in host_compute_core))
        metadata._set_attr('host_compute_core', attr_value)
    with ops.control_dependencies([metadata]):
        if use_tpu:
            compile_status = tpu_ops.tpu_compilation_result()
            op = compile_status.op
            attr_value = attr_value_pb2.AttrValue(s=compat.as_bytes(cluster_name))
            op._set_attr(_TPU_COMPILATION_STATUS_ATTR, attr_value)
        else:
            compile_status = control_flow_ops.no_op(name='compilation_status')
    if not output_tensors:
        return [compile_status, [control_flow_ops.group(control_deps, name='shard_%d' % i) for i in range(num_replicas)]]
    replicated_outputs = [[] for i in range(num_replicas)]
    for i, t in enumerate(output_tensors):
        if t is None:
            for replica in range(num_replicas):
                replicated_outputs[replica].append(None)
            continue
        ys = tpu_ops.tpu_replicated_output(t, num_replicas, name='output{}'.format(i))
        with ops.control_dependencies(control_deps):
            for replica in range(num_replicas):
                replicated_outputs[replica].append(array_ops.identity(ys[replica], name='output_%d_shard_%d' % (i, replica)))
    replicated_outputs = [nest.pack_sequence_as(pack_template, replica_outs, expand_composites=True) for replica_outs in replicated_outputs]
    return [compile_status, replicated_outputs]

Builds graph operators that runs compilation and replicated computation.

This is a lower level interface than replicate that returns a separate compile
and execute output tensor. In the generated graph the compile op feeds into
the execute op and no additional compilation is incurred when running the
compile op before the execute op. The compile op returns additional
information about the compilation but does not return the compiled program.

Args:
  computation: A Python function that builds the computation to replicate.
  inputs: A list of lists of input tensors or `None` (equivalent to
    `[[]]`), indexed by `[replica_num][input_num]`. All replicas must
    have the same number of inputs. Each input can be a nested structure
    containing values that are convertible to tensors. Note that passing an
    N-dimension list of compatible values will result in a N-dimension list of
    scalar tensors rather than a single Rank-N tensors. If you need different
    behavior, convert part of inputs to tensors with `tf.convert_to_tensor`.
  infeed_queue: If not `None`, the `InfeedQueue` from which to append a tuple
    of arguments as inputs to computation.
  device_assignment: If not `None`, a `DeviceAssignment` describing the
    mapping between logical cores in the computation with physical cores in
    the TPU topology. Uses a default device assignment if `None`. The
    `DeviceAssignment` may be omitted if each replica of the computation uses
    only one core, and there is either only one replica, or the number of
    replicas is equal to the number of cores in the TPU system.
  name: (Deprecated) Does nothing.
  use_tpu: When false, the input `computation` is executed on the XLA CPU/GPU
    backends. Currently, only supports a default placement (computation is
    placed on GPU if one is available, and on CPU if not).
  maximum_shapes: A nested structure of tf.TensorShape representing the shape
    to which the respective component of each input element in each replica
    should be padded. Any unknown dimensions (e.g.
    tf.compat.v1.Dimension(None) in a tf.TensorShape or -1 in a tensor-like
    object) will be padded to the maximum size of that dimension over all
    replicas. The structure of `maximum_shapes` needs to be the same as
    `inputs[0]`.
  padding_spec: An enum specified by `tf.tpu.PaddingSpec`. This describes the
    padding policy when the `inputs` to `tf.tpu.replicate` is dynamic.
    One usage is to enable automatic bucketizing on the inputs by setting the
    value to `tpu.PaddingSpec.POWER_OF_TWO`, which can help to reduce the
    recompilation in the XLA side.
  xla_options: An instance of `tpu.XLAOptions` which indicates the options
    passed to XLA compiler. Use `None` for default options.

Returns:
  A list of lists with the first list corresponding to the compile op and the
  second a list of output tensors, indexed by `[replica_num][output_num]`.
Raises:
  ValueError: If all replicas do not have equal numbers of input tensors.
  ValueError: If the number of inputs per replica does not match
    the number of formal parameters to `computation`.
  ValueError: If the static `inputs` dimensions don't match with the values
    given in `maximum_shapes`.
  ValueError: If the structure of inputs per replica does not match
    the structure of `maximum_shapes`.


class RunThresholdCriterion(beam.PTransform[beam.PCollection[NestedKeyedOutputT], beam.PCollection[NestedKeyedOutputT]]):

    def __init__(self, threshold_criterion: ThresholdFn):
        self._threshold_fn = threshold_criterion

    def expand(self, input: beam.PCollection[NestedKeyedOutputT]) -> beam.PCollection[NestedKeyedOutputT]:
        if self._threshold_fn.is_stateful:
            return input | beam.ParDo(_StatefulThresholdDoFn(self._threshold_fn.to_spec()))
        else:
            return input | beam.ParDo(_StatelessThresholdDoFn(self._threshold_fn.to_spec()))

Applies a threshold criterion to anomaly detection results.

This PTransform applies a `ThresholdFn` to the anomaly scores in
`AnomalyResult` objects, updating the prediction labels. It handles both
stateful and stateless `ThresholdFn` implementations.

Args:
  threshold_criterion: The `ThresholdFn` to apply.


class GPTSanJapaneseLayerSparseFF(nn.Module):

    def __init__(self, config: GPTSanJapaneseConfig):
        super().__init__()
        self.mlp = GPTSanJapaneseSparseMLP(config)
        self.soft_bypass_mlp = nn.Linear(config.d_model, config.d_model, bias=False)
        self.norm = nn.LayerNorm(config.d_model, eps=config.layer_norm_epsilon)

    def forward(self, hidden_states, output_router_logits):
        """
        Args:
            hidden_states (`torch.Tensor`) :
                [num_groups, tokens_per_group, hidden_dim] inputs to send to experts.
            output_router_logits (`bool`) :
                output experts router output.
        Returns:
            torch.Tensor[num_groups, tokens_per_group, hidden_dim]

        """
        forwarded_states, router_tuple = self.mlp(hidden_states)
        forwarded_states += torch.tanh(self.soft_bypass_mlp(hidden_states))
        output = hidden_states + self.norm(forwarded_states)
        if output_router_logits and router_tuple is not None:
            return (output, router_tuple)
        else:
            return output

Switch Transformers Feed Forward layer module. This is a wrapper around the Mixture of Experts module.

Parameters:
    config : ([`GPTSanJapaneseConfig`]): Model configuration class with all the parameters of the model.
        Initializing with a config file does not load the weights associated with the model, only the
        configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.


def is_scalar_batch(self, name='is_scalar_batch'):
    with self._name_scope(name):
        return ops.convert_to_tensor(self._is_scalar_helper(self.batch_shape, self.batch_shape_tensor), name='is_scalar_batch')

Indicates that `batch_shape == []`.

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

Returns:
  is_scalar_batch: `bool` scalar `Tensor`.


def _pop(self, key, indices=None, name=None):
    if name is None:
        name = '%s_get' % self._name
    indices, dtypes = self._get_indices_and_dtypes(indices)
    with ops.colocate_with(self._coloc_op):
        result = self._pop_fn(key, shared_name=self._name, indices=indices, dtypes=dtypes, name=name, capacity=self._capacity, memory_limit=self._memory_limit)
    return (key, self._get_return_value(result, indices))

Remove and return the associated (key, value) is returned from the staging area.

If the key is not in the staging area, this method will block until
the associated (key, value) is inserted.
Args:
    key: Key associated with the required data
    indices: Partial list of tensors to retrieve (optional).
            A list of integer or string indices.
            String indices are only valid if the Staging Area
            has names associated with it.
    name: A name for the operation (optional)

Returns:
    The created op


def register_serializable(package='Custom', name=None, predicate=None):

    def decorator(arg):
        """Registers a class with the serialization framework."""
        nonlocal predicate
        if not tf_inspect.isclass(arg):
            raise TypeError('Registered serializable must be a class: {}'.format(arg))
        class_name = name if name is not None else arg.__name__
        if predicate is None:
            predicate = lambda x: isinstance(x, arg)
        _class_registry.register(package, class_name, predicate, arg)
        return arg
    return decorator

Decorator for registering a serializable class.

THIS METHOD IS STILL EXPERIMENTAL AND MAY CHANGE AT ANY TIME.

Registered classes will be saved with a name generated by combining the
`package` and `name` arguments. When loading a SavedModel, modules saved with
this registered name will be created using the `_deserialize_from_proto`
method.

By default, only direct instances of the registered class will be saved/
restored with the `serialize_from_proto`/`deserialize_from_proto` methods. To
extend the registration to subclasses, use the `predicate argument`:

```python
class A(tf.Module):
  pass

register_serializable(
    package="Example", predicate=lambda obj: isinstance(obj, A))(A)
```

Args:
  package: The package that this class belongs to.
  name: The name to serialize this class under in this package. If None, the
    class's name will be used.
  predicate: An optional function that takes a single Trackable argument, and
    determines whether that object should be serialized with this `package`
    and `name`. The default predicate checks whether the object's type exactly
    matches the registered class. Predicates are executed in the reverse order
    that they are added (later registrations are checked first).

Returns:
  A decorator that registers the decorated class with the passed names and
  predicate.


def kill_test_logger(logger):
    for h in list(logger.handlers):
        logger.removeHandler(h)
        if isinstance(h, logging.FileHandler):
            h.close()

Cleans up a test logger object by removing all of its handlers.

Args:
  logger: The logging object to clean up.


class MeanSquaredError(LossFunctionWrapper):

    def __init__(self, reduction='sum_over_batch_size', name='mean_squared_error', dtype=None):
        super().__init__(mean_squared_error, name=name, reduction=reduction, dtype=dtype)

    def get_config(self):
        return Loss.get_config(self)

Computes the mean of squares of errors between labels and predictions.

Formula:

```python
loss = mean(square(y_true - y_pred))
```

Args:
    reduction: Type of reduction to apply to the loss. In almost all cases
        this should be `"sum_over_batch_size"`. Supported options are
        `"sum"`, `"sum_over_batch_size"`, `"mean"`,
        `"mean_with_sample_weight"` or `None`. `"sum"` sums the loss,
        `"sum_over_batch_size"` and `"mean"` sum the loss and divide by the
        sample size, and `"mean_with_sample_weight"` sums the loss and
        divides by the sum of the sample weights. `"none"` and `None`
        perform no aggregation. Defaults to `"sum_over_batch_size"`.
    name: Optional name for the loss instance.
    dtype: The dtype of the loss's computations. Defaults to `None`, which
        means using `keras.backend.floatx()`. `keras.backend.floatx()` is a
        `"float32"` unless set to different value
        (via `keras.backend.set_floatx()`). If a `keras.DTypePolicy` is
        provided, then the `compute_dtype` will be utilized.


def build_estimator(tf_transform_output, config, hidden_units=None):
    transformed_feature_spec = tf_transform_output.transformed_feature_spec().copy()
    transformed_feature_spec.pop(taxi.transformed_name(taxi.LABEL_KEY))
    real_valued_columns = [tf.feature_column.numeric_column(key, shape=()) for key in taxi.transformed_names(taxi.DENSE_FLOAT_FEATURE_KEYS)]
    categorical_columns = [tf.feature_column.categorical_column_with_identity(key, num_buckets=taxi.VOCAB_SIZE + taxi.OOV_SIZE, default_value=0) for key in taxi.transformed_names(taxi.VOCAB_FEATURE_KEYS)]
    categorical_columns += [tf.feature_column.categorical_column_with_identity(key, num_buckets=taxi.FEATURE_BUCKET_COUNT, default_value=0) for key in taxi.transformed_names(taxi.BUCKET_FEATURE_KEYS)]
    categorical_columns += [tf.feature_column.categorical_column_with_identity(key, num_buckets=num_buckets, default_value=0) for key, num_buckets in zip(taxi.transformed_names(taxi.CATEGORICAL_FEATURE_KEYS), taxi.MAX_CATEGORICAL_FEATURE_VALUES)]
    return tf_estimator.DNNLinearCombinedClassifier(config=config, linear_feature_columns=categorical_columns, dnn_feature_columns=real_valued_columns, dnn_hidden_units=hidden_units or [100, 70, 50, 25])

Build an estimator for predicting the tipping behavior of taxi riders.

Args:
  tf_transform_output: A TFTransformOutput.
  config: tf.contrib.learn.RunConfig defining the runtime environment for the
    estimator (including model_dir).
  hidden_units: [int], the layer sizes of the DNN (input layer first)

Returns:
  Resulting DNNLinearCombinedClassifier.


def crank_nicolson_step():

    def step_fn(time, next_time, coord_grid, value_grid, boundary_conditions, second_order_coeff_fn, first_order_coeff_fn, zeroth_order_coeff_fn, inner_second_order_coeff_fn, inner_first_order_coeff_fn, num_steps_performed, dtype=None, name=None):
        """Performs the step."""
        del num_steps_performed
        name = name or 'crank_nicolson_step'
        return parabolic_equation_step(time, next_time, coord_grid, value_grid, boundary_conditions, second_order_coeff_fn, first_order_coeff_fn, zeroth_order_coeff_fn, inner_second_order_coeff_fn, inner_first_order_coeff_fn, time_marching_scheme=crank_nicolson_scheme, dtype=dtype, name=name)
    return step_fn

Creates a stepper function with Crank-Nicolson time marching scheme.

Crank-Nicolson time marching scheme is one of the the most widely used schemes
for 1D PDEs. Given a space-discretized equation

```
du/dt = A(t) u(t) + b(t)
```
(here `u` is a value vector, `A` and `b` are the matrix and the vector defined
by the PDE), it approximates the right-hand side as an average of values taken
before and after the time step:

```
(u(t2) - u(t1)) / (t2 - t1) = (A(t1) u(t1) + b(t1) + A(t2) u(t2) + b(t2)) / 2.
```

Crank-Nicolson has second order accuracy and is stable.

More details can be found in `weighted_implicit_explicit.py` describing the
weighted implicit-explicit scheme - Crank-Nicolson scheme is a special case
with `theta = 0.5`.

Returns:
  Callable to be used in finite-difference PDE solvers (see fd_solvers.py).


def should_stop(self):
    if self._check_stop():
        return True
    if self._sess:
        return self._wrapped_is_stoppable and self._sess.should_stop()
    return True

Return true if this session should not be used anymore.

Always return True if the session was closed.

Returns:
  True if the session should stop, False otherwise.


class EsmForProteinFoldingOutput(ModelOutput):
    frames: Optional[torch.FloatTensor] = None
    sidechain_frames: Optional[torch.FloatTensor] = None
    unnormalized_angles: Optional[torch.FloatTensor] = None
    angles: Optional[torch.FloatTensor] = None
    positions: Optional[torch.FloatTensor] = None
    states: Optional[torch.FloatTensor] = None
    s_s: Optional[torch.FloatTensor] = None
    s_z: Optional[torch.FloatTensor] = None
    distogram_logits: Optional[torch.FloatTensor] = None
    lm_logits: Optional[torch.FloatTensor] = None
    aatype: Optional[torch.FloatTensor] = None
    atom14_atom_exists: Optional[torch.FloatTensor] = None
    residx_atom14_to_atom37: Optional[torch.FloatTensor] = None
    residx_atom37_to_atom14: Optional[torch.FloatTensor] = None
    atom37_atom_exists: Optional[torch.FloatTensor] = None
    residue_index: Optional[torch.FloatTensor] = None
    lddt_head: Optional[torch.FloatTensor] = None
    plddt: Optional[torch.FloatTensor] = None
    ptm_logits: Optional[torch.FloatTensor] = None
    ptm: Optional[torch.FloatTensor] = None
    aligned_confidence_probs: Optional[torch.FloatTensor] = None
    predicted_aligned_error: Optional[torch.FloatTensor] = None
    max_predicted_aligned_error: Optional[torch.FloatTensor] = None

Output type of [`EsmForProteinFoldingOutput`].

Args:
    frames (`torch.FloatTensor`):
        Output frames.
    sidechain_frames (`torch.FloatTensor`):
        Output sidechain frames.
    unnormalized_angles (`torch.FloatTensor`):
        Predicted unnormalized backbone and side chain torsion angles.
    angles (`torch.FloatTensor`):
        Predicted backbone and side chain torsion angles.
    positions (`torch.FloatTensor`):
        Predicted positions of the backbone and side chain atoms.
    states (`torch.FloatTensor`):
        Hidden states from the protein folding trunk.
    s_s (`torch.FloatTensor`):
        Per-residue embeddings derived by concatenating the hidden states of each layer of the ESM-2 LM stem.
    s_z (`torch.FloatTensor`):
        Pairwise residue embeddings.
    distogram_logits (`torch.FloatTensor`):
        Input logits to the distogram used to compute residue distances.
    lm_logits (`torch.FloatTensor`):
        Logits output by the ESM-2 protein language model stem.
    aatype (`torch.FloatTensor`):
        Input amino acids (AlphaFold2 indices).
    atom14_atom_exists (`torch.FloatTensor`):
        Whether each atom exists in the atom14 representation.
    residx_atom14_to_atom37 (`torch.FloatTensor`):
        Mapping between atoms in the atom14 and atom37 representations.
    residx_atom37_to_atom14 (`torch.FloatTensor`):
        Mapping between atoms in the atom37 and atom14 representations.
    atom37_atom_exists (`torch.FloatTensor`):
        Whether each atom exists in the atom37 representation.
    residue_index (`torch.FloatTensor`):
        The index of each residue in the protein chain. Unless internal padding tokens are used, this will just be
        a sequence of integers from 0 to `sequence_length`.
    lddt_head (`torch.FloatTensor`):
        Raw outputs from the lddt head used to compute plddt.
    plddt (`torch.FloatTensor`):
        Per-residue confidence scores. Regions of low confidence may indicate areas where the model's prediction is
        uncertain, or where the protein structure is disordered.
    ptm_logits (`torch.FloatTensor`):
        Raw logits used for computing ptm.
    ptm (`torch.FloatTensor`):
        TM-score output representing the model's high-level confidence in the overall structure.
    aligned_confidence_probs (`torch.FloatTensor`):
        Per-residue confidence scores for the aligned structure.
    predicted_aligned_error (`torch.FloatTensor`):
        Predicted error between the model's prediction and the ground truth.
    max_predicted_aligned_error (`torch.FloatTensor`):
        Per-sample maximum predicted error.


def fulltypes_for_flat_tensors(element_spec):
    specs = _specs_for_flat_tensors(element_spec)
    full_types_lists = [_translate_to_fulltype_for_flat_tensors(s) for s in specs]
    rval = nest.flatten(full_types_lists)
    return rval

Convert the element_spec for a dataset to a list of FullType Def.

Note that "flat" in this function and in `_flat_tensor_specs` is a nickname
for the "batchable tensor list" encoding used by datasets and map_fn.
The FullTypeDef created corresponds to this encoding (e.g. that uses variants
and not the FullTypeDef corresponding to the default "component" encoding).

This is intended for temporary internal use and expected to be removed
when type inference support is sufficient. See limitations of
`_translate_to_fulltype_for_flat_tensors`.

Args:
  element_spec: A nest of TypeSpec describing the elements of a dataset (or
    map_fn).

Returns:
  A list of FullTypeDef corresponding to ELEMENT_SPEC. The items
  in this list correspond to the items in `_flat_tensor_specs`.


def NeedsSeparatingHyphenHyphen(self, flag='help'):
    element = self.GetLastHealthyElement()
    component = element.component
    spec = inspectutils.GetFullArgSpec(component)
    return spec.varkw is not None or flag in spec.args or flag in spec.kwonlyargs

Returns whether a the trace need '--' before '--help'.

'--' is needed when the component takes keyword arguments, when the value of
flag matches one of the argument of the component, or the component takes in
keyword-only arguments(e.g. argument with default value).

Args:
  flag: the flag available for the trace

Returns:
  True for needed '--', False otherwise.


def matrix_transpose(a, name='matrix_transpose', conjugate=False):
    with ops.name_scope(name, values=[a]):
        a = ops.convert_to_tensor(a, name='a')
        a_shape = a.get_shape()
        ndims = a_shape.ndims
        if ndims is not None:
            if ndims < 2:
                raise ValueError(f'Argument `a` should be a (batch) matrix with rank >= 2.  Received `a` = {a} with shape: {a_shape}')
            perm = list(range(ndims - 2)) + [ndims - 1] + [ndims - 2]
        else:
            a_rank = rank(a)
            perm = concat((gen_math_ops._range(0, a_rank - 2, 1), [a_rank - 1, a_rank - 2]), 0)
        return transpose(a, perm=perm, conjugate=conjugate)

Transposes last two dimensions of tensor `a`.

For example:

```python
x = tf.constant([[1, 2, 3], [4, 5, 6]])
tf.linalg.matrix_transpose(x)  # [[1, 4],
                               #  [2, 5],
                               #  [3, 6]]

x = tf.constant([[1 + 1j, 2 + 2j, 3 + 3j],
                 [4 + 4j, 5 + 5j, 6 + 6j]])
tf.linalg.matrix_transpose(x, conjugate=True)  # [[1 - 1j, 4 - 4j],
                                               #  [2 - 2j, 5 - 5j],
                                               #  [3 - 3j, 6 - 6j]]

# Matrix with two batch dimensions.
# x.shape is [1, 2, 3, 4]
# tf.linalg.matrix_transpose(x) is shape [1, 2, 4, 3]
```

Note that `tf.matmul` provides kwargs allowing for transpose of arguments.
This is done with minimal cost, and is preferable to using this function. E.g.

```python
# Good!  Transpose is taken at minimal additional cost.
tf.matmul(matrix, b, transpose_b=True)

# Inefficient!
tf.matmul(matrix, tf.linalg.matrix_transpose(b))
```

@compatibility(numpy)
In `numpy` transposes are memory-efficient constant time operations as they
simply return a new view of the same data with adjusted `strides`.

TensorFlow does not support strides, `linalg.matrix_transpose` returns a new
tensor with the items permuted.
@end_compatibility

Args:
  a: A `Tensor` with `rank >= 2`.
  name: A name for the operation (optional).
  conjugate: Optional bool. Setting it to `True` is mathematically equivalent
    to tf.math.conj(tf.linalg.matrix_transpose(input)).

Returns:
  A transposed batch matrix `Tensor`.

Raises:
  ValueError:  If `a` is determined statically to have `rank < 2`.


def build_graph(device, dtype, data_format, input_shape, filter_shape, strides, padding, num_iters, warmup_iters):
    with ops.device('/%s:0' % device):
        inp = variable_v1.VariableV1(random_ops.truncated_normal(input_shape, dtype=dtype))
        filt = variable_v1.VariableV1(random_ops.truncated_normal(filter_shape, dtype=dtype))
        outputs = []
        conv2d_op = nn_ops.conv2d(inp, filt, strides, padding, data_format=data_format)
        outputs.append(conv2d_op)
        for _ in range(1, num_iters):
            with ops.control_dependencies([conv2d_op]):
                conv2d_op = nn_ops.conv2d(inp, filt, strides, padding, data_format=data_format)
                outputs.append(conv2d_op)
        warmup_groups = []
        warmup_conv2d_op = nn_ops.conv2d(inp, filt, strides, padding, data_format=data_format)
        warmup_groups.append(warmup_conv2d_op)
        for _ in range(1, warmup_iters):
            with ops.control_dependencies([warmup_conv2d_op]):
                warmup_conv2d_op = nn_ops.conv2d(inp, filt, strides, padding, data_format=data_format)
                warmup_groups.append(warmup_conv2d_op)
        return (control_flow_ops.group(*warmup_groups), control_flow_ops.group(*outputs))

builds a graph containing a sequence of conv2d operations.

Args:
  device: String, the device to run on.
  dtype: Data type for the convolution.
  data_format: A string from: "NHWC" or "NCHW". Data format for input and
               output data.
  input_shape: Shape of the input tensor.
  filter_shape: Shape of the filter tensor.
  strides: A list of ints. 1-D of length 4. The stride of sliding
           window for each dimension of input.
  padding: A string from: "SAME", "VALID". The type of padding
           algorithm to use.
  num_iters: number of iterations to run conv2d.
  warmup_iters: number of iterations for warmup runs.

Returns:
  An array of tensors to run()


def read(self, range_tracker):
    raise NotImplementedError

Returns an iterator that reads data from the source.

The returned set of data must respect the boundaries defined by the given
``RangeTracker`` object. For example:

  * Returned set of data must be for the range
    ``[range_tracker.start_position, range_tracker.stop_position)``. Note
    that a source may decide to return records that start after
    ``range_tracker.stop_position``. See documentation in class
    ``RangeTracker`` for more details. Also, note that framework might
    invoke ``range_tracker.try_split()`` to perform dynamic split
    operations. range_tracker.stop_position may be updated
    dynamically due to successful dynamic split operations.
  * Method ``range_tracker.try_split()`` must be invoked for every record
    that starts at a split point.
  * Method ``range_tracker.record_current_position()`` may be invoked for
    records that do not start at split points.

Args:
  range_tracker: a ``RangeTracker`` whose boundaries must be respected
                 when reading data from the source. A runner that reads this
                 source muss pass a ``RangeTracker`` object that is not
                 ``None``.
Returns:
  an iterator of data read by the source.


def serialize(optimizer):
    return serialize_keras_object(optimizer)

Serialize the optimizer configuration to JSON compatible python dict.

The configuration can be used for persistence and reconstruct the `Optimizer`
instance again.

>>> tf.keras.optimizers.serialize(tf.keras.optimizers.SGD())
{'class_name': 'SGD', 'config': {'name': 'SGD', 'learning_rate': 0.01,
                                 'decay': 0.0, 'momentum': 0.0,
                                 'nesterov': False}}

Args:
  optimizer: An `Optimizer` instance to serialize.

Returns:
  Python dict which contains the configuration of the input optimizer.


def __init__(self, name: Text, num_replicas: int, pivot: ops.Operation):
    super(TPUReplicateContext, self).__init__()
    self._num_replicas = num_replicas
    self._outer_device_function_stack = None
    self._oc_dev_fn_stack = None
    self._outside_compilation_cluster = None
    self._is_map_outside_compilation = False
    self._outside_compilation_v2_context = None
    self._outside_compilation_counter = 0
    self._in_gradient_colocation = None
    self._gradient_colocation_stack = []
    self._host_compute_core = []
    self._name = name
    self._tpu_replicate_attr = attr_value_pb2.AttrValue(s=compat.as_bytes(self._name))
    self._unsupported_ops = []
    self._pivot = pivot
    self._replicated_vars = {}

Builds a new TPUReplicateContext.

Args:
  name: a unique name for the context, used to populate the `_tpu_replicate`
    attribute.
  num_replicas: an integer that gives the number of replicas for the
    computation.
  pivot: a pivot node. Nodes in the TPUReplicateContext that do not have any
    inputs will have a control dependency on the pivot node. This ensures
    that nodes are correctly included in any enclosing control flow
    contexts.


class VitPoseConfig(PretrainedConfig):
    model_type = 'vitpose'

    def __init__(self, backbone_config: Optional[PretrainedConfig]=None, backbone: Optional[str]=None, use_pretrained_backbone: bool=False, use_timm_backbone: bool=False, backbone_kwargs: Optional[dict]=None, initializer_range: float=0.02, scale_factor: int=4, use_simple_decoder: bool=True, **kwargs):
        super().__init__(**kwargs)
        if use_pretrained_backbone:
            logger.info('`use_pretrained_backbone` is `True`. For the pure inference purpose of VitPose weight do not set this value.')
        if use_timm_backbone:
            raise ValueError('use_timm_backbone set `True` is not supported at the moment.')
        if backbone_config is None and backbone is None:
            logger.info('`backbone_config` is `None`. Initializing the config with the default `VitPose` backbone.')
            backbone_config = CONFIG_MAPPING['vitpose_backbone'](out_indices=[4])
        elif isinstance(backbone_config, dict):
            backbone_model_type = backbone_config.get('model_type')
            config_class = CONFIG_MAPPING[backbone_model_type]
            backbone_config = config_class.from_dict(backbone_config)
        verify_backbone_config_arguments(use_timm_backbone=use_timm_backbone, use_pretrained_backbone=use_pretrained_backbone, backbone=backbone, backbone_config=backbone_config, backbone_kwargs=backbone_kwargs)
        self.backbone_config = backbone_config
        self.backbone = backbone
        self.use_pretrained_backbone = use_pretrained_backbone
        self.use_timm_backbone = use_timm_backbone
        self.backbone_kwargs = backbone_kwargs
        self.initializer_range = initializer_range
        self.scale_factor = scale_factor
        self.use_simple_decoder = use_simple_decoder

This is the configuration class to store the configuration of a [`VitPoseForPoseEstimation`]. It is used to instantiate a
VitPose 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 VitPose
[usyd-community/vitpose-base-simple](https://huggingface.co/usyd-community/vitpose-base-simple) architecture.

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

Args:
    backbone_config (`PretrainedConfig` or `dict`, *optional*, defaults to `VitPoseBackboneConfig()`):
        The configuration of the backbone model. Currently, only `backbone_config` with `vitpose_backbone` as `model_type` is supported.
    backbone (`str`, *optional*):
        Name of backbone to use when `backbone_config` is `None`. If `use_pretrained_backbone` is `True`, this
        will load the corresponding pretrained weights from the timm or transformers library. If `use_pretrained_backbone`
        is `False`, this loads the backbone's config and uses that to initialize the backbone with random weights.
    use_pretrained_backbone (`bool`, *optional*, defaults to `False`):
        Whether to use pretrained weights for the backbone.
    use_timm_backbone (`bool`, *optional*, defaults to `False`):
        Whether to load `backbone` from the timm library. If `False`, the backbone is loaded from the transformers
        library.
    backbone_kwargs (`dict`, *optional*):
        Keyword arguments to be passed to AutoBackbone when loading from a checkpoint
        e.g. `{'out_indices': (0, 1, 2, 3)}`. Cannot be specified if `backbone_config` is set.
    initializer_range (`float`, *optional*, defaults to 0.02):
        The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
    scale_factor (`int`, *optional*, defaults to 4):
        Factor to upscale the feature maps coming from the ViT backbone.
    use_simple_decoder (`bool`, *optional*, defaults to `True`):
        Whether to use a `VitPoseSimpleDecoder` to decode the feature maps from the backbone into heatmaps. Otherwise it uses `VitPoseClassicDecoder`.

Example:

```python
>>> from transformers import VitPoseConfig, VitPoseForPoseEstimation

>>> # Initializing a VitPose configuration
>>> configuration = VitPoseConfig()

>>> # Initializing a model (with random weights) from the configuration
>>> model = VitPoseForPoseEstimation(configuration)

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


def __init__(self, job_to_label_mapping=None, tf_server_port=8470, rpc_layer='grpc', override_client=None, executable_location=ExecutableLocation.WITHIN_CLUSTER):
    try:
        from kubernetes import config as k8sconfig
        if not override_client:
            if executable_location == ExecutableLocation.OFF_CLUSTER:
                k8sconfig.load_kube_config()
            elif executable_location == ExecutableLocation.WITHIN_CLUSTER:
                k8sconfig.load_incluster_config()
            else:
                raise ValueError('The executable location provided is invalid.')
    except ImportError:
        if not override_client:
            raise ImportError('The Kubernetes Python client must be installed before using the Kubernetes Cluster Resolver. To install the Kubernetes Python client, run `pip install kubernetes` on your command line.')
    if not job_to_label_mapping:
        job_to_label_mapping = {'worker': ['job-name=tensorflow']}
    self._job_to_label_mapping = job_to_label_mapping
    self._tf_server_port = tf_server_port
    self._override_client = override_client
    self.task_type = None
    self.task_id = None
    self.rpc_layer = rpc_layer

Initializes a new KubernetesClusterResolver.

This initializes a new Kubernetes ClusterResolver. The ClusterResolver
will attempt to talk to the Kubernetes master to retrieve all the instances
of pods matching a label selector.

Args:
  job_to_label_mapping: A mapping of TensorFlow jobs to label selectors.
    This allows users to specify many TensorFlow jobs in one Cluster
    Resolver, and each job can have pods belong with different label
    selectors. For example, a sample mapping might be
    ```
    {'worker': ['job-name=worker-cluster-a', 'job-name=worker-cluster-b'],
     'ps': ['job-name=ps-1', 'job-name=ps-2']}
    ```
  tf_server_port: The port the TensorFlow server is listening on.
  rpc_layer: (Optional) The RPC layer TensorFlow should use to communicate
    between tasks in Kubernetes. Defaults to 'grpc'.
  override_client: The Kubernetes client (usually automatically retrieved
    using `from kubernetes import client as k8sclient`). If you pass this
    in, you are responsible for setting Kubernetes credentials manually and
    calling `k8sconfig.load_kube_config()` or
    `k8sconfig.load_incluster_config()` before using this ClusterResolver.
  executable_location: Parameter that specifies whether or not this
    TensorFlow code is running from within a K8S cluster or not.

Raises:
  ImportError: If the Kubernetes Python client is not installed and no
    `override_client` is passed in.
  RuntimeError: If autoresolve_task is not a boolean or a callable.
  ValueError: If `executable_location` is not a valid value.


def split_if_relative_reference(reference: message.Message) -> None:
    _validate_reference(reference)
    uri_field = reference.DESCRIPTOR.fields_by_name.get('uri')
    if not proto_utils.field_is_set(reference, uri_field):
        return
    uri = proto_utils.get_value_at_field(reference, uri_field)
    internal_match = re.fullmatch(_INTERNAL_REFERENCE_PATTERN, uri.value)
    if internal_match is not None:
        reference_id_field = get_reference_id_field_for_resource(reference, internal_match.group('resource_type'))
        reference_id = proto_utils.create_message_from_descriptor(reference_id_field.message_type)
        populate_typed_reference_id(reference_id, internal_match.group('resource_id'), internal_match.group('version'))
        proto_utils.copy_common_field(uri, reference_id, 'id')
        proto_utils.copy_common_field(uri, reference_id, 'extension')
        proto_utils.set_value_at_field(reference, reference_id_field, reference_id)
        return
    fragment_match = re.fullmatch(_FRAGMENT_REFERENCE_PATTERN, uri.value)
    if fragment_match is not None:
        fragment_field = reference.DESCRIPTOR.fields_by_name['fragment']
        fragment = proto_utils.create_message_from_descriptor(fragment_field.message_type)
        value_field = fragment.DESCRIPTOR.fields_by_name['value']
        proto_utils.set_value_at_field(fragment, value_field, uri.value[1:])
        proto_utils.copy_common_field(uri, fragment, 'id')
        proto_utils.copy_common_field(uri, fragment, 'extension')
        proto_utils.set_value_at_field(reference, fragment_field, fragment)
        return

If possible, parses a `Reference` `uri` into more structured fields.

This is only possible for two forms of reference uris:
* Relative references of the form $TYPE/$ID, e.g., "Patient/1234"
  In this case, this will be parsed to a proto of the form:
  {patient_id: {value: "1234"}}
* Fragments of the form "#$FRAGMENT", e.g., "#vs1".  In this case, this would
  be parsed into a proto of the form:
  {fragment: {value: "vs1"} }

If the reference URI matches one of these schemas, the `uri` field will be
cleared, and the appropriate structured fields set. Otherwise, the reference
will be unchanged.

Args:
  reference: The FHIR reference to potentially split.

Raises:
  ValueError: If the message is not a valid FHIR Reference proto.


class AltCLIPProcessor(ProcessorMixin):
    attributes = ['image_processor', 'tokenizer']
    image_processor_class = ('CLIPImageProcessor', 'CLIPImageProcessorFast')
    tokenizer_class = ('XLMRobertaTokenizer', 'XLMRobertaTokenizerFast')

    @deprecate_kwarg(old_name='feature_extractor', version='5.0.0', new_name='image_processor')
    def __init__(self, image_processor=None, tokenizer=None):
        if image_processor is None:
            raise ValueError('You need to specify an `image_processor`.')
        if tokenizer is None:
            raise ValueError('You need to specify a `tokenizer`.')
        super().__init__(image_processor, tokenizer)

    def __call__(self, images: ImageInput=None, text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]]=None, audio=None, videos=None, **kwargs: Unpack[AltClipProcessorKwargs]) -> BatchEncoding:
        """
        Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
        and `kwargs` arguments to XLMRobertaTokenizerFast's [`~XLMRobertaTokenizerFast.__call__`] if `text` is not
        `None` to encode the text. To prepare the image(s), this method forwards the `images` and `kwrags` arguments to
        CLIPImageProcessor's [`~CLIPImageProcessor.__call__`] if `images` is not `None`. Please refer to the docstring
        of the above two methods for more information.

        Args:

            images (`ImageInput`):
                The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
                tensor. Both channels-first and channels-last formats are supported.
            text (`TextInput`, `PreTokenizedInput`, `List[TextInput]`, `List[PreTokenizedInput]`):
                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
            return_tensors (`str` or [`~utils.TensorType`], *optional*):
                If set, will return tensors of a particular framework. Acceptable values are:
                    - `'tf'`: Return TensorFlow `tf.constant` objects.
                    - `'pt'`: Return PyTorch `torch.Tensor` objects.
                    - `'np'`: Return NumPy `np.ndarray` objects.
                    - `'jax'`: Return JAX `jnp.ndarray` objects.
        Returns:
            [`BatchEncoding`]: A [`BatchEncoding`] with the following fields:

            - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
              `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
              `None`).
            - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
        """
        if text is None and images is None:
            raise ValueError('You must specify either text or images.')
        if text is None and images is None:
            raise ValueError('You must specify either text or images.')
        output_kwargs = self._merge_kwargs(AltClipProcessorKwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, **kwargs)
        if text is not None:
            encoding = self.tokenizer(text, **output_kwargs['text_kwargs'])
        if images is not None:
            image_features = self.image_processor(images, **output_kwargs['images_kwargs'])
        if 'return_tensors' in output_kwargs['common_kwargs']:
            return_tensors = output_kwargs['common_kwargs'].pop('return_tensors', None)
        if text is not None and images is not None:
            encoding['pixel_values'] = image_features.pixel_values
            return encoding
        elif text is not None:
            return encoding
        else:
            return BatchEncoding(data=dict(**image_features), tensor_type=return_tensors)

    def batch_decode(self, *args, **kwargs):
        """
        This method forwards all its arguments to XLMRobertaTokenizerFast's [`~PreTrainedTokenizer.batch_decode`].
        Please refer to the docstring of this method for more information.
        """
        return self.tokenizer.batch_decode(*args, **kwargs)

    def decode(self, *args, **kwargs):
        """
        This method forwards all its arguments to XLMRobertaTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please
        refer to the docstring of this method for more information.
        """
        return self.tokenizer.decode(*args, **kwargs)

    @property
    def model_input_names(self):
        tokenizer_input_names = self.tokenizer.model_input_names
        image_processor_input_names = self.image_processor.model_input_names
        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))

Constructs a AltCLIP processor which wraps a CLIP image processor and a XLM-Roberta tokenizer into a single
processor.

[`AltCLIPProcessor`] offers all the functionalities of [`CLIPImageProcessor`] and [`XLMRobertaTokenizerFast`]. See
the [`~AltCLIPProcessor.__call__`] and [`~AltCLIPProcessor.decode`] for more information.

Args:
    image_processor ([`CLIPImageProcessor`], *optional*):
        The image processor is a required input.
    tokenizer ([`XLMRobertaTokenizerFast`], *optional*):
        The tokenizer is a required input.


def _get_strides_and_dilation_rate(num_spatial_dims, strides, dilation_rate):
    if dilation_rate is None:
        dilation_rate = [1] * num_spatial_dims
    elif len(dilation_rate) != num_spatial_dims:
        raise ValueError(f'`len(dilation_rate)` should be {num_spatial_dims}. Received: dilation_rate={dilation_rate} of length {len(dilation_rate)}')
    dilation_rate = np.array(dilation_rate, dtype=np.int32)
    if np.any(dilation_rate < 1):
        raise ValueError(f'all values of `dilation_rate` must be positive. Received: dilation_rate={dilation_rate}')
    if strides is None:
        strides = [1] * num_spatial_dims
    elif len(strides) != num_spatial_dims:
        raise ValueError(f'`len(strides)` should be {num_spatial_dims}. Received: strides={strides} of length {len(strides)}')
    strides = np.array(strides, dtype=np.int32)
    if np.any(strides < 1):
        raise ValueError(f'all values of `strides` must be positive. Received: strides={strides}')
    if np.any(strides > 1) and np.any(dilation_rate > 1):
        raise ValueError(f'`strides > 1` not supported in conjunction with `dilation_rate > 1`. Received: strides={strides} and dilation_rate={dilation_rate}')
    return (strides, dilation_rate)

Helper function for verifying strides and dilation_rate arguments.

This is used by `convolution` and `pool`.

Args:
  num_spatial_dims: int
  strides: Optional.  List of N ints >= 1.  Defaults to `[1]*N`.  If any value
    of strides is > 1, then all values of dilation_rate must be 1.
  dilation_rate: Optional.  List of N ints >= 1.  Defaults to `[1]*N`.  If any
    value of dilation_rate is > 1, then all values of strides must be 1.

Returns:
  Normalized (strides, dilation_rate) as int32 numpy arrays of shape
  [num_spatial_dims].

Raises:
  ValueError: if the parameters are invalid.


def from_schema(cls, schema: class_schema.Schema, module_name: str, name: str, qualname: Optional[str]=None, is_method: bool=True) -> 'Signature':
    arg_names = list(schema.metadata.get('init_arg_list', []))
    if arg_names and arg_names[-1].startswith('*'):
        vararg_name = arg_names[-1][1:]
        arg_names.pop(-1)
    else:
        vararg_name = None

    def get_arg_spec(arg_name):
        field = schema.get_field(arg_name)
        if not field:
            raise ValueError(f'Argument {arg_name!r} is not a symbolic field.')
        return field.value
    args = []
    if is_method:
        args.append(Argument.from_annotation('self', Argument.Kind.POSITIONAL_OR_KEYWORD))
    args.extend([Argument(n, Argument.Kind.POSITIONAL_OR_KEYWORD, get_arg_spec(n)) for n in arg_names])
    varargs = None
    if vararg_name:
        varargs = Argument(vararg_name, Argument.Kind.VAR_POSITIONAL, get_arg_spec(vararg_name))
    existing_names = set(arg_names)
    if vararg_name:
        existing_names.add(vararg_name)
    kwonlyargs = []
    varkw = None
    for key, field in schema.fields.items():
        if key not in existing_names and (not field.frozen):
            if key.is_const:
                kwonlyargs.append(Argument(str(key), Argument.Kind.KEYWORD_ONLY, field.value))
            else:
                varkw = Argument(schema.metadata.get('varkw_name', None) or 'kwargs', Argument.Kind.VAR_KEYWORD, class_schema.ValueSpec.DictType(field.value))
    return Signature(callable_type=CallableType.FUNCTION, name=name, module_name=module_name, qualname=qualname, description=schema.description, args=args, kwonlyargs=kwonlyargs, varargs=varargs, varkw=varkw, return_value=schema.metadata.get('returns', None))

Creates a signature from a schema object.

Args:
  schema: A `pg.typing.Schema` object associated with a `pg.Object`.
  module_name: Module name for the signature.
  name: Function or method name of the signature.
  qualname: Qualname of the signature.
  is_method: If True, `self` will be added in the signature as the first
    argument.

Returns:
  A signature object from the schema.


def update_checkpoint_state_internal(save_dir, model_checkpoint_path, all_model_checkpoint_paths=None, latest_filename=None, save_relative_paths=False, all_model_checkpoint_timestamps=None, last_preserved_timestamp=None):
    coord_checkpoint_filename = _GetCheckpointFilename(save_dir, latest_filename)
    if save_relative_paths:
        if os.path.isabs(model_checkpoint_path):
            rel_model_checkpoint_path = os.path.relpath(model_checkpoint_path, save_dir)
        else:
            rel_model_checkpoint_path = model_checkpoint_path
        rel_all_model_checkpoint_paths = []
        for p in all_model_checkpoint_paths:
            if os.path.isabs(p):
                rel_all_model_checkpoint_paths.append(os.path.relpath(p, save_dir))
            else:
                rel_all_model_checkpoint_paths.append(p)
        ckpt = generate_checkpoint_state_proto(save_dir, rel_model_checkpoint_path, all_model_checkpoint_paths=rel_all_model_checkpoint_paths, all_model_checkpoint_timestamps=all_model_checkpoint_timestamps, last_preserved_timestamp=last_preserved_timestamp)
    else:
        ckpt = generate_checkpoint_state_proto(save_dir, model_checkpoint_path, all_model_checkpoint_paths=all_model_checkpoint_paths, all_model_checkpoint_timestamps=all_model_checkpoint_timestamps, last_preserved_timestamp=last_preserved_timestamp)
    if coord_checkpoint_filename == ckpt.model_checkpoint_path:
        raise RuntimeError("Save path '%s' conflicts with path used for checkpoint state.  Please use a different save path." % model_checkpoint_path)
    file_io.atomic_write_string_to_file(coord_checkpoint_filename, text_format.MessageToString(ckpt))

Updates the content of the 'checkpoint' file.

This updates the checkpoint file containing a CheckpointState
proto.

Args:
  save_dir: Directory where the model was saved.
  model_checkpoint_path: The checkpoint file.
  all_model_checkpoint_paths: List of strings.  Paths to all not-yet-deleted
    checkpoints, sorted from oldest to newest.  If this is a non-empty list,
    the last element must be equal to model_checkpoint_path.  These paths
    are also saved in the CheckpointState proto.
  latest_filename: Optional name of the checkpoint file.  Default to
    'checkpoint'.
  save_relative_paths: If `True`, will write relative paths to the checkpoint
    state file.
  all_model_checkpoint_timestamps: Optional list of timestamps (floats,
    seconds since the Epoch) indicating when the checkpoints in
    `all_model_checkpoint_paths` were created.
  last_preserved_timestamp: A float, indicating the number of seconds since
    the Epoch when the last preserved checkpoint was written, e.g. due to a
    `keep_checkpoint_every_n_hours` parameter (see
    `tf.train.CheckpointManager` for an implementation).

Raises:
  RuntimeError: If any of the model checkpoint paths conflict with the file
    containing CheckpointSate.


def __init__(self, components):
    global _next_device_number, _next_device_number_lock
    self.components = tuple((device_util.canonicalize(d) for d in components))
    if not self.components:
        raise ValueError('ParallelDevice requires at least one component.')
    ctx = context.context()
    with _next_device_number_lock:
        self._name = '{}/device:CUSTOM:{}'.format(ctx.host_address_space(), _next_device_number)
        _next_device_number += 1
    device, device_info = _pywrap_parallel_device.GetParallelDeviceCapsules(self._name, self.components)
    context.register_custom_device(device, self._name, device_info)
    self._device_ids = None
    self._device_scope = None
    _all_parallel_devices[self._name] = self

Creates a device which executes operations in parallel on `components`.

Args:
  components: A list of device names. Each operation executed on the
    returned device executes on these component devices.

Returns:
  A string with the name of the newly created device.


class LogCoshError(MeanMetricWrapper):

    def __init__(self, name='logcosh', dtype=None):
        super(LogCoshError, self).__init__(logcosh, name, dtype=dtype)

Computes the logarithm of the hyperbolic cosine of the prediction error.

`logcosh = log((exp(x) + exp(-x))/2)`, where x is the error (y_pred - y_true)

Args:
  name: (Optional) string name of the metric instance.
  dtype: (Optional) data type of the metric result.

Standalone usage:

>>> m = tf.keras.metrics.LogCoshError()
>>> m.update_state([[0, 1], [0, 0]], [[1, 1], [0, 0]])
>>> m.result().numpy()
0.10844523

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

Usage with `compile()` API:

```python
model.compile(optimizer='sgd',
              loss='mse',
              metrics=[tf.keras.metrics.LogCoshError()])
```


def forward(self, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.Tensor]=None, encoder_hidden_states: Optional[torch.FloatTensor]=None, encoder_attention_mask: Optional[torch.FloatTensor]=None, head_mask: Optional[torch.Tensor]=None, cross_attn_head_mask: Optional[torch.Tensor]=None, past_key_values: Optional[List[torch.FloatTensor]]=None, inputs_embeds: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, cache_position: Optional[torch.LongTensor]=None) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    outputs = self.model.decoder(input_ids=input_ids, attention_mask=attention_mask, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, head_mask=head_mask, cross_attn_head_mask=cross_attn_head_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, cache_position=cache_position)
    logits = self.lm_head(outputs[0])
    loss = None
    if labels is not None:
        labels = labels.to(logits.device)
        loss_fct = CrossEntropyLoss()
        loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
    if not return_dict:
        output = (logits,) + outputs[1:]
        return (loss,) + output if loss is not None else output
    return CausalLMOutputWithCrossAttentions(loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, cross_attentions=outputs.cross_attentions)

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

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

Example:

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

>>> tokenizer = AutoTokenizer.from_pretrained("facebook/blenderbot_small-90M")
>>> model = BlenderbotSmallForCausalLM.from_pretrained("facebook/blenderbot_small-90M", add_cross_attention=False)
>>> assert model.config.is_decoder, f"{model.__class__} has to be configured as a decoder."
>>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
>>> outputs = model(**inputs)

>>> logits = outputs.logits
>>> expected_shape = [1, inputs.input_ids.shape[-1], model.config.vocab_size]
>>> list(logits.shape) == expected_shape
True
```