def get_image_processor_dict(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> tuple[dict[str, Any], dict[str, Any]]:
    cache_dir = kwargs.pop('cache_dir', None)
    force_download = kwargs.pop('force_download', False)
    resume_download = kwargs.pop('resume_download', None)
    proxies = kwargs.pop('proxies', None)
    token = kwargs.pop('token', None)
    use_auth_token = kwargs.pop('use_auth_token', None)
    local_files_only = kwargs.pop('local_files_only', False)
    revision = kwargs.pop('revision', None)
    subfolder = kwargs.pop('subfolder', '')
    image_processor_filename = kwargs.pop('image_processor_filename', IMAGE_PROCESSOR_NAME)
    from_pipeline = kwargs.pop('_from_pipeline', None)
    from_auto_class = kwargs.pop('_from_auto', False)
    if use_auth_token is not None:
        warnings.warn('The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.', FutureWarning)
        if token is not None:
            raise ValueError('`token` and `use_auth_token` are both specified. Please set only the argument `token`.')
        token = use_auth_token
    user_agent = {'file_type': 'image processor', 'from_auto_class': from_auto_class}
    if from_pipeline is not None:
        user_agent['using_pipeline'] = from_pipeline
    if is_offline_mode() and (not local_files_only):
        logger.info('Offline mode: forcing local_files_only=True')
        local_files_only = True
    pretrained_model_name_or_path = str(pretrained_model_name_or_path)
    is_local = os.path.isdir(pretrained_model_name_or_path)
    if os.path.isdir(pretrained_model_name_or_path):
        image_processor_file = os.path.join(pretrained_model_name_or_path, image_processor_filename)
    if os.path.isfile(pretrained_model_name_or_path):
        resolved_image_processor_file = pretrained_model_name_or_path
        is_local = True
    elif is_remote_url(pretrained_model_name_or_path):
        image_processor_file = pretrained_model_name_or_path
        resolved_image_processor_file = download_url(pretrained_model_name_or_path)
    else:
        image_processor_file = image_processor_filename
        try:
            resolved_image_processor_file = cached_file(pretrained_model_name_or_path, image_processor_file, cache_dir=cache_dir, force_download=force_download, proxies=proxies, resume_download=resume_download, local_files_only=local_files_only, token=token, user_agent=user_agent, revision=revision, subfolder=subfolder)
        except OSError:
            raise
        except Exception:
            raise OSError(f"Can't load image processor for '{pretrained_model_name_or_path}'. If you were trying to load it from 'https://huggingface.co/models', make sure you don't have a local directory with the same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory containing a {image_processor_filename} file")
    try:
        with open(resolved_image_processor_file, encoding='utf-8') as reader:
            text = reader.read()
        image_processor_dict = json.loads(text)
    except json.JSONDecodeError:
        raise OSError(f"It looks like the config file at '{resolved_image_processor_file}' is not a valid JSON file.")
    if is_local:
        logger.info(f'loading configuration file {resolved_image_processor_file}')
    else:
        logger.info(f'loading configuration file {image_processor_file} from cache at {resolved_image_processor_file}')
    return (image_processor_dict, kwargs)

From a `pretrained_model_name_or_path`, resolve to a dictionary of parameters, to be used for instantiating a
image processor of type [`~image_processor_utils.ImageProcessingMixin`] using `from_dict`.

Parameters:
    pretrained_model_name_or_path (`str` or `os.PathLike`):
        The identifier of the pre-trained checkpoint from which we want the dictionary of parameters.
    subfolder (`str`, *optional*, defaults to `""`):
        In case the relevant files are located inside a subfolder of the model repo on huggingface.co, you can
        specify the folder name here.
    image_processor_filename (`str`, *optional*, defaults to `"config.json"`):
        The name of the file in the model directory to use for the image processor config.

Returns:
    `Tuple[Dict, Dict]`: The dictionary(ies) that will be used to instantiate the image processor object.


def matches(self, regex: str) -> 'Builder':
    param_nodes = self._function_args_to_nodes(self.node, [regex])
    return self._to_builder(_evaluation.MatchesFunction(self.node.context, self.node, param_nodes))

The FHIRPath matches() function.

Args:
  regex: a regular expression to match against the parent element.

Returns:
  An expression that evaluates to True if the parent matches the given
  regular expression.


def _parse_getprop_output(self, output):
    output = output.decode('utf-8', errors='ignore').replace('\r\n', '\n')
    results = {}
    for line in output.split(']\n'):
        if not line:
            continue
        try:
            name, value = line.split(': ', 1)
        except ValueError:
            logging.debug('Failed to parse adb getprop line %s', line)
            continue
        name = name.strip()[1:-1]
        if value and value[0] == '[':
            value = value[1:]
        results[name] = value
    return results

Parses the raw output of `adb shell getprop` into a dictionary.

Args:
  output: byte str, the raw output of the `adb shell getprop` call.

Returns:
  dict, name-value pairs of the properties.


def get_grappler_config(optimizers_list):
    config = _config_pb2.ConfigProto()
    rewrite_options = config.graph_options.rewrite_options
    for optimizer in optimizers_list:
        rewrite_options.optimizers.append(optimizer)
    return config

Creates a tf.compat.v1.ConfigProto for configuring Grappler.

Args:
  optimizers_list: List of strings that represents the list of optimizers.

Returns:
  tf.ConfigProto.


def is_compatible(self, other: 'Schema') -> bool:
    if not isinstance(other, Schema):
        raise TypeError(f"Argument 'other' should be a Schema object. Encountered {other}.")
    for key_spec in other.keys():
        if key_spec not in self:
            return False
    for key_spec, field in self.items():
        if key_spec not in other:
            return False
        if not field.value.is_compatible(other[key_spec].value):
            return False
    return True

Returns whether current schema is compatible with the other schema.

NOTE(daiyip): schema A is compatible with schema B when:
schema A and schema B have the same keys, with compatible values specs.

Args:
  other: Other schema.

Returns:
  True if values that is acceptable to the other schema is acceptable to
    current schema.
Raises:
  TypeError: If `other` is not a schema object.


class CLIPImageProcessor(BaseImageProcessor):
    model_input_names = ['pixel_values']

    def __init__(self, do_resize: bool=True, size: Optional[Dict[str, int]]=None, resample: PILImageResampling=PILImageResampling.BICUBIC, do_center_crop: bool=True, crop_size: Optional[Dict[str, int]]=None, do_rescale: bool=True, rescale_factor: Union[int, float]=1 / 255, do_normalize: bool=True, image_mean: Optional[Union[float, List[float]]]=None, image_std: Optional[Union[float, List[float]]]=None, do_convert_rgb: bool=True, **kwargs) -> None:
        super().__init__(**kwargs)
        size = size if size is not None else {'shortest_edge': 224}
        size = get_size_dict(size, default_to_square=False)
        crop_size = crop_size if crop_size is not None else {'height': 224, 'width': 224}
        crop_size = get_size_dict(crop_size, default_to_square=True, param_name='crop_size')
        self.do_resize = do_resize
        self.size = size
        self.resample = resample
        self.do_center_crop = do_center_crop
        self.crop_size = crop_size
        self.do_rescale = do_rescale
        self.rescale_factor = rescale_factor
        self.do_normalize = do_normalize
        self.image_mean = image_mean if image_mean is not None else OPENAI_CLIP_MEAN
        self.image_std = image_std if image_std is not None else OPENAI_CLIP_STD
        self.do_convert_rgb = do_convert_rgb
        self._valid_processor_keys = ['images', 'do_resize', 'size', 'resample', 'do_center_crop', 'crop_size', 'do_rescale', 'rescale_factor', 'do_normalize', 'image_mean', 'image_std', 'do_convert_rgb', 'return_tensors', 'data_format', 'input_data_format']
        if 'use_square_size' in kwargs and kwargs['use_square_size']:
            self.size = {'height': size['shortest_edge'], 'width': size['shortest_edge']}
            delattr(self, 'use_square_size')

    def resize(self, image: np.ndarray, size: Dict[str, int], resample: PILImageResampling=PILImageResampling.BICUBIC, data_format: Optional[Union[str, ChannelDimension]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None, **kwargs) -> np.ndarray:
        """
        Resize an image. The shortest edge of the image is resized to size["shortest_edge"], with the longest edge
        resized to keep the input aspect ratio.

        Args:
            image (`np.ndarray`):
                Image to resize.
            size (`Dict[str, int]`):
                Size of the output image.
            resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BICUBIC`):
                Resampling filter to use when resiizing the image.
            data_format (`str` or `ChannelDimension`, *optional*):
                The channel dimension format of the image. If not provided, it will be the same as the input image.
            input_data_format (`ChannelDimension` or `str`, *optional*):
                The channel dimension format of the input image. If not provided, it will be inferred.
        """
        default_to_square = True
        if 'shortest_edge' in size:
            size = size['shortest_edge']
            default_to_square = False
        elif 'height' in size and 'width' in size:
            size = (size['height'], size['width'])
        else:
            raise ValueError("Size must contain either 'shortest_edge' or 'height' and 'width'.")
        output_size = get_resize_output_image_size(image, size=size, default_to_square=default_to_square, input_data_format=input_data_format)
        return resize(image, size=output_size, resample=resample, data_format=data_format, input_data_format=input_data_format, **kwargs)

    def preprocess(self, images: ImageInput, do_resize: Optional[bool]=None, size: Optional[Dict[str, int]]=None, resample: PILImageResampling=None, do_center_crop: Optional[bool]=None, crop_size: Optional[int]=None, do_rescale: Optional[bool]=None, rescale_factor: Optional[float]=None, do_normalize: Optional[bool]=None, image_mean: Optional[Union[float, List[float]]]=None, image_std: Optional[Union[float, List[float]]]=None, do_convert_rgb: Optional[bool]=None, return_tensors: Optional[Union[str, TensorType]]=None, data_format: Optional[ChannelDimension]=ChannelDimension.FIRST, input_data_format: Optional[Union[str, ChannelDimension]]=None, **kwargs) -> PIL.Image.Image:
        """
        Preprocess an image or batch of images.

        Args:
            images (`ImageInput`):
                Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
                passing in images with pixel values between 0 and 1, set `do_rescale=False`.
            do_resize (`bool`, *optional*, defaults to `self.do_resize`):
                Whether to resize the image.
            size (`Dict[str, int]`, *optional*, defaults to `self.size`):
                Size of the image after resizing. Shortest edge of the image is resized to size["shortest_edge"], with
                the longest edge resized to keep the input aspect ratio.
            resample (`int`, *optional*, defaults to `self.resample`):
                Resampling filter to use if resizing the image. This can be one of the enum `PILImageResampling`. Only
                has an effect if `do_resize` is set to `True`.
            do_center_crop (`bool`, *optional*, defaults to `self.do_center_crop`):
                Whether to center crop the image.
            crop_size (`Dict[str, int]`, *optional*, defaults to `self.crop_size`):
                Size of the center crop. Only has an effect if `do_center_crop` is set to `True`.
            do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
                Whether to rescale the image.
            rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
                Rescale factor to rescale the image by if `do_rescale` is set to `True`.
            do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
                Whether to normalize the image.
            image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
                Image mean to use for normalization. Only has an effect if `do_normalize` is set to `True`.
            image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
                Image standard deviation to use for normalization. Only has an effect if `do_normalize` is set to
                `True`.
            do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):
                Whether to convert the image to RGB.
            return_tensors (`str` or `TensorType`, *optional*):
                The type of tensors to return. Can be one of:
                - Unset: Return a list of `np.ndarray`.
                - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
                - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
                - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
                - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
            data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
                The channel dimension format for the output image. Can be one of:
                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
                - Unset: Use the channel dimension format of the input image.
            input_data_format (`ChannelDimension` or `str`, *optional*):
                The channel dimension format for the input image. If unset, the channel dimension format is inferred
                from the input image. Can be one of:
                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
                - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
        """
        do_resize = do_resize if do_resize is not None else self.do_resize
        size = size if size is not None else self.size
        size = get_size_dict(size, param_name='size', default_to_square=False)
        resample = resample if resample is not None else self.resample
        do_center_crop = do_center_crop if do_center_crop is not None else self.do_center_crop
        crop_size = crop_size if crop_size is not None else self.crop_size
        crop_size = get_size_dict(crop_size, param_name='crop_size', default_to_square=True)
        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
        image_mean = image_mean if image_mean is not None else self.image_mean
        image_std = image_std if image_std is not None else self.image_std
        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
        validate_kwargs(captured_kwargs=kwargs.keys(), valid_processor_keys=self._valid_processor_keys)
        images = make_flat_list_of_images(images)
        if not valid_images(images):
            raise ValueError('Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, torch.Tensor, tf.Tensor or jax.ndarray.')
        validate_preprocess_arguments(do_rescale=do_rescale, rescale_factor=rescale_factor, do_normalize=do_normalize, image_mean=image_mean, image_std=image_std, do_center_crop=do_center_crop, crop_size=crop_size, do_resize=do_resize, size=size, resample=resample)
        if do_convert_rgb:
            images = [convert_to_rgb(image) for image in images]
        images = [to_numpy_array(image) for image in images]
        if do_rescale and is_scaled_image(images[0]):
            logger.warning_once('It looks like you are trying to rescale already rescaled images. If the input images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again.')
        if input_data_format is None:
            input_data_format = infer_channel_dimension_format(images[0])
        all_images = []
        for image in images:
            if do_resize:
                image = self.resize(image=image, size=size, resample=resample, input_data_format=input_data_format)
            if do_center_crop:
                image = self.center_crop(image=image, size=crop_size, input_data_format=input_data_format)
            if do_rescale:
                image = self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format)
            if do_normalize:
                image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
            all_images.append(image)
        images = [to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) for image in all_images]
        data = {'pixel_values': images}
        return BatchFeature(data=data, tensor_type=return_tensors)

Constructs a CLIP image processor.

Args:
    do_resize (`bool`, *optional*, defaults to `True`):
        Whether to resize the image's (height, width) dimensions to the specified `size`. Can be overridden by
        `do_resize` in the `preprocess` method.
    size (`Dict[str, int]` *optional*, defaults to `{"shortest_edge": 224}`):
        Size of the image after resizing. The shortest edge of the image is resized to size["shortest_edge"], with
        the longest edge resized to keep the input aspect ratio. Can be overridden by `size` in the `preprocess`
        method.
    resample (`PILImageResampling`, *optional*, defaults to `Resampling.BICUBIC`):
        Resampling filter to use if resizing the image. Can be overridden by `resample` in the `preprocess` method.
    do_center_crop (`bool`, *optional*, defaults to `True`):
        Whether to center crop the image to the specified `crop_size`. Can be overridden by `do_center_crop` in the
        `preprocess` method.
    crop_size (`Dict[str, int]` *optional*, defaults to 224):
        Size of the output image after applying `center_crop`. Can be overridden by `crop_size` in the `preprocess`
        method.
    do_rescale (`bool`, *optional*, defaults to `True`):
        Whether to rescale the image by the specified scale `rescale_factor`. Can be overridden by `do_rescale` in
        the `preprocess` method.
    rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
        Scale factor to use if rescaling the image. Can be overridden by `rescale_factor` in the `preprocess`
        method.
    do_normalize (`bool`, *optional*, defaults to `True`):
        Whether to normalize the image. Can be overridden by `do_normalize` in the `preprocess` method.
    image_mean (`float` or `List[float]`, *optional*, defaults to `[0.48145466, 0.4578275, 0.40821073]`):
        Mean to use if normalizing the image. This is a float or list of floats the length of the number of
        channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method.
    image_std (`float` or `List[float]`, *optional*, defaults to `[0.26862954, 0.26130258, 0.27577711]`):
        Standard deviation to use if normalizing the image. This is a float or list of floats the length of the
        number of channels in the image. Can be overridden by the `image_std` parameter in the `preprocess` method.
        Can be overridden by the `image_std` parameter in the `preprocess` method.
    do_convert_rgb (`bool`, *optional*, defaults to `True`):
        Whether to convert the image to RGB.


class SuperModel(SuperPreTrainedModel):

    def __init__(self, config: SuperConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.layers = nn.ModuleList([SuperDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)])
        self.norm = SuperRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.rotary_emb = SuperRotaryEmbedding(config=config)
        self.gradient_checkpointing = False
        self.post_init()

    def get_input_embeddings(self):
        return self.embed_tokens

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

    @can_return_tuple
    @add_start_docstrings_to_model_forward(SUPER_INPUTS_DOCSTRING)
    def forward(self, input_ids: torch.LongTensor=None, attention_mask: Optional[torch.Tensor]=None, position_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Union[Cache, list[torch.FloatTensor]]]=None, inputs_embeds: Optional[torch.FloatTensor]=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, CausalLMOutputWithPast]:
        out = super().forward(input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict, cache_position)
        out.logits *= 2 ** 4
        return out

    def _update_causal_mask(self, attention_mask: Union[torch.Tensor, 'BlockMask'], input_tensor: torch.Tensor, cache_position: torch.Tensor, past_key_values: Cache, output_attentions: bool=False):
        if self.config._attn_implementation == 'flash_attention_2':
            if attention_mask is not None and (attention_mask == 0.0).any():
                return attention_mask
            return None
        if self.config._attn_implementation == 'flex_attention':
            if isinstance(attention_mask, torch.Tensor):
                attention_mask = make_flex_block_causal_mask(attention_mask)
            return attention_mask
        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
        using_static_cache = isinstance(past_key_values, StaticCache)
        if self.config._attn_implementation == 'sdpa' and (not using_static_cache) and (not output_attentions):
            if AttentionMaskConverter._ignore_causal_mask_sdpa(attention_mask, inputs_embeds=input_tensor, past_key_values_length=past_seen_tokens, is_training=self.training):
                return None
        dtype = input_tensor.dtype
        sequence_length = input_tensor.shape[1]
        if using_static_cache:
            target_length = past_key_values.get_max_cache_shape()
        else:
            target_length = attention_mask.shape[-1] if isinstance(attention_mask, torch.Tensor) else past_seen_tokens + sequence_length + 1
        causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(attention_mask, sequence_length=sequence_length, target_length=target_length, dtype=dtype, cache_position=cache_position, batch_size=input_tensor.shape[0])
        if self.config._attn_implementation == 'sdpa' and attention_mask is not None and (attention_mask.device.type in ['cuda', 'xpu', 'npu']) and (not output_attentions):
            min_dtype = torch.finfo(dtype).min
            causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)
        return causal_mask

    @staticmethod
    def _prepare_4d_causal_attention_mask_with_cache_position(attention_mask: torch.Tensor, sequence_length: int, target_length: int, dtype: torch.dtype, cache_position: torch.Tensor, batch_size: int, **kwargs):
        """
        Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
        `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.

        Args:
            attention_mask (`torch.Tensor`):
                A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape
                `(batch_size, 1, query_length, key_value_length)`.
            sequence_length (`int`):
                The sequence length being processed.
            target_length (`int`):
                The target length: when generating with static cache, the mask should be as long as the static cache,
                to account for the 0 padding, the part of the cache that is not filled yet.
            dtype (`torch.dtype`):
                The dtype to use for the 4D attention mask.
            cache_position (`torch.Tensor`):
                Indices depicting the position of the input sequence tokens in the sequence.
            batch_size (`torch.Tensor`):
                Batch size.
        """
        if attention_mask is not None and attention_mask.dim() == 4:
            causal_mask = attention_mask
        else:
            min_dtype = torch.finfo(dtype).min
            causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device)
            if sequence_length != 1:
                causal_mask = torch.triu(causal_mask, diagonal=1)
            causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1)
            causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
            if attention_mask is not None:
                causal_mask = causal_mask.clone()
                mask_length = attention_mask.shape[-1]
                padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(causal_mask.device)
                padding_mask = padding_mask == 0
                causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(padding_mask, min_dtype)
        return causal_mask

Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`SuperDecoderLayer`]

Args:
    config: SuperConfig


def potentially_ragged_concat(tensors):
    if len(tensors) == 1:
        return tensors[0]
    elif isinstance(tensors[0], tf.SparseTensor):
        return tf.sparse.concat(axis=0, sp_inputs=tensors)
    elif isinstance(tensors[0], tf.RaggedTensor):
        return tf.concat(tensors, axis=0)
    non_batch_shapes = tf.stack([tf.shape(tensor)[1:] for tensor in tensors])
    constant_dims = tf.math.reduce_all(non_batch_shapes == non_batch_shapes[:1], axis=0)
    if tf.math.reduce_all(constant_dims).numpy().item():
        if _is_scalar(tensors[0]):
            return tf.stack(tensors, axis=0)
        else:
            return tf.concat(tensors, axis=0)
    constant_inner_dimensions = constant_dims.numpy().tolist()[::-1].index(False)
    if constant_inner_dimensions == 0:
        constant_inner_shape = None
    else:
        constant_inner_shape = tensors[0].shape[-constant_inner_dimensions:]
    return tf.ragged.constant([tensor.numpy() for tensor in tensors], inner_shape=constant_inner_shape).merge_dims(0, 1)

Concats `Tensor`s along their first dimension.

Args:
    tensors: List of `Tensor`s.

Returns:
    Concatenation of the inputs along the first dimension -- of type
    `np.ndarray` if all input shapes are compatible, or `tf.RaggedTensor`
    if not.


def _partitioner(shape, dtype):
    if axis >= len(shape):
        raise ValueError(f'Cannot partition variable along axis {axis} when shape is only {shape}')
    dtype = dtypes.as_dtype(dtype)
    if dtype.base_dtype == dtypes.string:
        bytes_per_element = bytes_per_string_element
    else:
        bytes_per_element = dtype.size
    total_size_bytes = shape.num_elements() * bytes_per_element
    partitions = total_size_bytes / min_slice_size
    partitions_list = [1] * len(shape)
    partitions_list[axis] = max(1, min(shape.dims[axis].value, max_partitions, int(math.ceil(partitions))))
    return partitions_list

Partitioner that partitions list for a variable of given shape and type.

Ex: Consider partitioning a variable of type float32 with
  shape=[1024, 1024].
  If `max_partitions` >= 16, this function would return
    [(1024 * 1024 * 4) / (256 * 1024), 1] = [16, 1].
  If `max_partitions` < 16, this function would return
    [`max_partitions`, 1].

Args:
  shape: Shape of the variable.
  dtype: Type of the variable.

Returns:
  List of partitions for each axis (currently only one axis can be
  partitioned).

Raises:
  ValueError: If axis to partition along does not exist for the variable.


def lu_matrix_inverse(lower_upper, perm, validate_args=False, name=None):
    with ops.name_scope(name or 'lu_matrix_inverse'):
        lower_upper = ops.convert_to_tensor(lower_upper, dtype_hint=dtypes.float32, name='lower_upper')
        perm = ops.convert_to_tensor(perm, dtype_hint=dtypes.int32, name='perm')
        assertions = lu_reconstruct_assertions(lower_upper, perm, validate_args)
        if assertions:
            with ops.control_dependencies(assertions):
                lower_upper = array_ops.identity(lower_upper)
                perm = array_ops.identity(perm)
        shape = array_ops.shape(lower_upper)
        return lu_solve(lower_upper, perm, rhs=eye(shape[-1], batch_shape=shape[:-2], dtype=lower_upper.dtype), validate_args=False)

Computes the inverse given the LU decomposition(s) of one or more matrices.

This op is conceptually identical to,

```python
inv_X = tf.lu_matrix_inverse(*tf.linalg.lu(X))
tf.assert_near(tf.matrix_inverse(X), inv_X)
# ==> True
```

Note: this function does not verify the implied matrix is actually invertible
nor is this condition checked even when `validate_args=True`.

Args:
  lower_upper: `lu` as returned by `tf.linalg.lu`, i.e., if `matmul(P,
    matmul(L, U)) = X` then `lower_upper = L + U - eye`.
  perm: `p` as returned by `tf.linag.lu`, i.e., if `matmul(P, matmul(L, U)) =
    X` then `perm = argmax(P)`.
  validate_args: Python `bool` indicating whether arguments should be checked
    for correctness. Note: this function does not verify the implied matrix is
      actually invertible, even when `validate_args=True`.
    Default value: `False` (i.e., don't validate arguments).
  name: Python `str` name given to ops managed by this object.
    Default value: `None` (i.e., 'lu_matrix_inverse').

Returns:
  inv_x: The matrix_inv, i.e.,
    `tf.matrix_inverse(tf.linalg.lu_reconstruct(lu, perm))`.

#### Examples

```python
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

x = [[[3., 4], [1, 2]],
     [[7., 8], [3, 4]]]
inv_x = tf.linalg.lu_matrix_inverse(*tf.linalg.lu(x))
tf.assert_near(tf.matrix_inverse(x), inv_x)
# ==> True
```


def reducer_override(self, obj):
    t = type(obj)
    try:
        is_anyclass = issubclass(t, type)
    except TypeError:
        is_anyclass = False
    if is_anyclass:
        return _class_reduce(obj)
    elif isinstance(obj, types.FunctionType):
        return self._function_reduce(obj)
    else:
        return NotImplemented

Type-agnostic reducing callback for function and classes.

For performance reasons, subclasses of the C `pickle.Pickler` class
cannot register custom reducers for functions and classes in the
dispatch_table attribute. Reducers for such types must instead
implemented via the special `reducer_override` method.

Note that this method will be called for any object except a few
builtin-types (int, lists, dicts etc.), which differs from reducers
in the Pickler's dispatch_table, each of them being invoked for
objects of a specific type only.

This property comes in handy for classes: although most classes are
instances of the ``type`` metaclass, some of them can be instances
of other custom metaclasses (such as enum.EnumMeta for example). In
particular, the metaclass will likely not be known in advance, and
thus cannot be special-cased using an entry in the dispatch_table.
reducer_override, among other things, allows us to register a
reducer that will be called for any class, independently of its
type.

Notes:

* reducer_override has the priority over dispatch_table-registered
reducers.
* reducer_override can be used to fix other limitations of
  cloudpickle for other types that suffered from type-specific
  reducers, such as Exceptions. See
  https://github.com/cloudpipe/cloudpickle/issues/248


def dynamic_slice(operand, start_indices, slice_sizes):
    operand = tf_np.asarray(operand).data
    start_indices = tf_np.asarray(start_indices, np.int32).data
    idx = _get_dynamic_indices(operand, start_indices, slice_sizes)
    if idx is not None:
        operand = array_ops.gather_nd(operand, idx)
    return tf_np.asarray(operand)

Slicing operation where the indices can be dynamic vlaues.

See the docstring of `jax.lax.dynamic_slice`
(https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.dynamic_slice.html)
for details.

Args:
  operand: an array to slice.
  start_indices: a vector of integers, one per dimension. The starts of the
    slice. The vector can be dynamic.
  slice_sizes: a list of integers, one per dimension. The sizes of the slice.

Returns:
  An array containing the slice, with shape equal to `slice_sizes`.


def _dict_to_feed(self, parse=True):
    if not self.raw_feed:
        return []
    headers = self.raw_feed[0]
    row = 1
    for item in self.feed:
        for key in iter(item.keys()):
            if key in headers:
                column = headers.index(key)
                value = self._parse_value(item[key]) if parse else item[key]
                while column >= len(self.raw_feed[row]):
                    self.raw_feed[row].append('')
                self.raw_feed[row][column] = value
        row += 1
    return self.raw_feed

Turns a feed into a list of strings to be written back to the feed.

Args:
  feed: Dictionary list to convert into a list of lists of strings.

Returns:
 List of list of strings representing the values of the feed.


def get_special_tokens_mask(self, token_ids_0: List[int], token_ids_1: Optional[List[int]]=None, already_has_special_tokens: bool=False) -> List[int]:
    if already_has_special_tokens:
        return super().get_special_tokens_mask(token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True)
    prefix_ones = [1] * len(self.prefix_tokens)
    suffix_ones = [1] * len(self.suffix_tokens)
    if token_ids_1 is None:
        return prefix_ones + [0] * len(token_ids_0) + suffix_ones
    return prefix_ones + [0] * len(token_ids_0) + [0] * len(token_ids_1) + suffix_ones

Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
special tokens using the tokenizer `prepare_for_model` method.

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

Returns:
    `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.


def _parse_lambda(lam):
    mod = inspect.getmodule(lam)
    f = inspect.getsourcefile(lam)
    def_line = lam.__code__.co_firstlineno
    lines = linecache.getlines(f, mod.__dict__)
    source = ''.join(lines)
    all_nodes = parse(source, preamble_len=0, single_node=False)
    search_nodes = []
    for node in all_nodes:
        if getattr(node, 'lineno', def_line) <= def_line:
            search_nodes.append(node)
        else:
            break
    lambda_nodes = []
    for node in search_nodes:
        lambda_nodes.extend((n for n in gast.walk(node) if isinstance(n, gast.Lambda)))
    candidates = []
    for ln in lambda_nodes:
        minl, maxl = (MAX_SIZE, 0)
        for n in gast.walk(ln):
            minl = min(minl, getattr(n, 'lineno', minl))
            lineno = getattr(n, 'lineno', maxl)
            end_lineno = getattr(n, 'end_lineno', None)
            if end_lineno is not None:
                lineno = end_lineno
            maxl = max(maxl, lineno)
        if minl <= def_line <= maxl:
            candidates.append((ln, minl, maxl))
    if len(candidates) == 1:
        (node, minl, maxl), = candidates
        return _without_context(node, lines, minl, maxl)
    elif not candidates:
        lambda_codes = '\n'.join([unparse(l) for l in lambda_nodes])
        raise errors.UnsupportedLanguageElementError(f'could not parse the source code of {lam}: no matching AST found among candidates:\n{lambda_codes}')
    matches = [v for v in candidates if _node_matches_argspec(v[0], lam)]
    if len(matches) == 1:
        (node, minl, maxl), = matches
        return _without_context(node, lines, minl, maxl)
    matches = '\n'.join(('Match {}:\n{}\n'.format(i, unparse(node, include_encoding_marker=False)) for i, (node, _, _) in enumerate(matches)))
    raise errors.UnsupportedLanguageElementError(f'could not parse the source code of {lam}: found multiple definitions with identical signatures at the location. This error may be avoided by defining each lambda on a single line and with unique argument names. The matching definitions were:\n{matches}')

Returns the AST and source code of given lambda function.

Args:
  lam: types.LambdaType, Python function/method/class

Returns:
  gast.AST, Text: the parsed AST node; the source code that was parsed to
  generate the AST (including any prefixes that this function may have added).


def SplitLine(self, line, width):
    lines = []
    chunk = ''
    w = 0
    keep = False
    for normal, control in self.SplitIntoNormalAndControl(line):
        keep = True
        while True:
            n = width - w
            w += len(normal)
            if w <= width:
                break
            lines.append(chunk + normal[:n])
            chunk = ''
            keep = False
            w = 0
            normal = normal[n:]
        chunk += normal + control
    if chunk or keep:
        lines.append(chunk)
    return lines

Splits line into width length chunks.

Args:
  line: The line to split.
  width: The width of each chunk except the last which could be smaller than
    width.

Returns:
  A list of chunks, all but the last with display width == width.


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

Output type of [`ElectraForPreTraining`].

Args:
    loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):
        Total loss of the ELECTRA objective.
    logits (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):
        Prediction scores of the head (scores for each token before SoftMax).
    hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
        shape `(batch_size, sequence_length, hidden_size)`.

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

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


def cond(self, name='cond'):
    with self._name_scope(name):
        return self._cond()

Returns the condition number of this linear operator.

Args:
  name:  A name for this `Op`.

Returns:
  Shape `[B1,...,Bb]` `Tensor` of same `dtype` as `self`.


def __str__(self):
    info = {'section': self._section, 'config': self.config, 'req_type': self._req_type, 'req': str(self.req), 'range': str(self.range), 'exclude': str(self.exclude), 'include': str(self.include), 'init': str(self._initialized)}
    req_str = '\n >>> _Reqs Instance <<<\n'
    req_str += 'Section: {section}\n'
    req_str += 'Configuration name: {config}\n'
    req_str += 'Requirement type: {req_type}\n'
    req_str += 'Requirement: {req}\n'
    req_str += 'Range: {range}\n'
    req_str += 'Exclude: {exclude}\n'
    req_str += 'Include: {include}\n'
    req_str += 'Initialized: {init}\n\n'
    return req_str.format(**info)

Prints a requirement and its components.

Returns:
  String that has concatenated information about a requirement.


def read_file_to_string(filename, binary_mode=False):
    if binary_mode:
        f = FileIO(filename, mode='rb')
    else:
        f = FileIO(filename, mode='r')
    return f.read()

Reads the entire contents of a file to a string.

Args:
  filename: string, path to a file
  binary_mode: whether to open the file in binary mode or not. This changes
    the type of the object returned.

Returns:
  contents of the file as a string or bytes.

Raises:
  errors.OpError: Raises variety of errors that are subtypes e.g.
  `NotFoundError` etc.


def GenerateTableHtml(items, keys_to_print, display_index=True):
    html = ''
    html += '

\n'
    html += '
\n'
    if display_index:
        html += 'index	'
    for h, mapper in keys_to_print:
        html += '%s	' % h
    html += '
\n'
    for idx, tensor in enumerate(items):
        html += '
\n'
        if display_index:
            html += '%d	' % idx
        for h, mapper in keys_to_print:
            val = tensor[h] if h in tensor else None
            val = val if mapper is None else mapper(val)
            html += '%s	\n' % val
        html += '
\n'
    html += '
\n'
    return html

Given a list of object values and keys to print, make an HTML table.

Args:
  items: Items to print an array of dicts.
  keys_to_print: (key, display_fn). `key` is a key in the object. i.e.
    items[0][key] should exist. display_fn is the mapping function on display.
    i.e. the displayed html cell will have the string returned by
    `mapping_fn(items[0][key])`.
  display_index: add a column which is the index of each row in `items`.

Returns:
  An html table.


class GenericType(Type):
    base_type: NamedType | ClassType | LateType
    parameters: tuple[TypeU, ...]

    @property
    def name(self):
        return self.base_type.name

    @property
    def element_type(self):
        """Type of the contained type, assuming we only have one type parameter."""
        element_type, = self.parameters
        return element_type

Generic type. Takes a base type and type parameters.

This is used for homogeneous tuples, lists, dictionaries, user classes, etc.

Attributes:
  base_type: The base type. Instance of Type.
  parameters: Type parameters. Tuple of instances of Type.


def advise(self, options):
    advise_pb = tfprof_output_pb2.AdviceProto()
    opts = _build_advisor_options(options)
    advise_pb.ParseFromString(print_mdl.Profile('advise'.encode('utf-8'), opts.SerializeToString()))
    return advise_pb

Automatically detect problems and generate reports.

Args:
  options: A dict of options. See ALL_ADVICE example above.

Returns:
  An Advise proto that contains the reports from all checkers.


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

Creates a regularizer from its config.

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

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

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

Returns:
    A regularizer instance.


def post_process_segmentation(self, outputs, target_sizes, threshold=0.9, mask_threshold=0.5):
    logger.warning_once('`post_process_segmentation` is deprecated and will be removed in v5 of Transformers, please use `post_process_semantic_segmentation`.')
    out_logits, raw_masks = (outputs.logits, outputs.pred_masks)
    empty_label = out_logits.shape[-1] - 1
    preds = []

    def to_tuple(tup):
        if isinstance(tup, tuple):
            return tup
        return tuple(tup.tolist())
    for cur_logits, cur_masks, size in zip(out_logits, raw_masks, target_sizes):
        cur_scores, cur_labels = cur_logits.softmax(-1).max(-1)
        keep = cur_labels.ne(empty_label) & (cur_scores > threshold)
        cur_scores = cur_scores[keep]
        cur_labels = cur_labels[keep]
        cur_masks = cur_masks[keep]
        cur_masks = nn.functional.interpolate(cur_masks[:, None], to_tuple(size), mode='bilinear').squeeze(1)
        cur_masks = (cur_masks.sigmoid() > mask_threshold) * 1
        predictions = {'scores': cur_scores, 'labels': cur_labels, 'masks': cur_masks}
        preds.append(predictions)
    return preds

Converts the output of [`DetrForSegmentation`] into image segmentation predictions. Only supports PyTorch.

Args:
    outputs ([`DetrSegmentationOutput`]):
        Raw outputs of the model.
    target_sizes (`torch.Tensor` of shape `(batch_size, 2)` or `List[Tuple]` of length `batch_size`):
        Torch Tensor (or list) corresponding to the requested final size (h, w) of each prediction.
    threshold (`float`, *optional*, defaults to 0.9):
        Threshold to use to filter out queries.
    mask_threshold (`float`, *optional*, defaults to 0.5):
        Threshold to use when turning the predicted masks into binary values.
Returns:
    `List[Dict]`: A list of dictionaries, each dictionary containing the scores, labels, and masks for an image
    in the batch as predicted by the model.


class KLDivergence(MeanMetricWrapper):

    def __init__(self, name='kullback_leibler_divergence', dtype=None):
        super(KLDivergence, self).__init__(kullback_leibler_divergence, name, dtype=dtype)

Computes Kullback-Leibler divergence metric between `y_true` and `y_pred`.

`metric = y_true * log(y_true / y_pred)`

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

Standalone usage:

>>> m = tf.keras.metrics.KLDivergence()
>>> m.update_state([[0, 1], [0, 0]], [[0.6, 0.4], [0.4, 0.6]])
>>> m.result().numpy()
0.45814306

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

Usage with `compile()` API:

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


def forward(self, pixel_values: Optional[torch.Tensor]=None, bool_masked_pos: Optional[torch.BoolTensor]=None, head_mask: Optional[torch.Tensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, interpolate_pos_encoding: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[tuple, BaseModelOutput]:
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    if pixel_values is None:
        raise ValueError('You have to specify pixel_values')
    head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
    embedding_output = self.embeddings(pixel_values, bool_masked_pos=bool_masked_pos, interpolate_pos_encoding=interpolate_pos_encoding)
    encoder_outputs = self.encoder(embedding_output, head_mask=head_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    sequence_output = encoder_outputs[0]
    sequence_output = self.layernorm(sequence_output)
    if not return_dict:
        head_outputs = (sequence_output,)
        return head_outputs + encoder_outputs[1:]
    return BaseModelOutput(last_hidden_state=sequence_output, hidden_states=encoder_outputs.hidden_states, attentions=encoder_outputs.attentions)

bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`, *optional*):
    Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).

Examples:

```python
>>> from transformers import AutoImageProcessor, ViTMSNModel
>>> import torch
>>> from PIL import Image
>>> import requests

>>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
>>> image = Image.open(requests.get(url, stream=True).raw)

>>> image_processor = AutoImageProcessor.from_pretrained("facebook/vit-msn-small")
>>> model = ViTMSNModel.from_pretrained("facebook/vit-msn-small")
>>> inputs = image_processor(images=image, return_tensors="pt")
>>> with torch.no_grad():
...     outputs = model(**inputs)
>>> last_hidden_states = outputs.last_hidden_state
```


def resize(self, images: 'torch.Tensor', size: SizeDict, interpolation: Optional['F.InterpolationMode']=None, size_divisor: Optional[int]=None) -> 'torch.Tensor':
    if interpolation is None:
        interpolation = self.resample
    shorter = size.shortest_edge
    longer = int(MAX_LONGER_EDGE / MAX_SHORTER_EDGE * shorter)
    heights = images.shape[-2]
    widths = images.shape[-1]
    if heights < widths:
        new_heights = shorter
        new_widths = widths * (shorter / heights)
    else:
        new_heights = heights * (shorter / widths)
        new_widths = shorter
    if max(new_heights, new_widths) > longer:
        scale = longer / max(new_heights, new_widths)
        new_heights = new_heights * scale
        new_widths = new_widths * scale
    new_heights = int(new_heights + 0.5)
    new_widths = int(new_widths + 0.5)
    if size_divisor is not None:
        new_heights = new_heights // size_divisor * size_divisor
        new_widths = new_widths // size_divisor * size_divisor
    return F.resize(images, [new_heights, new_widths], interpolation=interpolation)

Resize an image or batch of images to specified size.

Args:
    images (`torch.Tensor`): Image or batch of images to resize.
    size (`Dict[str, int]`): Size dictionary with shortest_edge key.
    interpolation (`F.InterpolationMode`, *optional*): Interpolation method to use.
    size_divisor (`int`, *optional*): Value to ensure height/width are divisible by.

Returns:
    `torch.Tensor`: Resized image or batch of images.


def gather_gpu_devices():
    try:
        dev_info = _gather_gpu_devices_proc()
        if not dev_info:
            raise ValueError('No devices found')
        return dev_info
    except (IOError, ValueError, errors.OpError):
        pass
    try:
        return _gather_gpu_devices_cudart()
    except (OSError, ValueError, NotImplementedError, errors.OpError):
        return []

Gather gpu device info.

Returns:
  A list of test_log_pb2.GPUInfo messages.


def _serialize_io(value, debug_path: Optional[str]=None, use_repr: bool=True, path_to_value: Optional[str]=None):
    if isinstance(value, (list, tuple)):
        return [_serialize_io(v, debug_path=debug_path, use_repr=use_repr, path_to_value=f'{path_to_value}_{i}') for i, v in enumerate(value)]
    if isinstance(value, dict):
        return {k: _serialize_io(v, debug_path=debug_path, use_repr=use_repr, path_to_value=f'{path_to_value}_{k}') for k, v in value.items()}
    if hasattr(value, '_local_tensor'):
        return _serialize_tensor_like_io(value._local_tensor, debug_path=debug_path, use_repr=use_repr, path_to_value=path_to_value)
    if isinstance(value, torch.Tensor):
        return _serialize_tensor_like_io(value, debug_path=debug_path, use_repr=use_repr, path_to_value=path_to_value)
    return _sanitize_repr_for_diff(repr(value))

Recursively build a JSON-serializable Python structure from `value`.
Tensors and DTensors become either sanitized repr strings, or are saved to disk as SafeTensors files and their
relative paths are recorded in the returned Python structure.
Lists/tuples/dicts are recursed into.
All memory addresses are replaced with a stable placeholder.

Args:
    value: Any Python object, often including torch Tensors, lists, dicts, etc.
    debug_path (`str`, *optional*, defaults to `None`): Directory to dump debug JSON and SafeTensors files.
    use_repr (bool, *optional*, defaults to `True`): Whether to save a `repr()`-ized version of the tensors as the
        `value` property in the asscoiated FULL_TENSORS.json file, or to store full tensors in separate SafeTensors
        files and store the relative path to that file in the `value` property.
    path_to_value (`str`, *optional*, defaults to `None`): The file name for the SafeTensors file holding the full
        tensor value if `use_repr=False`.

Returns:
    A nested Python structure (list, dict, or sanitized string) that is safe to json.dump.


def set_default_language(self, language: str):
    if language not in self.config.languages:
        raise ValueError(f'{self} does not have an adapter for {language}. Supported languages: {list(self.config.languages)}')
    self.config.default_language = language

Set the default language code for the model. This is used when the language is not specified in the input.

Args:
    language (`str`): The language code, such as `"en_XX"` or `"de_DE"`.


def atrous_conv2d_transpose(value, filters, output_shape, rate, padding, name=None):
    with ops.name_scope(name, 'atrous_conv2d_transpose', [value, filters, output_shape]) as name:
        value = ops.convert_to_tensor(value, name='value')
        filters = ops.convert_to_tensor(filters, name='filters')
        if not value.get_shape().dims[3].is_compatible_with(filters.get_shape()[3]):
            raise ValueError(f'`value` channel count must be compatible with `filters` input channel count. Received: value.shape={value.get_shape()} with channel count {value.get_shape()[3]} and filters.shape={filters.get_shape()} with input channel count {filters.get_shape()[3]}.')
        if rate < 1:
            raise ValueError(f'`rate` cannot be less than one. Received: rate={rate}')
        if rate == 1:
            return conv2d_transpose(value, filters, output_shape, strides=[1, 1, 1, 1], padding=padding, data_format='NHWC')
        output_shape_ = ops.convert_to_tensor(output_shape, name='output_shape')
        if not output_shape_.get_shape().is_compatible_with(tensor_shape.TensorShape([4])):
            raise ValueError(f'`output_shape` must have shape (4,). Received: output_shape={output_shape_.get_shape()}')
        if isinstance(output_shape, tuple):
            output_shape = list(output_shape)
        if isinstance(output_shape, (list, np.ndarray)):
            if not filters.get_shape().dims[2].is_compatible_with(output_shape[3]):
                raise ValueError(f'`output_shape` channel count must be compatible with `filters` output channel count. Received: output_shape={output_shape} with channel count {output_shape[3]} and filters.shape={filters.get_shape()} with output channel count {filters.get_shape()[3]}.')
        if padding == 'SAME':
            if filters.get_shape().is_fully_defined():
                filter_shape = filters.get_shape().as_list()
            else:
                filter_shape = array_ops.shape(filters)
            filter_height, filter_width = (filter_shape[0], filter_shape[1])
            filter_height_up = filter_height + (filter_height - 1) * (rate - 1)
            filter_width_up = filter_width + (filter_width - 1) * (rate - 1)
            pad_height = filter_height_up - 1
            pad_width = filter_width_up - 1
            pad_top = pad_height // 2
            pad_bottom = pad_height - pad_top
            pad_left = pad_width // 2
            pad_right = pad_width - pad_left
        elif padding == 'VALID':
            pad_top = 0
            pad_bottom = 0
            pad_left = 0
            pad_right = 0
        else:
            raise ValueError(f"`padding` must be either 'VALID' or 'SAME'. Received: padding={padding}")
        in_height = output_shape[1] + pad_top + pad_bottom
        in_width = output_shape[2] + pad_left + pad_right
        pad_bottom_extra = (rate - in_height % rate) % rate
        pad_right_extra = (rate - in_width % rate) % rate
        space_to_batch_pad = [[0, pad_bottom_extra], [0, pad_right_extra]]
        value = array_ops.space_to_batch(input=value, paddings=space_to_batch_pad, block_size=rate)
        input_sizes = [rate * rate * output_shape[0], (in_height + pad_bottom_extra) // rate, (in_width + pad_right_extra) // rate, output_shape[3]]
        value = gen_nn_ops.conv2d_backprop_input(input_sizes=input_sizes, filter=filters, out_backprop=value, strides=[1, 1, 1, 1], padding='VALID', data_format='NHWC')
        batch_to_space_crop = [[pad_top, pad_bottom + pad_bottom_extra], [pad_left, pad_right + pad_right_extra]]
        return array_ops.batch_to_space(input=value, crops=batch_to_space_crop, block_size=rate)

The transpose of `atrous_conv2d`.

This operation is sometimes called "deconvolution" after
(Zeiler et al., 2010), but is really the transpose (gradient) of
`atrous_conv2d` rather than an actual deconvolution.

Args:
  value: A 4-D `Tensor` of type `float`. It needs to be in the default `NHWC`
    format. Its shape is `[batch, in_height, in_width, in_channels]`.
  filters: A 4-D `Tensor` with the same type as `value` and shape
    `[filter_height, filter_width, out_channels, in_channels]`. `filters`'
    `in_channels` dimension must match that of `value`. Atrous convolution is
    equivalent to standard convolution with upsampled filters with effective
    height `filter_height + (filter_height - 1) * (rate - 1)` and effective
    width `filter_width + (filter_width - 1) * (rate - 1)`, produced by
    inserting `rate - 1` zeros along consecutive elements across the
    `filters`' spatial dimensions.
  output_shape: A 1-D `Tensor` of shape representing the output shape of the
    deconvolution op, of form `[batch, out_height, out_width, out_channels]`.
  rate: A positive int32. The stride with which we sample input values across
    the `height` and `width` dimensions. Equivalently, the rate by which we
    upsample the filter values by inserting zeros across the `height` and
    `width` dimensions. In the literature, the same parameter is sometimes
    called `input stride` or `dilation`.
  padding: A string, either `'VALID'` or `'SAME'`. The padding algorithm. See
    [here](https://www.tensorflow.org/api_docs/python/tf/nn#notes_on_padding_2)
    for more information.
  name: Optional name for the returned tensor.

Returns:
  A `Tensor` with the same type as `value`.

Raises:
  ValueError: If input/output depth does not match `filters`' shape, or if
    padding is other than `'VALID'` or `'SAME'`, or if the `rate` is less
    than one, or if the output_shape is not a tensor with 4 elements.

References:
  Deconvolutional Networks:
    [Zeiler et al., 2010]
    (https://ieeexplore.ieee.org/abstract/document/5539957)
    ([pdf]
    (http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.232.4023&rep=rep1&type=pdf))


def save(self, file_prefix, session=None, options=None):
    graph_building = not context.executing_eagerly()
    if graph_building:
        if ops.inside_function():
            raise NotImplementedError('Calling tf.train.Checkpoint.save() from a function is not supported, as save() modifies saving metadata in ways not supported by TensorFlow Operations. Consider using tf.train.Checkpoint.write(), a lower-level API which does not update metadata. tf.train.latest_checkpoint and related APIs will not see this checkpoint.')
        if session is None:
            session = get_session()
        if self._save_counter is None:
            session.run(self.save_counter.initializer)
    if not graph_building or self._save_assign_op is None:
        with ops.colocate_with(self.save_counter):
            assign_op = self.save_counter.assign_add(1, read_value=True)
        if graph_building:
            self._save_assign_op = data_structures.NoDependency(assign_op)
    if graph_building:
        checkpoint_number = session.run(self._save_assign_op)
    else:
        checkpoint_number = assign_op.numpy()
    file_path = self.write('%s-%d' % (file_prefix, checkpoint_number), session=session, options=options)
    checkpoint_management.update_checkpoint_state_internal(save_dir=os.path.dirname(file_prefix), model_checkpoint_path=file_path, all_model_checkpoint_paths=[file_path], save_relative_paths=True)
    return file_path

Saves a training checkpoint and provides basic checkpoint management.

The saved checkpoint includes variables created by this object and any
trackable objects it depends on at the time `Checkpoint.save()` is
called.

`save` is a basic convenience wrapper around the `write` method,
sequentially numbering checkpoints using `save_counter` and updating the
metadata used by `tf.train.latest_checkpoint`. More advanced checkpoint
management, for example garbage collection and custom numbering, may be
provided by other utilities which also wrap `write`
(`tf.train.CheckpointManager` for example).

Args:
  file_prefix: A prefix to use for the checkpoint filenames
    (/path/to/directory/and_a_prefix). Names are generated based on this
    prefix and `Checkpoint.save_counter`.
  session: The session to evaluate variables in. Ignored when executing
    eagerly. If not provided when graph building, the default session is
    used.
  options: Optional `tf.train.CheckpointOptions` object.

Returns:
  The full path to the checkpoint.


class TFSemanticSegmenterOutput(ModelOutput):
    loss: tf.Tensor | None = None
    logits: Optional[tf.Tensor] = None
    hidden_states: Tuple[tf.Tensor] | None = None
    attentions: Tuple[tf.Tensor] | None = None

Base class for outputs of semantic segmentation models.

Args:
    loss (`tf.Tensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
        Classification (or regression if config.num_labels==1) loss.
    logits (`tf.Tensor` of shape `(batch_size, config.num_labels, logits_height, logits_width)`):
        Classification scores for each pixel.

        <Tip warning={true}>

        The logits returned do not necessarily have the same size as the `pixel_values` passed as inputs. This is
        to avoid doing two interpolations and lose some quality when a user needs to resize the logits to the
        original image size as post-processing. You should always check your logits shape and resize as needed.

        </Tip>

    hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `tf.Tensor` (one for the output of the embeddings, if the model has an embedding layer, + one for
        the output of each layer) of shape `(batch_size, patch_size, hidden_size)`.

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

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


def can_return_loss(model_class):
    framework = infer_framework(model_class)
    if framework == 'tf':
        signature = inspect.signature(model_class.call)
    elif framework == 'pt':
        signature = inspect.signature(model_class.forward)
    else:
        signature = inspect.signature(model_class.__call__)
    for p in signature.parameters:
        if p == 'return_loss' and signature.parameters[p].default is True:
            return True
    return False

Check if a given model can return loss.

Args:
    model_class (`type`): The class of the model.


def FromTimestampToDateTime(ts):
    return datetime.datetime.utcfromtimestamp(ts)

Converts internal nss_cache timestamp to datetime object.

Args:
  ts: number of seconds since epoch
Returns:
  datetime object


def set_render_option(self, render_option):
    self._render_option = render_option

Sets the rendering option.

Args:
  render_option: (str) this parameter decides how the pipeline graph is
      rendered. See display.pipeline_graph_renderer for available options.


def forward(self, input_values: Optional[torch.Tensor], attention_mask: Optional[torch.Tensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[Tuple, UniSpeechSatForPreTrainingOutput]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    outputs = self.unispeech_sat(input_values, attention_mask=attention_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    transformer_features = outputs[0]
    extract_features = self.dropout_features(outputs[1])
    logits = extract_features
    loss = quantized_features = codevector_perplexity = None
    if not return_dict:
        if loss is not None:
            return (loss, logits, transformer_features, quantized_features, codevector_perplexity) + outputs[2:]
        return (logits, transformer_features, quantized_features, codevector_perplexity) + outputs[2:]
    return UniSpeechSatForPreTrainingOutput(loss=loss, logits=logits, projected_states=transformer_features, projected_quantized_states=quantized_features, codevector_perplexity=codevector_perplexity, hidden_states=outputs.hidden_states, attentions=outputs.attentions)

Example:

```python
>>> import torch
>>> from transformers import AutoFeatureExtractor, UniSpeechSatForPreTraining
>>> from transformers.models.unispeech_sat.modeling_unispeech_sat import _compute_mask_indices

>>> feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/unispeech-sat-base")
>>> model = UniSpeechSatForPreTraining.from_pretrained("microsoft/unispeech-sat-base")
>>> # TODO: Add full pretraining example
```


def kvlist_str(kvlist: List[Tuple[str, Any, Any]], compact: bool=True, verbose: bool=False, root_indent: int=0, *, label: Optional[str]=None, bracket_type: BracketType=BracketType.ROUND, custom_format: Optional[CustomFormatFn]=None, memo: Optional[Set[int]]=None, **kwargs) -> str:
    s = io.StringIO()
    is_first = True
    bracket_start, bracket_end = bracket_chars(bracket_type)
    child_indent = root_indent + 1 if label else root_indent
    body = io.StringIO()
    for k, v, d in kvlist:
        if isinstance(d, tuple):
            include_pair = True
            for sd in d:
                if sd == v:
                    include_pair = False
                    break
        else:
            include_pair = v != d
        if include_pair:
            if not is_first:
                body.write(',')
                body.write(' ' if compact else '\n')
            v = format(v, compact=compact, verbose=verbose, root_indent=child_indent, custom_format=custom_format, memo=memo, **kwargs)
            if not compact:
                body.write(_indent('', child_indent))
            if k:
                body.write(f'{k}={str_ext(v, custom_format, child_indent)}')
            else:
                body.write(str_ext(v, custom_format, child_indent))
            is_first = False
    if label and (not is_first) and (not compact):
        body.write('\n')
    body = body.getvalue()
    if label is None:
        return body
    else:
        s.write(label)
        s.write(bracket_start)
        if body:
            if not compact:
                s.write('\n')
            s.write(body)
            if not compact:
                s.write(_indent('', root_indent))
        s.write(bracket_end)
        return s.getvalue()

Formats a list key/value pairs into a comma delimited string.

Args:
  kvlist: List of tuples in format of
    (key, value, default_value or a tuple of default values)
  compact: If True, format value in kvlist in compact form.
  verbose: If True, format value in kvlist in verbose.
  root_indent: The indent should be applied for values in kvlist if they are
    multi-line.
  label: (Optional) If not None, add label to brace all kv pairs.
  bracket_type: Bracket type used for embracing the kv pairs. Applicable only
    when `name` is not None.
  custom_format: An optional custom format function, which will be applied to
    each value (and child values) in kvlist. If the function returns None, it
    will fall back to the default `pg.format`.
  memo: A set of object ids that have been formatted. Used to avoid
    infinite recursion in the formatting process.
  **kwargs: Keyword arguments that will be passed through unto child
    ``Formattable`` objects.
Returns:
  A formatted string from a list of key/value pairs delimited by comma.


def _cbnd(dh, dk, rho):
    dtype = rho.dtype
    h = tf.cast(-dh, dtype=dtype)
    k = tf.cast(-dk, dtype=dtype)
    hk = h * k
    bvn = tf.zeros_like(hk)
    hs = (h * h + k * k) / 2
    asr = tf.math.asin(rho)

    def transformed_bvn(hk, hs, asr):

        def transformed_bvn_distribution(x):
            hk_exp = tf.expand_dims(hk, axis=-1)
            hs_exp = tf.expand_dims(hs, axis=-1)
            asr_exp = tf.expand_dims(asr, axis=-1)
            sn1 = tf.math.sin(asr_exp * (1 * x + 1) / 2)
            return tf.math.exp((hk_exp * sn1 - hs_exp) / (1 - sn1 * sn1))
        ones = tf.ones_like(hk)
        res = integration.gauss_legendre(func=transformed_bvn_distribution, lower=-ones, upper=ones, num_points=20, dtype=dtype)
        return res
    bvn = bvn + transformed_bvn(hk, hs, asr)
    bvn = bvn * asr / (4 * np.pi)
    bvn = bvn + _ncdf(-h) * _ncdf(-k)
    return bvn

Computes values for the cumulative standard bivariate normal distribution.

More specifically, compultes `P(x > dh, y > dk)` where `x` and `y` are
standard normal variables with correlation `rho`.

Args:
  dh: A real `Tensor` representing lower integration limits for `x`.
  dk: A `Tensor` of the same dtype as `dh` and of compatible shape
    representing lower integration limits `y`.
  rho: A `Tensor` of the same dtype as `dh` and of compatible shape
    representing correlation coefficients.

Returns:
  A `Tensor` of cumulative distribution function values.

#### References:
[1] Genz, A., Numerical Computation of Rectangular Bivariate and Trivariate
      Normal and t Probabilities, 2004
  http://www.math.wsu.edu/faculty/genz/papers/bvnt.pdf


def gaussian_blur(images, kernel_size=(3, 3), sigma=(1.0, 1.0), data_format=None):
    if any_symbolic_tensors((images,)):
        return GaussianBlur(kernel_size=kernel_size, sigma=sigma, data_format=data_format).symbolic_call(images)
    return backend.image.gaussian_blur(images, kernel_size=kernel_size, sigma=sigma, data_format=data_format)

Applies a Gaussian blur to the image(s).

Args:
    images: Input image or batch of images. Must be 3D or 4D.
    kernel_size: A tuple of two integers, specifying the height and width
        of the Gaussian kernel.
    sigma: A tuple of two floats, specifying the standard deviation of
        the Gaussian kernel along height and width.
    data_format: A string specifying the data format of the input tensor.
        It can be either `"channels_last"` or `"channels_first"`.
        `"channels_last"` corresponds to inputs with shape
        `(batch, height, width, channels)`, while `"channels_first"`
        corresponds to inputs with shape `(batch, channels, height, width)`.
        If not specified, the value will default to
        `keras.config.image_data_format`.

Returns:
    Blurred image or batch of images.

Examples:

>>> x = np.random.random((2, 64, 80, 3))  # batch of 2 RGB images
>>> y = keras.ops.image.gaussian_blur(x)
>>> y.shape
(2, 64, 80, 3)

>>> x = np.random.random((64, 80, 3))  # single RGB image
>>> y = keras.ops.image.gaussian_blur(x)
>>> y.shape
(64, 80, 3)

>>> x = np.random.random((2, 3, 64, 80))  # batch of 2 RGB images
>>> y = keras.ops.image.gaussian_blur(
...     x, data_format="channels_first")
>>> y.shape
(2, 3, 64, 80)


def requires_grad(self) -> bool:
    if self._rot_mats is not None:
        return self._rot_mats.requires_grad
    elif self._quats is not None:
        return self._quats.requires_grad
    else:
        raise ValueError('Both rotations are None')

Returns the requires_grad property of the underlying rotation

Returns:
    The requires_grad property of the underlying tensor


def dispatch(op, args, kwargs):
    for dispatcher in getattr(op, FALLBACK_DISPATCH_ATTR):
        result = dispatcher.handle(args, kwargs)
        if result is not OpDispatcher.NOT_SUPPORTED:
            return result
    for dispatcher in _GLOBAL_DISPATCHERS:
        result = dispatcher.handle(op, args, kwargs)
        if result is not OpDispatcher.NOT_SUPPORTED:
            return result
    return OpDispatcher.NOT_SUPPORTED

Returns the result from the first successful dispatcher for a given op.

Calls the `handle` method of each `OpDispatcher` that has been registered
to handle `op`, and returns the value from the first successful handler.

Args:
  op: Python function: the operation to dispatch for.
  args: The arguments to the operation.
  kwargs: They keyword arguments to the operation.

Returns:
  The result of the operation, or `NOT_SUPPORTED` if no registered
  dispatcher can handle the given arguments.


def __init__(self, value, translator):
    self.value = value
    self.translator = translator

Creates a NestedValueProvider that wraps the provided ValueProvider.

Args:
  value: ValueProvider object to wrap
  translator: function that is applied to the ValueProvider
Raises:
  ``RuntimeValueProviderError``: if any of the provided objects are not
    accessible.


def _compute_offsets(self, token_ids, time_precision=0.02, segment_size=1500):
    offsets = []
    if 'torch' in str(type(token_ids)) and (hasattr(token_ids, 'cpu') and callable(token_ids.cpu)):
        token_ids = token_ids.cpu()
    token_ids = np.array(token_ids)
    if token_ids.shape[0] > 1 and len(token_ids.shape) > 1:
        raise ValueError('Can only process a single input at a time')
    timestamp_begin = self.all_special_ids[-1] + 1
    timestamp_tokens = token_ids >= timestamp_begin
    consecutive = np.where(timestamp_tokens[:-1] & timestamp_tokens[1:])[0] + 1
    if consecutive.shape[0] == 0 and timestamp_tokens.sum() <= 1:
        return []
    elif np.where(timestamp_tokens)[0][-1] + 1 not in consecutive:
        consecutive = np.append(consecutive, np.where(timestamp_tokens)[0][-1] + 1)
    last_slice = np.where(timestamp_tokens)[0][0]
    cur_max_timestamp = 0
    prev_segments_len = 0
    for current_slice in consecutive:
        sliced_tokens = token_ids[last_slice:current_slice]
        if len(sliced_tokens) > 1:
            start_timestamp_position = sliced_tokens[0].item() - timestamp_begin
            end_timestamp_position = sliced_tokens[-1].item() - timestamp_begin
            if start_timestamp_position < cur_max_timestamp:
                is_single_ending = last_slice >= 2 and (not (token_ids[last_slice - 2] >= timestamp_begin and token_ids[last_slice - 1] >= timestamp_begin))
                if is_single_ending:
                    prev_segments_len += segment_size
                else:
                    prev_segments_len += cur_max_timestamp
            cur_max_timestamp = end_timestamp_position
            sliced_tokens = self._preprocess_token_ids(sliced_tokens)
            text = self._decode(sliced_tokens)
            text = self._filter_timestamp_ids(text)
            offsets.append({'text': text, 'timestamp': (start_timestamp_position * time_precision + prev_segments_len * time_precision, end_timestamp_position * time_precision + prev_segments_len * time_precision)})
        last_slice = current_slice
    return offsets

Compute offsets for a given tokenized input

Args:
    token_ids (`Union[int, List[int], np.ndarray, torch.Tensor, tf.Tensor]`):
        List of tokenized input ids. Can be obtained using the `__call__` method.
    time_precision (`float`, *optional*, defaults to 0.02):
        The time ratio to convert from token to time.
    segment_size (`int`, *optional*, defaults to 1500):
        The number of features in the input mel spectrogram.


def _as_graph_element(obj):
    conv_fn = getattr(obj, '_as_graph_element', None)
    if conv_fn and callable(conv_fn):
        return conv_fn()
    return None

Convert `obj` to a graph element if possible, otherwise return `None`.

Args:
  obj: Object to convert.

Returns:
  The result of `obj._as_graph_element()` if that method is available;
      otherwise `None`.


def get_optimizer_cls_and_kwargs(args: TrainingArguments, model: Optional[PreTrainedModel]=None) -> tuple[Any, Any]:
    optim_args = {}
    if args.optim_args:
        for mapping in args.optim_args.replace(' ', '').split(','):
            key, value = mapping.split('=')
            optim_args[key] = value
    optimizer_kwargs = {'lr': args.learning_rate}
    adam_kwargs = {'betas': (args.adam_beta1, args.adam_beta2), 'eps': args.adam_epsilon}

    def setup_low_rank_optimizer(optimizer_name: str, optimizer_mapping: dict[str, Any], optim_kwargs: dict[str, Any], is_layerwise_supported: bool=True) -> tuple[Any, Any]:
        """
            Helper function to set up low-rank optimizers like GaLore and Apollo.

            Args:
                optimizer_name (str): Name of the optimizer.
                optimizer_mapping (dict): Mapping of optimizer names to their classes.
                optim_kwargs (dict): Keyword arguments for the optimizer.
                is_layerwise_supported (bool): Whether layerwise optimization is supported.

            Returns:
                Tuple[Any, Any]: Optimizer class and updated optimizer kwargs.
            """
        is_layerwise = optimizer_name.lower().endswith('layerwise')
        if is_layerwise and args.parallel_mode == ParallelMode.DISTRIBUTED and is_layerwise_supported:
            raise NotImplementedError(f'Layer-wise {optimizer_name} does not support DDP at this time')
        optimizer_cls = optimizer_mapping[optimizer_name]
        if args.optim_target_modules is None:
            raise ValueError(f'You need to define `optim_target_modules` to use {optimizer_name} optimizers')
        if not isinstance(args.optim_target_modules, (list, str)):
            raise ValueError(f"`optim_target_modules` must be a list of strings, a regex string, or 'all-linear'. Got: {args.optim_target_modules}")
        if model is None:
            raise ValueError(f'You need to pass a model to initialize {optimizer_name} optimizer.')
        all_linear = isinstance(args.optim_target_modules, str) and args.optim_target_modules.replace('_', '-') == 'all-linear'
        target_params_names = []
        for module_name, module in model.named_modules():
            target_module_exists, is_regex = check_target_module_exists(args.optim_target_modules, module_name, return_is_regex=True)
            if not isinstance(module, nn.Linear):
                if target_module_exists and (not is_regex):
                    logger.warning(f'{module_name} matched but ignored. {optimizer_name} only supports linear layers.')
                continue
            if not target_module_exists and (not all_linear):
                continue
            target_params_names.append(module_name + '.weight')
        if len(target_params_names) == 0:
            raise ValueError(f'No target modules found for {optimizer_name} ({args.optim_target_modules}).')
        target_params = [p for n, p in model.named_parameters() if n in target_params_names]
        non_target_params = [p for n, p in model.named_parameters() if n not in target_params_names]
        optim_kwargs.update(optim_args)
        param_groups = [{'params': non_target_params}, {'params': target_params, **optim_kwargs}]
        if is_layerwise:
            if args.gradient_accumulation_steps != 1:
                raise ValueError(f'Layerwise {optimizer_name} does not support gradient accumulation!')
            optimizer_dict = {}
            for param in non_target_params:
                optimizer_dict[param] = optimizer_cls([{'params': [param]}], **optimizer_kwargs)
            for param in target_params:
                optimizer_dict[param] = optimizer_cls([{'params': [param], **optim_kwargs}], **optimizer_kwargs)

            def optimizer_hook(param):
                if param.grad is not None:
                    optimizer_dict[param].step()
                    optimizer_dict[param].zero_grad()
            for param in model.parameters():
                if param.requires_grad:
                    param.register_post_accumulate_grad_hook(optimizer_hook)
            optimizer_cls = LayerWiseDummyOptimizer
            optimizer_kwargs.update({'optimizer_dict': optimizer_dict})
        optimizer_kwargs.update({'params': param_groups})
        return (optimizer_cls, optimizer_kwargs)
    if args.optim == OptimizerNames.ADAFACTOR:
        optimizer_cls = Adafactor
        optimizer_kwargs.update({'scale_parameter': False, 'relative_step': False})
    elif args.optim in [OptimizerNames.ADAMW_TORCH, OptimizerNames.ADAMW_TORCH_FUSED]:
        from torch.optim import AdamW
        optimizer_cls = AdamW
        optimizer_kwargs.update(adam_kwargs)
        if args.optim == OptimizerNames.ADAMW_TORCH_FUSED:
            optimizer_kwargs.update({'fused': True})
    elif args.optim == OptimizerNames.ADAMW_TORCH_XLA:
        try:
            from torch_xla.amp.syncfree import AdamW
            optimizer_cls = AdamW
            optimizer_kwargs.update(adam_kwargs)
        except ImportError:
            raise ValueError('Trainer failed to import syncfree AdamW from torch_xla.')
    elif args.optim == OptimizerNames.ADAMW_TORCH_NPU_FUSED:
        try:
            from torch_npu.optim import NpuFusedAdamW
            optimizer_cls = NpuFusedAdamW
            optimizer_kwargs.update(adam_kwargs)
        except ImportError:
            raise ValueError('Trainer failed to import FusedAdamW from torch_npu.')
    elif args.optim == OptimizerNames.ADAMW_APEX_FUSED:
        try:
            from apex.optimizers import FusedAdam
            optimizer_cls = FusedAdam
            optimizer_kwargs.update(adam_kwargs)
        except ImportError:
            raise ValueError('Trainer tried to instantiate apex FusedAdam but apex is not installed!')
    elif args.optim in [OptimizerNames.ADAMW_BNB, OptimizerNames.ADAMW_8BIT, OptimizerNames.PAGED_ADAMW, OptimizerNames.PAGED_ADAMW_8BIT, OptimizerNames.ADEMAMIX, OptimizerNames.ADEMAMIX_8BIT, OptimizerNames.PAGED_ADEMAMIX, OptimizerNames.PAGED_ADEMAMIX_8BIT, OptimizerNames.LION, OptimizerNames.LION_8BIT, OptimizerNames.PAGED_LION, OptimizerNames.PAGED_LION_8BIT, OptimizerNames.RMSPROP_BNB, OptimizerNames.RMSPROP_8BIT, OptimizerNames.RMSPROP_32BIT]:
        try:
            from bitsandbytes.optim import AdamW, Lion, RMSprop
            is_paged = False
            optim_bits = 32
            optimizer_cls = None
            additional_optim_kwargs = adam_kwargs
            if 'paged' in args.optim:
                is_paged = True
            if '8bit' in args.optim:
                optim_bits = 8
            if 'adam' in args.optim:
                optimizer_cls = AdamW
            elif 'lion' in args.optim:
                optimizer_cls = Lion
                additional_optim_kwargs = {'betas': (args.adam_beta1, args.adam_beta2)}
            elif 'rmsprop' in args.optim:
                optimizer_cls = RMSprop
                additional_optim_kwargs = optim_args
            elif 'ademamix' in args.optim:
                if is_bitsandbytes_available() and version.parse(importlib.metadata.version('bitsandbytes')) < version.parse('0.44.0'):
                    raise ValueError('The AdEMAMix optimizer is not supported by your current version of `bitsandbytes`. Please install `bitsandbytes` >= 0.44.0.')
                from bitsandbytes.optim import AdEMAMix
                optimizer_cls = AdEMAMix
                additional_optim_kwargs = {'betas': (float(optim_args.get('beta1', args.adam_beta1)), float(optim_args.get('beta2', args.adam_beta2)), float(optim_args.get('beta3', 0.9999))), 'alpha': float(optim_args.get('alpha', 5.0)), 'eps': float(optim_args.get('eps', args.adam_epsilon))}
                if 't_alpha' in optim_args:
                    additional_optim_kwargs['t_alpha'] = int(optim_args['t_alpha'])
                if 't_beta3' in optim_args:
                    additional_optim_kwargs['t_beta3'] = int(optim_args['t_beta3'])
            bnb_kwargs = {'optim_bits': optim_bits}
            if 'rmsprop' not in args.optim:
                bnb_kwargs['is_paged'] = is_paged
            optimizer_kwargs.update(additional_optim_kwargs)
            optimizer_kwargs.update(bnb_kwargs)
        except ImportError:
            raise ValueError('Trainer tried to instantiate bnb optimizer but `bitsandbytes` is not installed!')
        if is_bitsandbytes_available() and version.parse(importlib.metadata.version('bitsandbytes')) < version.parse('0.41.1'):
            logger.warning('You are using 8-bit optimizers with a version of `bitsandbytes` < 0.41.1. It is recommended to update your version as a major bug has been fixed in 8-bit optimizers.')
    elif args.optim == OptimizerNames.ADAMW_ANYPRECISION:
        try:
            from torchdistx.optimizers import AnyPrecisionAdamW
            optimizer_cls = AnyPrecisionAdamW
            optimizer_kwargs.update(adam_kwargs)
            optimizer_kwargs.update({'use_kahan_summation': strtobool(optim_args.get('use_kahan_summation', 'False')), 'momentum_dtype': getattr(torch, optim_args.get('momentum_dtype', 'float32')), 'variance_dtype': getattr(torch, optim_args.get('variance_dtype', 'float32')), 'compensation_buffer_dtype': getattr(torch, optim_args.get('compensation_buffer_dtype', 'bfloat16'))})
        except ImportError:
            raise ValueError('Please install https://github.com/pytorch/torchdistx')
    elif args.optim == OptimizerNames.SGD:
        optimizer_cls = torch.optim.SGD
    elif args.optim == OptimizerNames.ADAGRAD:
        optimizer_cls = torch.optim.Adagrad
    elif args.optim == OptimizerNames.RMSPROP:
        optimizer_cls = torch.optim.RMSprop
    elif args.optim in [OptimizerNames.GALORE_ADAMW, OptimizerNames.GALORE_ADAMW_8BIT, OptimizerNames.GALORE_ADAFACTOR, OptimizerNames.GALORE_ADAMW_LAYERWISE, OptimizerNames.GALORE_ADAMW_8BIT_LAYERWISE, OptimizerNames.GALORE_ADAFACTOR_LAYERWISE]:
        if not is_galore_torch_available():
            raise ImportError('You need to install `galore_torch` in order to use GaLore optimizers install it with `pip install git+https://github.com/jiaweizzhao/GaLore`')
        from galore_torch import GaLoreAdafactor, GaLoreAdamW, GaLoreAdamW8bit
        optimizer_mapping = {OptimizerNames.GALORE_ADAMW: GaLoreAdamW, OptimizerNames.GALORE_ADAMW_8BIT: GaLoreAdamW8bit, OptimizerNames.GALORE_ADAFACTOR: GaLoreAdafactor, OptimizerNames.GALORE_ADAMW_LAYERWISE: GaLoreAdamW, OptimizerNames.GALORE_ADAMW_8BIT_LAYERWISE: GaLoreAdamW8bit, OptimizerNames.GALORE_ADAFACTOR_LAYERWISE: GaLoreAdafactor}
        galore_optim_kwargs = {'rank': int(optim_args.pop('rank', 128)), 'update_proj_gap': int(optim_args.pop('update_proj_gap', 200)), 'scale': float(optim_args.pop('scale', 0.25)), 'proj_type': optim_args.pop('proj_type', 'std')}
        optimizer_cls, optimizer_kwargs = setup_low_rank_optimizer(args.optim, optimizer_mapping, galore_optim_kwargs)
        if args.optim == OptimizerNames.GALORE_ADAFACTOR:
            optimizer_kwargs.update({'scale_parameter': False, 'relative_step': False})
    elif args.optim in [OptimizerNames.APOLLO_ADAMW, OptimizerNames.APOLLO_ADAMW_LAYERWISE]:
        if not is_apollo_torch_available():
            raise ImportError('You need to install `apollo_torch` in order to use APOLLO optimizers install it with `pip install git+https://github.com/zhuhanqing/APOLLO`')
        from apollo_torch import APOLLOAdamW
        optimizer_mapping = {OptimizerNames.APOLLO_ADAMW: APOLLOAdamW, OptimizerNames.APOLLO_ADAMW_LAYERWISE: APOLLOAdamW}
        apollo_optim_kwargs = {'rank': int(optim_args.pop('rank', 128)), 'proj': optim_args.pop('proj', 'random'), 'scale_type': optim_args.pop('scale_type', 'channel'), 'update_proj_gap': int(optim_args.pop('update_proj_gap', 200)), 'scale': float(optim_args.pop('scale', 1.0)), 'proj_type': optim_args.pop('proj_type', 'std')}
        apollo_optim_kwargs.update(adam_kwargs)
        optimizer_cls, optimizer_kwargs = setup_low_rank_optimizer(args.optim, optimizer_mapping, apollo_optim_kwargs)
    elif args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]:
        if not is_lomo_available():
            raise ImportError('You need to install `lomo_optim` in order to use LOMO optimizers install it with `pip install lomo-optim`')
        if not is_accelerate_available('0.30.0'):
            raise ImportError('You need to have `accelerate>=0.30.0` to be able to use LOMO optimizers')
        if model is None:
            raise ValueError('You need to pass a `model` in order to correctly initialize a LOMO optimizer.')
        from lomo_optim import AdaLomo, Lomo
        if 'ada' in args.optim:
            optimizer_cls = AdaLomo
        else:
            optimizer_cls = Lomo
        optimizer_kwargs.update({'model': model})
    elif args.optim == OptimizerNames.GROKADAMW:
        if not is_grokadamw_available():
            raise ValueError('Please install grokadamw with `pip install grokadamw`')
        from grokadamw import GrokAdamW
        optimizer_cls = GrokAdamW
        optimizer_kwargs.update({'alpha_init': float(optim_args.get('alpha_init', 0.98)), 'lamb': float(optim_args.get('lamb', 2.0)), 'gamma': float(optim_args.get('gamma', 0.1)), 'grokking_signal_decay_rate': float(optim_args.get('grokking_signal_decay_rate', 0.1)), 'gradient_clipping': float(optim_args.get('gradient_clipping', 1.0))})
    elif args.optim in [OptimizerNames.ADAMW_TORCH_4BIT, OptimizerNames.ADAMW_TORCH_8BIT]:
        if not is_torchao_available() or version.parse(importlib.metadata.version('torchao')) < version.parse('0.4.0'):
            raise ImportError('You need to have `torchao>=0.4.0` in order to use torch 4-bit optimizers.Install it with `pip install torchao` or follow the instructions here: https://github.com/pytorch/ao')
        if version.parse(importlib.metadata.version('torch')) <= version.parse('2.4'):
            raise ImportError('You need to have `torch>2.4` in order to use torch 4-bit optimizers. Install it with `pip install --upgrade torch` it is available on pipy. Otherwise, you need to install torch nightly.')
        if version.parse(importlib.metadata.version('torchao')) >= version.parse('0.11.0'):
            from torchao.optim import AdamW4bit, AdamW8bit
        else:
            from torchao.prototype.low_bit_optim import AdamW4bit, AdamW8bit
        if args.optim == OptimizerNames.ADAMW_TORCH_4BIT:
            optimizer_cls = AdamW4bit
        elif args.optim == OptimizerNames.ADAMW_TORCH_8BIT:
            optimizer_cls = AdamW8bit
        else:
            raise ValueError('Invalid optimizer')
        optimizer_kwargs.update(adam_kwargs)
    elif args.optim in [OptimizerNames.SCHEDULE_FREE_RADAM, OptimizerNames.SCHEDULE_FREE_ADAMW, OptimizerNames.SCHEDULE_FREE_SGD]:
        if not is_schedulefree_available():
            raise ImportError('You need to install `schedulefree` in order to use schedulefree optimizers. Install it with `pip install schedulefree.`')
        if not is_accelerate_available('0.30.0'):
            raise ImportError('You need to have `accelerate>=0.30.0` to be able to use schedulefree optimizers')
        from schedulefree import AdamWScheduleFree, SGDScheduleFree
        additional_optim_kwargs = {}
        require_warmup = True
        if args.optim == OptimizerNames.SCHEDULE_FREE_RADAM:
            if not is_schedulefree_available('1.4.0'):
                raise ImportError('You need to install `schedulefree>=1.4.0` in order to use RAdamScheduleFree optimizer. Install it with `pip install schedulefree.`')
            from schedulefree import RAdamScheduleFree
            optimizer_cls = RAdamScheduleFree
            additional_optim_kwargs = adam_kwargs
            require_warmup = False
        elif args.optim == OptimizerNames.SCHEDULE_FREE_ADAMW:
            optimizer_cls = AdamWScheduleFree
            additional_optim_kwargs = adam_kwargs
        elif args.optim == OptimizerNames.SCHEDULE_FREE_SGD:
            optimizer_cls = SGDScheduleFree
        else:
            raise ValueError('Invalid schedulefree optimizer')
        additional_optim_kwargs['weight_decay'] = args.weight_decay
        if require_warmup:
            additional_optim_kwargs['warmup_steps'] = args.warmup_steps
        additional_optim_kwargs.update({'weight_lr_power': float(optim_args.get('weight_lr_power', 2.0)), 'r': float(optim_args.get('r', 0.0))})
        optimizer_kwargs.update(additional_optim_kwargs)
    else:
        raise ValueError(f'Trainer cannot instantiate unsupported optimizer: {args.optim}')
    return (optimizer_cls, optimizer_kwargs)

Returns the optimizer class and optimizer parameters based on the training arguments.

Args:
    args (`transformers.training_args.TrainingArguments`):
        The training arguments for the training session.


def init_cache(self, batch_size, max_length, encoder_outputs):
    decoder_input_ids = jnp.ones((batch_size, max_length), dtype='i4')
    decoder_attention_mask = jnp.ones_like(decoder_input_ids)
    decoder_position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(decoder_input_ids).shape[-1]), decoder_input_ids.shape)

    def _decoder_forward(module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs):
        decoder_module = module._get_decoder_module()
        return decoder_module(decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs)
    init_variables = self.module.init(jax.random.PRNGKey(0), decoder_input_ids=decoder_input_ids, decoder_attention_mask=decoder_attention_mask, decoder_position_ids=decoder_position_ids, encoder_hidden_states=encoder_outputs[0], init_cache=True, method=_decoder_forward)
    return unfreeze(init_variables['cache'])

Args:
    batch_size (`int`):
        batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache.
    max_length (`int`):
        maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized
        cache.
    encoder_outputs (`Union[FlaxBaseModelOutput, tuple(tuple(jnp.ndarray)]`):
        `encoder_outputs` consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*:
        `attentions`). `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`, *optional*)
        is a sequence of hidden-states at the output of the last layer of the encoder. Used in the
        cross-attention of the decoder.


def extract_glimpse(input, size, offsets, centered=True, normalized=True, uniform_noise=True, name=None):
    return gen_image_ops.extract_glimpse(input=input, size=size, offsets=offsets, centered=centered, normalized=normalized, uniform_noise=uniform_noise, name=name)

Extracts a glimpse from the input tensor.

Returns a set of windows called glimpses extracted at location
`offsets` from the input tensor. If the windows only partially
overlaps the inputs, the non-overlapping areas will be filled with
random noise.

The result is a 4-D tensor of shape `[batch_size, glimpse_height,
glimpse_width, channels]`. The channels and batch dimensions are the
same as that of the input tensor. The height and width of the output
windows are specified in the `size` parameter.

The argument `normalized` and `centered` controls how the windows are built:

* If the coordinates are normalized but not centered, 0.0 and 1.0
  correspond to the minimum and maximum of each height and width
  dimension.
* If the coordinates are both normalized and centered, they range from
  -1.0 to 1.0. The coordinates (-1.0, -1.0) correspond to the upper
  left corner, the lower right corner is located at (1.0, 1.0) and the
  center is at (0, 0).
* If the coordinates are not normalized they are interpreted as
  numbers of pixels.

Usage Example:

>>> x = [[[[0.0],
...           [1.0],
...           [2.0]],
...          [[3.0],
...           [4.0],
...           [5.0]],
...          [[6.0],
...           [7.0],
...           [8.0]]]]
>>> tf.compat.v1.image.extract_glimpse(x, size=(2, 2), offsets=[[1, 1]],
...                                    centered=False, normalized=False)
<tf.Tensor: shape=(1, 2, 2, 1), dtype=float32, numpy=
array([[[[0.],
         [1.]],
        [[3.],
         [4.]]]], dtype=float32)>

Args:
  input: A `Tensor` of type `float32`. A 4-D float tensor of shape
    `[batch_size, height, width, channels]`.
  size: A `Tensor` of type `int32`. A 1-D tensor of 2 elements containing the
    size of the glimpses to extract.  The glimpse height must be specified
    first, following by the glimpse width.
  offsets: A `Tensor` of type `float32`. A 2-D integer tensor of shape
    `[batch_size, 2]` containing the y, x locations of the center of each
    window.
  centered: An optional `bool`. Defaults to `True`. indicates if the offset
    coordinates are centered relative to the image, in which case the (0, 0)
    offset is relative to the center of the input images. If false, the (0,0)
    offset corresponds to the upper left corner of the input images.
  normalized: An optional `bool`. Defaults to `True`. indicates if the offset
    coordinates are normalized.
  uniform_noise: An optional `bool`. Defaults to `True`. indicates if the
    noise should be generated using a uniform distribution or a Gaussian
    distribution.
  name: A name for the operation (optional).

Returns:
  A `Tensor` of type `float32`.


def _shape_tuple(self) -> NoReturn:
    raise NotImplementedError()

The shape of this Tensor, as a tuple.

This is more performant than tuple(shape().as_list()) as it avoids
two list and one object creation. Marked private for now as from an API
perspective, it would be better to have a single performant way of
getting a shape rather than exposing shape() and shape_tuple()
(and heaven forbid, shape_list() etc. as well!). Punting on that for now,
but ideally one would work things out and remove the need for this method.

Returns:
  tuple with the shape.


Return the default float type, as a string.

E.g. `'bfloat16'`, `'float16'`, `'float32'`, `'float64'`.

Returns:
    String, the current default float type.

Example:

>>> keras.config.floatx()
'float32'