class OneFormerImageProcessor(BaseImageProcessor):
    model_input_names = ['pixel_values', 'pixel_mask', 'task_inputs']

    @deprecate_kwarg('reduce_labels', new_name='do_reduce_labels', version='4.44.0')
    @deprecate_kwarg('max_size', version='4.27.0', warn_if_greater_or_equal_version=True)
    @filter_out_non_signature_kwargs(extra=['max_size', 'metadata', *INIT_SERVICE_KWARGS])
    def __init__(self, do_resize: bool=True, size: Optional[Dict[str, int]]=None, resample: PILImageResampling=PILImageResampling.BILINEAR, do_rescale: bool=True, rescale_factor: float=1 / 255, do_normalize: bool=True, image_mean: Optional[Union[float, List[float]]]=None, image_std: Optional[Union[float, List[float]]]=None, ignore_index: Optional[int]=None, do_reduce_labels: bool=False, repo_path: Optional[str]='shi-labs/oneformer_demo', class_info_file: Optional[str]=None, num_text: Optional[int]=None, num_labels: Optional[int]=None, **kwargs):
        super().__init__(**kwargs)
        self._max_size = kwargs.pop('max_size', 1333)
        size = size if size is not None else {'shortest_edge': 800, 'longest_edge': self._max_size}
        size = get_size_dict(size, max_size=self._max_size, default_to_square=False)
        if class_info_file is None:
            raise ValueError('You must provide a `class_info_file`')
        self.do_resize = do_resize
        self.size = size
        self.resample = resample
        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 IMAGENET_DEFAULT_MEAN
        self.image_std = image_std if image_std is not None else IMAGENET_DEFAULT_STD
        self.ignore_index = ignore_index
        self.do_reduce_labels = do_reduce_labels
        self.class_info_file = class_info_file
        self.repo_path = repo_path
        self.metadata = prepare_metadata(load_metadata(repo_path, class_info_file))
        self.num_text = num_text
        self.num_labels = num_labels

    @classmethod
    def from_dict(cls, image_processor_dict: Dict[str, Any], **kwargs):
        """
        Overrides the `from_dict` method from the base class to save support of deprecated `reduce_labels` in old configs
        """
        image_processor_dict = image_processor_dict.copy()
        if 'reduce_labels' in image_processor_dict:
            image_processor_dict['do_reduce_labels'] = image_processor_dict.pop('reduce_labels')
        return super().from_dict(image_processor_dict, **kwargs)

    def to_dict(self) -> Dict[str, Any]:
        """
        Serializes this instance to a Python dictionary. This method calls the superclass method and then removes the
        `_max_size` attribute from the dictionary.
        """
        image_processor_dict = super().to_dict()
        image_processor_dict.pop('_max_size', None)
        return image_processor_dict

    @deprecate_kwarg('max_size', version='4.27.0', warn_if_greater_or_equal_version=True)
    @filter_out_non_signature_kwargs(extra=['max_size'])
    def resize(self, image: np.ndarray, size: Dict[str, int], resample: PILImageResampling=PILImageResampling.BILINEAR, data_format=None, input_data_format: Optional[Union[str, ChannelDimension]]=None, **kwargs) -> np.ndarray:
        """
        Resize the image to the given size. Size can be min_size (scalar) or `(height, width)` tuple. If size is an
        int, smaller edge of the image will be matched to this number.
        """
        max_size = kwargs.pop('max_size', None)
        size = get_size_dict(size, max_size=max_size, default_to_square=False)
        if 'shortest_edge' in size and 'longest_edge' in size:
            size, max_size = (size['shortest_edge'], size['longest_edge'])
        elif 'height' in size and 'width' in size:
            size = (size['height'], size['width'])
            max_size = None
        else:
            raise ValueError(f"Size must contain 'height' and 'width' keys or 'shortest_edge' and 'longest_edge' keys. Got {size.keys()}.")
        size = get_oneformer_resize_output_image_size(image=image, size=size, max_size=max_size, default_to_square=False, input_data_format=input_data_format)
        image = resize(image, size=size, resample=resample, data_format=data_format, input_data_format=input_data_format)
        return image

    def rescale(self, image: np.ndarray, rescale_factor: float, data_format: Optional[Union[str, ChannelDimension]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None) -> np.ndarray:
        """
        Rescale the image by the given factor. image = image * rescale_factor.

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

    def convert_segmentation_map_to_binary_masks(self, segmentation_map: 'np.ndarray', instance_id_to_semantic_id: Optional[Dict[int, int]]=None, ignore_index: Optional[int]=None, do_reduce_labels: bool=False):
        do_reduce_labels = do_reduce_labels if do_reduce_labels is not None else self.do_reduce_labels
        ignore_index = ignore_index if ignore_index is not None else self.ignore_index
        return convert_segmentation_map_to_binary_masks(segmentation_map=segmentation_map, instance_id_to_semantic_id=instance_id_to_semantic_id, ignore_index=ignore_index, do_reduce_labels=do_reduce_labels)

    def __call__(self, images, task_inputs=None, segmentation_maps=None, **kwargs) -> BatchFeature:
        return self.preprocess(images, task_inputs=task_inputs, segmentation_maps=segmentation_maps, **kwargs)

    def _preprocess(self, image: ImageInput, do_resize: Optional[bool]=None, size: Optional[Dict[str, int]]=None, resample: PILImageResampling=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, input_data_format: Optional[Union[str, ChannelDimension]]=None):
        if do_resize:
            image = self.resize(image, size=size, resample=resample, input_data_format=input_data_format)
        if do_rescale:
            image = self.rescale(image, rescale_factor=rescale_factor, input_data_format=input_data_format)
        if do_normalize:
            image = self.normalize(image, mean=image_mean, std=image_std, input_data_format=input_data_format)
        return image

    def _preprocess_image(self, image: ImageInput, do_resize: Optional[bool]=None, size: Optional[Dict[str, int]]=None, resample: PILImageResampling=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, data_format: Optional[Union[str, ChannelDimension]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None) -> np.ndarray:
        """Preprocesses a single image."""
        image = to_numpy_array(image)
        if do_rescale and is_scaled_image(image):
            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(image)
        image = self._preprocess(image=image, do_resize=do_resize, size=size, resample=resample, do_rescale=do_rescale, rescale_factor=rescale_factor, do_normalize=do_normalize, image_mean=image_mean, image_std=image_std, input_data_format=input_data_format)
        if data_format is not None:
            image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
        return image

    def _preprocess_mask(self, segmentation_map: ImageInput, do_resize: Optional[bool]=None, size: Optional[Dict[str, int]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None) -> np.ndarray:
        """Preprocesses a single mask."""
        segmentation_map = to_numpy_array(segmentation_map)
        if segmentation_map.ndim == 2:
            added_channel_dim = True
            segmentation_map = segmentation_map[None, ...]
            input_data_format = ChannelDimension.FIRST
        else:
            added_channel_dim = False
            if input_data_format is None:
                input_data_format = infer_channel_dimension_format(segmentation_map, num_channels=1)
        segmentation_map = self._preprocess(image=segmentation_map, do_resize=do_resize, resample=PILImageResampling.NEAREST, size=size, do_rescale=False, do_normalize=False, input_data_format=input_data_format)
        if added_channel_dim:
            segmentation_map = segmentation_map.squeeze(0)
        return segmentation_map

    @filter_out_non_signature_kwargs()
    def preprocess(self, images: ImageInput, task_inputs: Optional[List[str]]=None, segmentation_maps: Optional[ImageInput]=None, instance_id_to_semantic_id: Optional[Dict[int, int]]=None, do_resize: Optional[bool]=None, size: Optional[Dict[str, int]]=None, resample: PILImageResampling=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, ignore_index: Optional[int]=None, do_reduce_labels: Optional[bool]=None, return_tensors: Optional[Union[str, TensorType]]=None, data_format: Union[str, ChannelDimension]=ChannelDimension.FIRST, input_data_format: Optional[Union[str, ChannelDimension]]=None) -> BatchFeature:
        if task_inputs is None:
            task_inputs = ['panoptic']
        do_resize = do_resize if do_resize is not None else self.do_resize
        size = size if size is not None else self.size
        size = get_size_dict(size, default_to_square=False, max_size=self._max_size)
        resample = resample if resample is not None else self.resample
        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
        ignore_index = ignore_index if ignore_index is not None else self.ignore_index
        do_reduce_labels = do_reduce_labels if do_reduce_labels is not None else self.do_reduce_labels
        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_resize=do_resize, size=size, resample=resample)
        if segmentation_maps is not None and (not valid_images(segmentation_maps)):
            raise ValueError('Invalid segmentation map type. Must be of type PIL.Image.Image, numpy.ndarray, torch.Tensor, tf.Tensor or jax.ndarray.')
        images = make_list_of_images(images)
        if segmentation_maps is not None:
            segmentation_maps = make_list_of_images(segmentation_maps, expected_ndims=2)
        if segmentation_maps is not None and len(images) != len(segmentation_maps):
            raise ValueError('Images and segmentation maps must have the same length.')
        images = [self._preprocess_image(image, do_resize=do_resize, size=size, resample=resample, do_rescale=do_rescale, rescale_factor=rescale_factor, do_normalize=do_normalize, image_mean=image_mean, image_std=image_std, data_format=data_format, input_data_format=input_data_format) for image in images]
        if segmentation_maps is not None:
            segmentation_maps = [self._preprocess_mask(segmentation_map, do_resize, size, input_data_format=input_data_format) for segmentation_map in segmentation_maps]
        encoded_inputs = self.encode_inputs(images, task_inputs, segmentation_maps, instance_id_to_semantic_id, ignore_index, do_reduce_labels, return_tensors, input_data_format=data_format)
        return encoded_inputs

    def _pad_image(self, image: np.ndarray, output_size: Tuple[int, int], constant_values: Union[float, Iterable[float]]=0, data_format: Optional[ChannelDimension]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None) -> np.ndarray:
        """
        Pad an image with zeros to the given size.
        """
        input_height, input_width = get_image_size(image, channel_dim=input_data_format)
        output_height, output_width = output_size
        pad_bottom = output_height - input_height
        pad_right = output_width - input_width
        padding = ((0, pad_bottom), (0, pad_right))
        padded_image = pad(image, padding, mode=PaddingMode.CONSTANT, constant_values=constant_values, data_format=data_format, input_data_format=input_data_format)
        return padded_image

    def pad(self, images: List[np.ndarray], constant_values: Union[float, Iterable[float]]=0, return_pixel_mask: bool=True, return_tensors: Optional[Union[str, TensorType]]=None, data_format: Optional[ChannelDimension]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None) -> BatchFeature:
        """
        Pads a batch of images to the bottom and right of the image with zeros to the size of largest height and width
        in the batch and optionally returns their corresponding pixel mask.

        Args:
            image (`np.ndarray`):
                Image to pad.
            constant_values (`float` or `Iterable[float]`, *optional*):
                The value to use for the padding if `mode` is `"constant"`.
            return_pixel_mask (`bool`, *optional*, defaults to `True`):
                Whether to return a pixel mask.
            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 (`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.
        """
        pad_size = get_max_height_width(images, input_data_format=input_data_format)
        padded_images = [self._pad_image(image, pad_size, constant_values=constant_values, data_format=data_format, input_data_format=input_data_format) for image in images]
        data = {'pixel_values': padded_images}
        if return_pixel_mask:
            masks = [make_pixel_mask(image=image, output_size=pad_size, input_data_format=input_data_format) for image in images]
            data['pixel_mask'] = masks
        return BatchFeature(data=data, tensor_type=return_tensors)

    def get_semantic_annotations(self, label, num_class_obj):
        annotation_classes = label['classes']
        annotation_masks = label['masks']
        texts = ['a semantic photo'] * self.num_text
        classes = []
        masks = []
        for idx in range(len(annotation_classes)):
            class_id = annotation_classes[idx]
            mask = annotation_masks[idx]
            if not np.all(mask is False):
                if class_id not in classes:
                    cls_name = self.metadata[str(class_id)]
                    classes.append(class_id)
                    masks.append(mask)
                    num_class_obj[cls_name] += 1
                else:
                    idx = classes.index(class_id)
                    masks[idx] += mask
                    masks[idx] = np.clip(masks[idx], 0, 1)
        num = 0
        for i, cls_name in enumerate(self.metadata['class_names']):
            if num_class_obj[cls_name] > 0:
                for _ in range(num_class_obj[cls_name]):
                    if num >= len(texts):
                        break
                    texts[num] = f'a photo with a {cls_name}'
                    num += 1
        classes = np.array(classes)
        masks = np.array(masks)
        return (classes, masks, texts)

    def get_instance_annotations(self, label, num_class_obj):
        annotation_classes = label['classes']
        annotation_masks = label['masks']
        texts = ['an instance photo'] * self.num_text
        classes = []
        masks = []
        for idx in range(len(annotation_classes)):
            class_id = annotation_classes[idx]
            mask = annotation_masks[idx]
            if class_id in self.metadata['thing_ids']:
                if not np.all(mask is False):
                    cls_name = self.metadata[str(class_id)]
                    classes.append(class_id)
                    masks.append(mask)
                    num_class_obj[cls_name] += 1
        num = 0
        for i, cls_name in enumerate(self.metadata['class_names']):
            if num_class_obj[cls_name] > 0:
                for _ in range(num_class_obj[cls_name]):
                    if num >= len(texts):
                        break
                    texts[num] = f'a photo with a {cls_name}'
                    num += 1
        classes = np.array(classes)
        masks = np.array(masks)
        return (classes, masks, texts)

    def get_panoptic_annotations(self, label, num_class_obj):
        annotation_classes = label['classes']
        annotation_masks = label['masks']
        texts = ['an panoptic photo'] * self.num_text
        classes = []
        masks = []
        for idx in range(len(annotation_classes)):
            class_id = annotation_classes[idx]
            mask = annotation_masks[idx].data
            if not np.all(mask is False):
                cls_name = self.metadata[str(class_id)]
                classes.append(class_id)
                masks.append(mask)
                num_class_obj[cls_name] += 1
        num = 0
        for i, cls_name in enumerate(self.metadata['class_names']):
            if num_class_obj[cls_name] > 0:
                for _ in range(num_class_obj[cls_name]):
                    if num >= len(texts):
                        break
                    texts[num] = f'a photo with a {cls_name}'
                    num += 1
        classes = np.array(classes)
        masks = np.array(masks)
        return (classes, masks, texts)

    def encode_inputs(self, pixel_values_list: List[ImageInput], task_inputs: List[str], segmentation_maps: ImageInput=None, instance_id_to_semantic_id: Optional[Union[List[Dict[int, int]], Dict[int, int]]]=None, ignore_index: Optional[int]=None, do_reduce_labels: bool=False, return_tensors: Optional[Union[str, TensorType]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None):
        """
        Pad images up to the largest image in a batch and create a corresponding `pixel_mask`.

        OneFormer addresses semantic segmentation with a mask classification paradigm, thus input segmentation maps
        will be converted to lists of binary masks and their respective labels. Let's see an example, assuming
        `segmentation_maps = [[2,6,7,9]]`, the output will contain `mask_labels =
        [[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,1]]` (four binary masks) and `class_labels = [2,6,7,9]`, the labels for
        each mask.

        Args:
            pixel_values_list (`List[ImageInput]`):
                List of images (pixel values) to be padded. Each image should be a tensor of shape `(channels, height,
                width)`.

            task_inputs (`List[str]`):
                List of task values.

            segmentation_maps (`ImageInput`, *optional*):
                The corresponding semantic segmentation maps with the pixel-wise annotations.

             (`bool`, *optional*, defaults to `True`):
                Whether or not to pad images up to the largest image in a batch and create a pixel mask.

                If left to the default, will return a pixel mask that is:

                - 1 for pixels that are real (i.e. **not masked**),
                - 0 for pixels that are padding (i.e. **masked**).

            instance_id_to_semantic_id (`List[Dict[int, int]]` or `Dict[int, int]`, *optional*):
                A mapping between object instance ids and class ids. If passed, `segmentation_maps` is treated as an
                instance segmentation map where each pixel represents an instance id. Can be provided as a single
                dictionary with a global/dataset-level mapping or as a list of dictionaries (one per image), to map
                instance ids in each image separately.

            return_tensors (`str` or [`~file_utils.TensorType`], *optional*):
                If set, will return tensors instead of NumPy arrays. If set to `'pt'`, return PyTorch `torch.Tensor`
                objects.

            input_data_format (`str` or `ChannelDimension`, *optional*):
                The channel dimension format of the input image. If not provided, it will be inferred from the input
                image.

        Returns:
            [`BatchFeature`]: A [`BatchFeature`] with the following fields:

            - **pixel_values** -- Pixel values to be fed to a model.
            - **pixel_mask** -- Pixel mask to be fed to a model (when `=True` or if `pixel_mask` is in
              `self.model_input_names`).
            - **mask_labels** -- Optional list of mask labels of shape `(labels, height, width)` to be fed to a model
              (when `annotations` are provided).
            - **class_labels** -- Optional list of class labels of shape `(labels)` to be fed to a model (when
              `annotations` are provided). They identify the labels of `mask_labels`, e.g. the label of
              `mask_labels[i][j]` if `class_labels[i][j]`.
            - **text_inputs** -- Optional list of text string entries to be fed to a model (when `annotations` are
              provided). They identify the binary masks present in the image.
        """
        ignore_index = self.ignore_index if ignore_index is None else ignore_index
        do_reduce_labels = self.do_reduce_labels if do_reduce_labels is None else do_reduce_labels
        pixel_values_list = [to_numpy_array(pixel_values) for pixel_values in pixel_values_list]
        if input_data_format is None:
            input_data_format = infer_channel_dimension_format(pixel_values_list[0])
        pad_size = get_max_height_width(pixel_values_list, input_data_format=input_data_format)
        encoded_inputs = self.pad(pixel_values_list, return_tensors=return_tensors, input_data_format=input_data_format)
        annotations = None
        if segmentation_maps is not None:
            segmentation_maps = map(np.array, segmentation_maps)
            annotations = []
            for idx, segmentation_map in enumerate(segmentation_maps):
                if isinstance(instance_id_to_semantic_id, list):
                    instance_id = instance_id_to_semantic_id[idx]
                else:
                    instance_id = instance_id_to_semantic_id
                masks, classes = self.convert_segmentation_map_to_binary_masks(segmentation_map, instance_id, ignore_index=ignore_index, do_reduce_labels=do_reduce_labels)
                annotations.append({'masks': masks, 'classes': classes})
        if annotations is not None:
            mask_labels = []
            class_labels = []
            text_inputs = []
            num_class_obj = {}
            for cls_name in self.metadata['class_names']:
                num_class_obj[cls_name] = 0
            for i, label in enumerate(annotations):
                task = task_inputs[i]
                if task == 'semantic':
                    classes, masks, texts = self.get_semantic_annotations(label, num_class_obj)
                elif task == 'instance':
                    classes, masks, texts = self.get_instance_annotations(label, num_class_obj)
                elif task == 'panoptic':
                    classes, masks, texts = self.get_panoptic_annotations(label, num_class_obj)
                else:
                    raise ValueError(f'{task} was not expected, expected `semantic`, `instance` or `panoptic`')
                masks = [mask[None, ...] for mask in masks]
                masks = [self._pad_image(image=mask, output_size=pad_size, constant_values=ignore_index) for mask in masks]
                masks = np.concatenate(masks, axis=0)
                mask_labels.append(torch.from_numpy(masks))
                class_labels.append(torch.from_numpy(classes).long())
                text_inputs.append(texts)
            encoded_inputs['mask_labels'] = mask_labels
            encoded_inputs['class_labels'] = class_labels
            encoded_inputs['text_inputs'] = text_inputs
        encoded_inputs['task_inputs'] = [f'the task is {task_input}' for task_input in task_inputs]
        return encoded_inputs

    def post_process_semantic_segmentation(self, outputs, target_sizes: Optional[List[Tuple[int, int]]]=None) -> 'torch.Tensor':
        """
        Converts the output of [`MaskFormerForInstanceSegmentation`] into semantic segmentation maps. Only supports
        PyTorch.

        Args:
            outputs ([`MaskFormerForInstanceSegmentation`]):
                Raw outputs of the model.
            target_sizes (`List[Tuple[int, int]]`, *optional*):
                List of length (batch_size), where each list item (`Tuple[int, int]]`) corresponds to the requested
                final size (height, width) of each prediction. If left to None, predictions will not be resized.
        Returns:
            `List[torch.Tensor]`:
                A list of length `batch_size`, where each item is a semantic segmentation map of shape (height, width)
                corresponding to the target_sizes entry (if `target_sizes` is specified). Each entry of each
                `torch.Tensor` correspond to a semantic class id.
        """
        class_queries_logits = outputs.class_queries_logits
        masks_queries_logits = outputs.masks_queries_logits
        masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1]
        masks_probs = masks_queries_logits.sigmoid()
        segmentation = torch.einsum('bqc, bqhw -> bchw', masks_classes, masks_probs)
        batch_size = class_queries_logits.shape[0]
        if target_sizes is not None:
            if batch_size != len(target_sizes):
                raise ValueError('Make sure that you pass in as many target sizes as the batch dimension of the logits')
            semantic_segmentation = []
            for idx in range(batch_size):
                resized_logits = torch.nn.functional.interpolate(segmentation[idx].unsqueeze(dim=0), size=target_sizes[idx], mode='bilinear', align_corners=False)
                semantic_map = resized_logits[0].argmax(dim=0)
                semantic_segmentation.append(semantic_map)
        else:
            semantic_segmentation = segmentation.argmax(dim=1)
            semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])]
        return semantic_segmentation

    def post_process_instance_segmentation(self, outputs, task_type: str='instance', is_demo: bool=True, threshold: float=0.5, mask_threshold: float=0.5, overlap_mask_area_threshold: float=0.8, target_sizes: Optional[List[Tuple[int, int]]]=None, return_coco_annotation: Optional[bool]=False):
        """
        Converts the output of [`OneFormerForUniversalSegmentationOutput`] into image instance segmentation
        predictions. Only supports PyTorch.

        Args:
            outputs ([`OneFormerForUniversalSegmentationOutput`]):
                The outputs from [`OneFormerForUniversalSegmentationOutput`].
            task_type (`str`, *optional*, defaults to "instance"):
                The post processing depends on the task token input. If the `task_type` is "panoptic", we need to
                ignore the stuff predictions.
            is_demo (`bool`, *optional)*, defaults to `True`):
                Whether the model is in demo mode. If true, use threshold to predict final masks.
            threshold (`float`, *optional*, defaults to 0.5):
                The probability score threshold to keep predicted instance masks.
            mask_threshold (`float`, *optional*, defaults to 0.5):
                Threshold to use when turning the predicted masks into binary values.
            overlap_mask_area_threshold (`float`, *optional*, defaults to 0.8):
                The overlap mask area threshold to merge or discard small disconnected parts within each binary
                instance mask.
            target_sizes (`List[Tuple]`, *optional*):
                List of length (batch_size), where each list item (`Tuple[int, int]]`) corresponds to the requested
                final size (height, width) of each prediction in batch. If left to None, predictions will not be
                resized.
            return_coco_annotation (`bool`, *optional)*, defaults to `False`):
                Whether to return predictions in COCO format.

        Returns:
            `List[Dict]`: A list of dictionaries, one per image, each dictionary containing two keys:
            - **segmentation** -- a tensor of shape `(height, width)` where each pixel represents a `segment_id`, set
              to `None` if no mask if found above `threshold`. If `target_sizes` is specified, segmentation is resized
              to the corresponding `target_sizes` entry.
            - **segments_info** -- A dictionary that contains additional information on each segment.
                - **id** -- an integer representing the `segment_id`.
                - **label_id** -- An integer representing the label / semantic class id corresponding to `segment_id`.
                - **was_fused** -- a boolean, `True` if `label_id` was in `label_ids_to_fuse`, `False` otherwise.
                  Multiple instances of the same class / label were fused and assigned a single `segment_id`.
                - **score** -- Prediction score of segment with `segment_id`.
        """
        class_queries_logits = outputs.class_queries_logits
        masks_queries_logits = outputs.masks_queries_logits
        device = masks_queries_logits.device
        batch_size = class_queries_logits.shape[0]
        num_queries = class_queries_logits.shape[1]
        num_classes = class_queries_logits.shape[-1] - 1
        results: List[Dict[str, torch.Tensor]] = []
        for i in range(batch_size):
            scores = torch.nn.functional.softmax(class_queries_logits[i], dim=-1)[:, :-1]
            labels = torch.arange(num_classes, device=device).unsqueeze(0).repeat(num_queries, 1).flatten(0, 1)
            scores_per_image, topk_indices = scores.flatten(0, 1).topk(num_queries, sorted=False)
            labels_per_image = labels[topk_indices]
            topk_indices = torch.div(topk_indices, num_classes, rounding_mode='floor')
            mask_pred = masks_queries_logits[i][topk_indices]
            if is_demo:
                keep = scores_per_image > threshold
                scores_per_image = scores_per_image[keep]
                labels_per_image = labels_per_image[keep]
                mask_pred = mask_pred[keep]
            if task_type == 'panoptic':
                keep = torch.zeros_like(scores_per_image).bool()
                for j, lab in enumerate(labels_per_image):
                    keep[j] = lab in self.metadata['thing_ids']
                scores_per_image = scores_per_image[keep]
                labels_per_image = labels_per_image[keep]
                mask_pred = mask_pred[keep]
            if mask_pred.shape[0] <= 0:
                height, width = target_sizes[i] if target_sizes is not None else mask_pred.shape[1:]
                segmentation = torch.zeros((height, width)) - 1
                results.append({'segmentation': segmentation, 'segments_info': []})
                continue
            if 'ade20k' in self.class_info_file and (not is_demo) and ('instance' in task_type):
                for j in range(labels_per_image.shape[0]):
                    labels_per_image[j] = self.metadata['thing_ids'].index(labels_per_image[j].item())
            target_size = target_sizes[i] if target_sizes is not None else None
            segmentation, segments = compute_segments(mask_pred, scores_per_image, labels_per_image, mask_threshold, overlap_mask_area_threshold, set(), target_size)
            if return_coco_annotation:
                segmentation = convert_segmentation_to_rle(segmentation)
            results.append({'segmentation': segmentation, 'segments_info': segments})
        return results

    def post_process_panoptic_segmentation(self, outputs, threshold: float=0.5, mask_threshold: float=0.5, overlap_mask_area_threshold: float=0.8, label_ids_to_fuse: Optional[Set[int]]=None, target_sizes: Optional[List[Tuple[int, int]]]=None) -> List[Dict]:
        """
        Converts the output of [`MaskFormerForInstanceSegmentationOutput`] into image panoptic segmentation
        predictions. Only supports PyTorch.

        Args:
            outputs ([`MaskFormerForInstanceSegmentationOutput`]):
                The outputs from [`MaskFormerForInstanceSegmentation`].
            threshold (`float`, *optional*, defaults to 0.5):
                The probability score threshold to keep predicted instance masks.
            mask_threshold (`float`, *optional*, defaults to 0.5):
                Threshold to use when turning the predicted masks into binary values.
            overlap_mask_area_threshold (`float`, *optional*, defaults to 0.8):
                The overlap mask area threshold to merge or discard small disconnected parts within each binary
                instance mask.
            label_ids_to_fuse (`Set[int]`, *optional*):
                The labels in this state will have all their instances be fused together. For instance we could say
                there can only be one sky in an image, but several persons, so the label ID for sky would be in that
                set, but not the one for person.
            target_sizes (`List[Tuple]`, *optional*):
                List of length (batch_size), where each list item (`Tuple[int, int]]`) corresponds to the requested
                final size (height, width) of each prediction in batch. If left to None, predictions will not be
                resized.

        Returns:
            `List[Dict]`: A list of dictionaries, one per image, each dictionary containing two keys:
            - **segmentation** -- a tensor of shape `(height, width)` where each pixel represents a `segment_id`, set
              to `None` if no mask if found above `threshold`. If `target_sizes` is specified, segmentation is resized
              to the corresponding `target_sizes` entry.
            - **segments_info** -- A dictionary that contains additional information on each segment.
                - **id** -- an integer representing the `segment_id`.
                - **label_id** -- An integer representing the label / semantic class id corresponding to `segment_id`.
                - **was_fused** -- a boolean, `True` if `label_id` was in `label_ids_to_fuse`, `False` otherwise.
                  Multiple instances of the same class / label were fused and assigned a single `segment_id`.
                - **score** -- Prediction score of segment with `segment_id`.
        """
        if label_ids_to_fuse is None:
            logger.warning('`label_ids_to_fuse` unset. No instance will be fused.')
            label_ids_to_fuse = set()
        class_queries_logits = outputs.class_queries_logits
        masks_queries_logits = outputs.masks_queries_logits
        batch_size = class_queries_logits.shape[0]
        num_labels = class_queries_logits.shape[-1] - 1
        mask_probs = masks_queries_logits.sigmoid()
        pred_scores, pred_labels = nn.functional.softmax(class_queries_logits, dim=-1).max(-1)
        results: List[Dict[str, TensorType]] = []
        for i in range(batch_size):
            mask_probs_item, pred_scores_item, pred_labels_item = remove_low_and_no_objects(mask_probs[i], pred_scores[i], pred_labels[i], threshold, num_labels)
            if mask_probs_item.shape[0] <= 0:
                height, width = target_sizes[i] if target_sizes is not None else mask_probs_item.shape[1:]
                segmentation = torch.zeros((height, width)) - 1
                results.append({'segmentation': segmentation, 'segments_info': []})
                continue
            target_size = target_sizes[i] if target_sizes is not None else None
            segmentation, segments = compute_segments(mask_probs=mask_probs_item, pred_scores=pred_scores_item, pred_labels=pred_labels_item, mask_threshold=mask_threshold, overlap_mask_area_threshold=overlap_mask_area_threshold, label_ids_to_fuse=label_ids_to_fuse, target_size=target_size)
            results.append({'segmentation': segmentation, 'segments_info': segments})
        return results

Constructs a OneFormer image processor. The image processor can be used to prepare image(s), task input(s) and
optional text inputs and targets for the model.

This image processor inherits from [`BaseImageProcessor`] which contains most of the main methods. Users should
refer to this superclass for more information regarding those methods.

Args:
    do_resize (`bool`, *optional*, defaults to `True`):
        Whether to resize the input to a certain `size`.
    size (`int`, *optional*, defaults to 800):
        Resize the input to the given size. Only has an effect if `do_resize` is set to `True`. If size is a
        sequence like `(width, height)`, output size will be matched to this. If size is an int, smaller edge of
        the image will be matched to this number. i.e, if `height > width`, then image will be rescaled to `(size *
        height / width, size)`.
    resample (`int`, *optional*, defaults to `Resampling.BILINEAR`):
        An optional resampling filter. This can be one of `PIL.Image.Resampling.NEAREST`,
        `PIL.Image.Resampling.BOX`, `PIL.Image.Resampling.BILINEAR`, `PIL.Image.Resampling.HAMMING`,
        `PIL.Image.Resampling.BICUBIC` or `PIL.Image.Resampling.LANCZOS`. Only has an effect if `do_resize` is set
        to `True`.
    do_rescale (`bool`, *optional*, defaults to `True`):
        Whether to rescale the input to a certain `scale`.
    rescale_factor (`float`, *optional*, defaults to `1/ 255`):
        Rescale the input by the given factor. Only has an effect if `do_rescale` is set to `True`.
    do_normalize (`bool`, *optional*, defaults to `True`):
        Whether or not to normalize the input with mean and standard deviation.
    image_mean (`int`, *optional*, defaults to `[0.485, 0.456, 0.406]`):
        The sequence of means for each channel, to be used when normalizing images. Defaults to the ImageNet mean.
    image_std (`int`, *optional*, defaults to `[0.229, 0.224, 0.225]`):
        The sequence of standard deviations for each channel, to be used when normalizing images. Defaults to the
        ImageNet std.
    ignore_index (`int`, *optional*):
        Label to be assigned to background pixels in segmentation maps. If provided, segmentation map pixels
        denoted with 0 (background) will be replaced with `ignore_index`.
    do_reduce_labels (`bool`, *optional*, defaults to `False`):
        Whether or not to decrement all label values of segmentation maps by 1. Usually used for datasets where 0
        is used for background, and background itself is not included in all classes of a dataset (e.g. ADE20k).
        The background label will be replaced by `ignore_index`.
    repo_path (`str`, *optional*, defaults to `"shi-labs/oneformer_demo"`):
        Path to hub repo or local directory containing the JSON file with class information for the dataset.
        If unset, will look for `class_info_file` in the current working directory.
    class_info_file (`str`, *optional*):
        JSON file containing class information for the dataset. See `shi-labs/oneformer_demo/cityscapes_panoptic.json` for an example.
    num_text (`int`, *optional*):
        Number of text entries in the text input list.
    num_labels (`int`, *optional*):
        The number of labels in the segmentation map.


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

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

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


def EfficientNetV2(width_coefficient, depth_coefficient, default_size, dropout_rate=0.2, drop_connect_rate=0.2, depth_divisor=8, min_depth=8, bn_momentum=0.9, activation='swish', blocks_args='default', name='efficientnetv2', include_top=True, weights='imagenet', input_tensor=None, input_shape=None, pooling=None, classes=1000, classifier_activation='softmax', include_preprocessing=True, weights_name=None):
    if blocks_args == 'default':
        blocks_args = DEFAULT_BLOCKS_ARGS[name]
    if not (weights in {'imagenet', None} or file_utils.exists(weights)):
        raise ValueError(f'The `weights` argument should be either `None` (random initialization), `imagenet` (pre-training on ImageNet), or the path to the weights file to be loaded.Received: weights={weights}')
    if weights == 'imagenet' and include_top and (classes != 1000):
        raise ValueError('If using `weights="imagenet"` with `include_top` as true, `classes` should be 1000')
    input_shape = imagenet_utils.obtain_input_shape(input_shape, default_size=default_size, min_size=32, data_format=backend.image_data_format(), require_flatten=include_top, weights=weights)
    if input_tensor is None:
        img_input = layers.Input(shape=input_shape)
    elif not backend.is_keras_tensor(input_tensor):
        img_input = layers.Input(tensor=input_tensor, shape=input_shape)
    else:
        img_input = input_tensor
    bn_axis = 3 if backend.image_data_format() == 'channels_last' else 1
    x = img_input
    if include_preprocessing:
        num_channels = input_shape[bn_axis - 1]
        if name.split('-')[-1].startswith('b') and num_channels == 3:
            x = layers.Rescaling(scale=1.0 / 255)(x)
            if backend.image_data_format() == 'channels_first':
                mean = [[[[0.485]], [[0.456]], [[0.406]]]]
                variance = [[[[0.229 ** 2]], [[0.224 ** 2]], [[0.225 ** 2]]]]
            else:
                mean = [0.485, 0.456, 0.406]
                variance = [0.229 ** 2, 0.224 ** 2, 0.225 ** 2]
            x = layers.Normalization(mean=mean, variance=variance, axis=bn_axis)(x)
        else:
            x = layers.Rescaling(scale=1.0 / 128.0, offset=-1)(x)
    stem_filters = round_filters(filters=blocks_args[0]['input_filters'], width_coefficient=width_coefficient, min_depth=min_depth, depth_divisor=depth_divisor)
    x = layers.Conv2D(filters=stem_filters, kernel_size=3, strides=2, kernel_initializer=CONV_KERNEL_INITIALIZER, padding='same', use_bias=False, name='stem_conv')(x)
    x = layers.BatchNormalization(axis=bn_axis, momentum=bn_momentum, name='stem_bn')(x)
    x = layers.Activation(activation, name='stem_activation')(x)
    blocks_args = copy.deepcopy(blocks_args)
    b = 0
    blocks = float(sum((args['num_repeat'] for args in blocks_args)))
    for i, args in enumerate(blocks_args):
        assert args['num_repeat'] > 0
        args['input_filters'] = round_filters(filters=args['input_filters'], width_coefficient=width_coefficient, min_depth=min_depth, depth_divisor=depth_divisor)
        args['output_filters'] = round_filters(filters=args['output_filters'], width_coefficient=width_coefficient, min_depth=min_depth, depth_divisor=depth_divisor)
        block = {0: MBConvBlock, 1: FusedMBConvBlock}[args.pop('conv_type')]
        repeats = round_repeats(repeats=args.pop('num_repeat'), depth_coefficient=depth_coefficient)
        for j in range(repeats):
            if j > 0:
                args['strides'] = 1
                args['input_filters'] = args['output_filters']
            x = block(activation=activation, bn_momentum=bn_momentum, survival_probability=drop_connect_rate * b / blocks, name=f'block{i + 1}{chr(j + 97)}_', **args)(x)
            b += 1
    top_filters = round_filters(filters=1280, width_coefficient=width_coefficient, min_depth=min_depth, depth_divisor=depth_divisor)
    x = layers.Conv2D(filters=top_filters, kernel_size=1, strides=1, kernel_initializer=CONV_KERNEL_INITIALIZER, padding='same', data_format=backend.image_data_format(), use_bias=False, name='top_conv')(x)
    x = layers.BatchNormalization(axis=bn_axis, momentum=bn_momentum, name='top_bn')(x)
    x = layers.Activation(activation=activation, name='top_activation')(x)
    if include_top:
        x = layers.GlobalAveragePooling2D(name='avg_pool')(x)
        if dropout_rate > 0:
            x = layers.Dropout(dropout_rate, name='top_dropout')(x)
        imagenet_utils.validate_activation(classifier_activation, weights)
        x = layers.Dense(classes, activation=classifier_activation, kernel_initializer=DENSE_KERNEL_INITIALIZER, bias_initializer=initializers.Constant(0.0), name='predictions')(x)
    elif pooling == 'avg':
        x = layers.GlobalAveragePooling2D(name='avg_pool')(x)
    elif pooling == 'max':
        x = layers.GlobalMaxPooling2D(name='max_pool')(x)
    if input_tensor is not None:
        inputs = operation_utils.get_source_inputs(input_tensor)
    else:
        inputs = img_input
    model = Functional(inputs, x, name=name)
    if weights == 'imagenet':
        if include_top:
            file_suffix = '.h5'
            file_hash = WEIGHTS_HASHES[weights_name][0]
        else:
            file_suffix = '_notop.h5'
            file_hash = WEIGHTS_HASHES[weights_name][1]
        file_name = name + file_suffix
        weights_path = file_utils.get_file(file_name, BASE_WEIGHTS_PATH + file_name, cache_subdir='models', file_hash=file_hash)
        model.load_weights(weights_path)
    elif weights is not None:
        model.load_weights(weights)
    return model

Instantiates the EfficientNetV2 architecture using given scaling
coefficients.

Args:
    width_coefficient: float, scaling coefficient for network width.
    depth_coefficient: float, scaling coefficient for network depth.
    default_size: integer, default input image size.
    dropout_rate: float, dropout rate before final classifier layer.
    drop_connect_rate: float, dropout rate at skip connections.
    depth_divisor: integer, a unit of network width.
    min_depth: integer, minimum number of filters.
    bn_momentum: float. Momentum parameter for Batch Normalization layers.
    activation: activation function.
    blocks_args: list of dicts, parameters to construct block modules.
    name: string, model name.
    include_top: whether to include the fully-connected layer at the top of
        the network.
    weights: one of `None` (random initialization), `"imagenet"`
        (pre-training on ImageNet),
        or the path to the weights file to be loaded.
    input_tensor: optional Keras tensor (i.e. output of `layers.Input()`) or
        numpy array to use as image input for the model.
    input_shape: optional shape tuple, only to be specified if `include_top`
        is `False`. It should have exactly 3 inputs channels.
    pooling: optional pooling mode for feature extraction when `include_top`
        is `False`.
        - `None` means that the output of the model will be the
            4D tensor output of the last convolutional layer.
        - "avg" means that global average pooling will be applied to
            the output of the last convolutional layer,
            and thus the output of the model will be a 2D tensor.
        - `"max"` means that global max pooling will be applied.
    classes: optional number of classes to classify images into,
        only to be specified if `include_top` is `True`, and if no `weights`
        argument is specified.
    classifier_activation: A string or callable. The activation function to
        use on the "top" layer. Ignored unless `include_top=True`. Set
        `classifier_activation=None` to return the logits of the "top"
        layer.
    include_preprocessing: Boolean, whether to include the preprocessing
        layer (`Rescaling`) at the bottom of the network.
        Defaults to `True`.

Returns:
    A model instance.


def convert_coco_poly_to_mask(segmentations, height: int, width: int) -> np.ndarray:
    try:
        from pycocotools import mask as coco_mask
    except ImportError:
        raise ImportError('Pycocotools is not installed in your environment.')
    masks = []
    for polygons in segmentations:
        rles = coco_mask.frPyObjects(polygons, height, width)
        mask = coco_mask.decode(rles)
        if len(mask.shape) < 3:
            mask = mask[..., None]
        mask = np.asarray(mask, dtype=np.uint8)
        mask = np.any(mask, axis=2)
        masks.append(mask)
    if masks:
        masks = np.stack(masks, axis=0)
    else:
        masks = np.zeros((0, height, width), dtype=np.uint8)
    return masks

Convert a COCO polygon annotation to a mask.

Args:
    segmentations (`List[List[float]]`):
        List of polygons, each polygon represented by a list of x-y coordinates.
    height (`int`):
        Height of the mask.
    width (`int`):
        Width of the mask.


def pack_eager_tensors(tensors, ctx=None) -> EagerTensor:
    if not isinstance(tensors, list):
        raise TypeError(f'tensors must be a list, but got a {type(tensors)}')
    if not tensors:
        raise ValueError('Cannot pack an empty list of tensors.')
    dtype = tensors[0].dtype
    shape = tensors[0].shape
    handle_data = tensors[0]._handle_data
    is_resource = dtype == dtypes.resource
    for i in range(len(tensors)):
        t = tensors[i]
        if not isinstance(t, EagerTensor):
            raise TypeError(f'All tensors being packed must be EagerTensor. Found an item of type {type(t)}.')
        if t.dtype != dtype:
            raise ValueError(f'All tensors being packed should have the same dtype {dtype}, but the {i}-th tensor is of dtype {t.dtype}')
        if t.shape != shape:
            raise ValueError(f'All tensors being packed should have the same shape {shape}, but the {i}-th tensor is of shape {t.shape}')
        if is_resource and t._handle_data != handle_data:
            raise ValueError(f'All tensors being packed should have the same handle data {handle_data}, but the {i}-th tensor is of handle data {t._handle_data}')
    if ctx is None:
        ctx = context.context()
    packed_tensor = ctx.pack_eager_tensors(tensors)
    if handle_data is not None:
        packed_tensor._handle_data = handle_data

    def grad_fun(_):
        raise ValueError('Computing gradients through pack_eager_tensors is not supported.')
    record.record_operation('pack_eager_tensors', [packed_tensor], tensors, grad_fun)
    return packed_tensor

Pack multiple `EagerTensor`s of the same dtype and shape.

Args:
  tensors: a list of EagerTensors to pack.
  ctx: context.context().

Returns:
  A packed EagerTensor.


class BinaryFocalCrossentropy(LossFunctionWrapper):

    def __init__(self, apply_class_balancing=False, alpha=0.25, gamma=2.0, from_logits=False, label_smoothing=0.0, axis=-1, reduction='sum_over_batch_size', name='binary_focal_crossentropy', dtype=None):
        super().__init__(binary_focal_crossentropy, name=name, reduction=reduction, dtype=dtype, apply_class_balancing=apply_class_balancing, alpha=alpha, gamma=gamma, from_logits=from_logits, label_smoothing=label_smoothing, axis=axis)
        self.from_logits = from_logits
        self.label_smoothing = label_smoothing
        self.axis = axis
        self.apply_class_balancing = apply_class_balancing
        self.alpha = alpha
        self.gamma = gamma

    def get_config(self):
        config = Loss.get_config(self)
        config.update({'from_logits': self.from_logits, 'label_smoothing': self.label_smoothing, 'axis': self.axis, 'apply_class_balancing': self.apply_class_balancing, 'alpha': self.alpha, 'gamma': self.gamma})
        return config

Computes focal cross-entropy loss between true labels and predictions.

Binary cross-entropy loss is often used for binary (0 or 1) classification
tasks. The loss function requires the following inputs:

- `y_true` (true label): This is either 0 or 1.
- `y_pred` (predicted value): This is the model's prediction, i.e, a single
    floating-point value which either represents a
    [logit](https://en.wikipedia.org/wiki/Logit), (i.e, value in [-inf, inf]
    when `from_logits=True`) or a probability (i.e, value in `[0., 1.]` when
    `from_logits=False`).

According to [Lin et al., 2018](https://arxiv.org/pdf/1708.02002.pdf), it
helps to apply a "focal factor" to down-weight easy examples and focus more
on hard examples. By default, the focal tensor is computed as follows:

`focal_factor = (1 - output) ** gamma` for class 1
`focal_factor = output ** gamma` for class 0
where `gamma` is a focusing parameter. When `gamma=0`, this function is
equivalent to the binary crossentropy loss.

Args:
    apply_class_balancing: A bool, whether to apply weight balancing on the
        binary classes 0 and 1.
    alpha: A weight balancing factor for class 1, default is `0.25` as
        mentioned in reference [Lin et al., 2018](
        https://arxiv.org/pdf/1708.02002.pdf).  The weight for class 0 is
        `1.0 - alpha`.
    gamma: A focusing parameter used to compute the focal factor, default is
        `2.0` as mentioned in the reference
        [Lin et al., 2018](https://arxiv.org/pdf/1708.02002.pdf).
    from_logits: Whether to interpret `y_pred` as a tensor of
        [logit](https://en.wikipedia.org/wiki/Logit) values. By default, we
        assume that `y_pred` are probabilities (i.e., values in `[0, 1]`).
    label_smoothing: Float in `[0, 1]`. When `0`, no smoothing occurs.
        When > `0`, we compute the loss between the predicted labels
        and a smoothed version of the true labels, where the smoothing
        squeezes the labels towards `0.5`.
        Larger values of `label_smoothing` correspond to heavier smoothing.
    axis: The axis along which to compute crossentropy (the features axis).
        Defaults to `-1`.
    reduction: Type of reduction to apply to the loss. In almost all cases
        this should be `"sum_over_batch_size"`. Supported options are
        `"sum"`, `"sum_over_batch_size"`, `"mean"`,
        `"mean_with_sample_weight"` or `None`. `"sum"` sums the loss,
        `"sum_over_batch_size"` and `"mean"` sum the loss and divide by the
        sample size, and `"mean_with_sample_weight"` sums the loss and
        divides by the sum of the sample weights. `"none"` and `None`
        perform no aggregation. Defaults to `"sum_over_batch_size"`.
    name: Optional name for the loss instance.
    dtype: The dtype of the loss's computations. Defaults to `None`, which
        means using `keras.backend.floatx()`. `keras.backend.floatx()` is a
        `"float32"` unless set to different value
        (via `keras.backend.set_floatx()`). If a `keras.DTypePolicy` is
        provided, then the `compute_dtype` will be utilized.

Examples:

With the `compile()` API:

```python
model.compile(
    loss=keras.losses.BinaryFocalCrossentropy(
        gamma=2.0, from_logits=True),
    ...
)
```

As a standalone function:

>>> # Example 1: (batch_size = 1, number of samples = 4)
>>> y_true = np.array([0, 1, 0, 0])
>>> y_pred = np.array([-18.6, 0.51, 2.94, -12.8])
>>> loss = keras.losses.BinaryFocalCrossentropy(
...    gamma=2, from_logits=True)
>>> loss(y_true, y_pred)
0.691

>>> # Apply class weight
>>> loss = keras.losses.BinaryFocalCrossentropy(
...     apply_class_balancing=True, gamma=2, from_logits=True)
>>> loss(y_true, y_pred)
0.51

>>> # Example 2: (batch_size = 2, number of samples = 4)
>>> y_true = np.array([[0, 1], [0, 0]])
>>> y_pred = np.array([[-18.6, 0.51], [2.94, -12.8]])
>>> # Using default 'auto'/'sum_over_batch_size' reduction type.
>>> loss = keras.losses.BinaryFocalCrossentropy(
...     gamma=3, from_logits=True)
>>> loss(y_true, y_pred)
0.647

>>> # Apply class weight
>>> loss = keras.losses.BinaryFocalCrossentropy(
...      apply_class_balancing=True, gamma=3, from_logits=True)
>>> loss(y_true, y_pred)
0.482

>>> # Using 'sample_weight' attribute with focal effect
>>> loss = keras.losses.BinaryFocalCrossentropy(
...     gamma=3, from_logits=True)
>>> loss(y_true, y_pred, sample_weight=[0.8, 0.2])
0.133

>>> # Apply class weight
>>> loss = keras.losses.BinaryFocalCrossentropy(
...      apply_class_balancing=True, gamma=3, from_logits=True)
>>> loss(y_true, y_pred, sample_weight=[0.8, 0.2])
0.097

>>> # Using 'sum' reduction` type.
>>> loss = keras.losses.BinaryFocalCrossentropy(
...     gamma=4, from_logits=True,
...     reduction="sum")
>>> loss(y_true, y_pred)
1.222

>>> # Apply class weight
>>> loss = keras.losses.BinaryFocalCrossentropy(
...     apply_class_balancing=True, gamma=4, from_logits=True,
...     reduction="sum")
>>> loss(y_true, y_pred)
0.914

>>> # Using 'none' reduction type.
>>> loss = keras.losses.BinaryFocalCrossentropy(
...     gamma=5, from_logits=True,
...     reduction=None)
>>> loss(y_true, y_pred)
array([0.0017 1.1561], dtype=float32)

>>> # Apply class weight
>>> loss = keras.losses.BinaryFocalCrossentropy(
...     apply_class_balancing=True, gamma=5, from_logits=True,
...     reduction=None)
>>> loss(y_true, y_pred)
array([0.0004 0.8670], dtype=float32)


def _get_relationships(self, dna: pg.DNA) -> Tuple[List[pg.DNA], List[Optional[pg.DNA]], List[Optional[int]]]:

    def is_mutable_node(obj):
        return self._is_mutable_node(obj)
    results = pg.query(dna, where=is_mutable_node, enter_selected=True)
    child_nodes = list(results.values())
    parent_nodes = [n.parent_dna for n in child_nodes]
    child_indexes = [n.sym_path.key if n.parent_dna else None for n in child_nodes]
    return (child_nodes, parent_nodes, child_indexes)

Extracts the parent-child node relationships in a DNA.

Note that PyGlove represents the nodes in a DNA instance as DNA instances
themselves.

Args:
  dna: the DNA that will be mutated.

Returns:
  A tuple of 3 lists of the same length with corresponding elements:
  -child_nodes: a list of every node in the DNA.
  -parent_nodes: a list of the parent node of the corresponding node in
    `child_nodes`.
  -child_indexes: a list of indexes. For all j, child_nodes[j] is the i-th
    child of parent_nodes[j], where i = child_indexes[j].
  Note that the root is included as a "child" with a `None` parent.


def batch_normalization(x, mean, var, beta, gamma, axis=-1, epsilon=0.001):
    if ndim(x) == 4:
        if axis == 1 or axis == -3:
            tf_data_format = 'NCHW'
        elif axis == 3 or axis == -1:
            tf_data_format = 'NHWC'
        else:
            tf_data_format = None
        if tf_data_format == 'NHWC' or (tf_data_format == 'NCHW' and _has_nchw_support()):
            if ndim(mean) > 1:
                mean = array_ops.reshape(mean, [-1])
            if ndim(var) > 1:
                var = array_ops.reshape(var, [-1])
            if beta is None:
                beta = zeros_like(mean)
            elif ndim(beta) > 1:
                beta = array_ops.reshape(beta, [-1])
            if gamma is None:
                gamma = ones_like(mean)
            elif ndim(gamma) > 1:
                gamma = array_ops.reshape(gamma, [-1])
        y, _, _ = nn.fused_batch_norm(x, gamma, beta, epsilon=epsilon, mean=mean, variance=var, data_format=tf_data_format, is_training=False)
        return y
    return nn.batch_normalization(x, mean, var, beta, gamma, epsilon)

Applies batch normalization on x given mean, var, beta and gamma.

I.e. returns:
`output = (x - mean) / (sqrt(var) + epsilon) * gamma + beta`

Args:
    x: Input tensor or variable.
    mean: Mean of batch.
    var: Variance of batch.
    beta: Tensor with which to center the input.
    gamma: Tensor by which to scale the input.
    axis: Integer, the axis that should be normalized.
        (typically the features axis).
    epsilon: Fuzz factor.

Returns:
    A tensor.


def _infer_graph(self, inputs, clusters):
    assert isinstance(inputs, list)
    scores = self._distance_graph(inputs, clusters, self._distance_metric)
    output = []
    if self._distance_metric == COSINE_DISTANCE and (not self._clusters_l2_normalized()):
        with ops.colocate_with(clusters, ignore_existing=True):
            clusters = nn_impl.l2_normalize(clusters, axis=1)
    for inp, score in zip(inputs, scores):
        with ops.colocate_with(inp, ignore_existing=True):
            indices, distances = gen_clustering_ops.nearest_neighbors(inp, clusters, 1)
            if self._distance_metric == COSINE_DISTANCE:
                distances *= 0.5
            output.append((score, array_ops.squeeze(distances, [-1]), array_ops.squeeze(indices, [-1])))
    return zip(*output)

Maps input to closest cluster and the score.

Args:
  inputs: list of input Tensors.
  clusters: Tensor of cluster centers.

Returns:
  List of tuple, where each value in tuple corresponds to a value in inp.
  The tuple has following three elements:
  all_scores: distance of each input to each cluster center.
  score: distance of each input to closest cluster center.
  cluster_idx: index of cluster center closest to the corresponding input.


class MMBTForClassification(nn.Module):

    def __init__(self, config, transformer, encoder):
        super().__init__()
        self.num_labels = config.num_labels
        self.mmbt = MMBTModel(config, transformer, encoder)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)

    def forward(self, input_modal, input_ids=None, modal_start_tokens=None, modal_end_tokens=None, attention_mask=None, token_type_ids=None, modal_token_type_ids=None, position_ids=None, modal_position_ids=None, head_mask=None, inputs_embeds=None, labels=None, return_dict=None):
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        outputs = self.mmbt(input_modal=input_modal, input_ids=input_ids, modal_start_tokens=modal_start_tokens, modal_end_tokens=modal_end_tokens, attention_mask=attention_mask, token_type_ids=token_type_ids, modal_token_type_ids=modal_token_type_ids, position_ids=position_ids, modal_position_ids=modal_position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, return_dict=return_dict)
        pooled_output = outputs[1]
        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)
        loss = None
        if labels is not None:
            if self.num_labels == 1:
                loss_fct = MSELoss()
                loss = loss_fct(logits.view(-1), labels.view(-1))
            else:
                loss_fct = CrossEntropyLoss()
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
        if not return_dict:
            output = (logits,) + outputs[2:]
            return (loss,) + output if loss is not None else output
        return SequenceClassifierOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions)

**labels**: (*optional*) `torch.LongTensor` of shape `(batch_size,)`:
    Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
    config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
    `config.num_labels > 1` a classification loss is computed (Cross-Entropy).

Returns: *Tuple* comprising various elements depending on the configuration (config) and inputs: **loss**:
(*optional*, returned when `labels` is provided) `torch.FloatTensor` of shape `(1,)`: Classification (or
regression if config.num_labels==1) loss. **logits**:
    `torch.FloatTensor` of shape `(batch_size, config.num_labels)` Classification (or regression if
    config.num_labels==1) scores (before SoftMax).
**hidden_states**: (*optional*, returned when `output_hidden_states=True`) list of `torch.FloatTensor` (one for
the output of each layer + the output of the embeddings) 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**:
(*optional*, returned when `output_attentions=True`) list 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.

Examples:

```python
# For example purposes. Not runnable.
transformer = BertModel.from_pretrained("google-bert/bert-base-uncased")
encoder = ImageEncoder(args)
model = MMBTForClassification(config, transformer, encoder)
outputs = model(input_modal, input_ids, labels=labels)
loss, logits = outputs[:2]
```


def _remove_lines(self, lines, sublist_lengths, num_to_remove):
    curr = 0
    result = []
    for offset in sublist_lengths:
        end = curr + offset
        start = min(curr + num_to_remove, end)
        result += lines[start:end]
        curr += offset
    return result

Utility function to remove num_to_remove lines from each sublist.

Args:
  lines: list of items.
  sublist_lengths: list of integers representing length of sublist
    corresponding to each source file.
  num_to_remove: number of lines to remove from each sublist.
Returns:
  remaining lines.


def __init__(self, start_date, end_date, coupon_spec, first_coupon_date=None, penultimate_coupon_date=None, dtype=None, name=None):
    super(FixedCashflowStream, self).__init__()
    self._name = name or 'fixed_cashflow_stream'
    with tf.name_scope(self._name):
        self._start_date = dates.convert_to_date_tensor(start_date)
        self._end_date = dates.convert_to_date_tensor(end_date)
        self._batch_size = self._start_date.shape[0]
        self._dtype = dtype
        if first_coupon_date is None:
            self._first_coupon_date = None
        else:
            self._first_coupon_date = dates.convert_to_date_tensor(first_coupon_date)
        if penultimate_coupon_date is None:
            self._penultimate_coupon_date = None
        else:
            self._penultimate_coupon_date = dates.convert_to_date_tensor(penultimate_coupon_date)
        self._setup(coupon_spec)

Initialize a batch of fixed cashflow streams.

Args:
  start_date: A rank 1 `DateTensor` specifying the starting dates of the
    accrual of the first coupon of the cashflow stream. The shape of the
    input correspond to the numbercof streams being created.
  end_date: A rank 1 `DateTensor` specifying the end dates for accrual of
    the last coupon in each cashflow stream. The shape of the input should
    be the same as that of `start_date`.
  coupon_spec: A scalar or a list of `FixedCouponSpecs` specifying the
    details of the coupon payment for the cashflow stream. If specified as
    a list then the length of the list should be the same as the number of
    streams being created and each coupon within the list must have the
    same daycount_convention and businessday_rule. If specified as
    a scalar, then the elements of the namedtuple must be of the same shape
    as (or compatible to) the shape of `start_date`.
  first_coupon_date: An optional rank 1 `DateTensor` specifying the payment
    dates of the first coupon of the cashflow stream. Use this input for
    cashflows with irregular first coupon.
    Default value: None which implies regular first coupon.
  penultimate_coupon_date: An optional rank 1 `DateTensor` specifying the
    payment date of the penultimate (next to last) coupon of the cashflow
    stream. Use this input for cashflows with irregular last coupon.
    Default value: None which implies regular last coupon.
  dtype: `tf.Dtype`. If supplied the dtype for the real variables or ops
    either supplied to the FixedCashflowStream object or created by the
    object.
    Default value: None which maps to the default dtype inferred by
    TensorFlow.
  name: Python str. The name to give to the ops created by this class.
    Default value: `None` which maps to 'fixed_cashflow_stream'.


def list_adb_devices_by_usb_id():
    out = adb.AdbProxy().devices(['-l'])
    clean_lines = str(out, 'utf-8').strip().split('\n')
    results = []
    for line in clean_lines:
        tokens = line.strip().split()
        if len(tokens) > 2 and tokens[1] == 'device':
            results.append(tokens[2])
    return results

List the usb id of all android devices connected to the computer that
are detected by adb.

Returns:
  A list of strings that are android device usb ids. Empty if there's
  none.


def pad(tensor, paddings, mode='CONSTANT', name=None, constant_values=0):
    mode = mode.upper()
    if mode == 'CONSTANT':
        if not tensor_util.is_tf_type(constant_values) and np.ndim(constant_values) == 0 and (constant_values == np.zeros_like(constant_values)):
            result = gen_array_ops.pad(tensor, paddings, name=name)
        else:
            result = gen_array_ops.pad_v2(tensor, paddings, constant_values, name=name)
    elif mode == 'REFLECT':
        result = gen_array_ops.mirror_pad(tensor, paddings, mode='REFLECT', name=name)
    elif mode == 'SYMMETRIC':
        result = gen_array_ops.mirror_pad(tensor, paddings, mode='SYMMETRIC', name=name)
    else:
        raise ValueError(f'Value of argument `mode` expected to be one of "CONSTANT", "REFLECT", or "SYMMETRIC". Received `mode` = {mode}')
    if not context.executing_eagerly():
        paddings_constant = _get_paddings_constant(paddings)
        input_shape = tensor_shape.TensorShape(tensor.shape) if isinstance(tensor, tensor_lib.Tensor) else result.op.inputs[0].shape
        if input_shape.ndims is not None and (not result.shape.is_fully_defined()) and (paddings_constant is not None):
            new_shape = []
            for padding, dim in zip(paddings_constant, input_shape.as_list()):
                if padding is None or dim is None or any((x is None for x in padding)):
                    new_shape.append(None)
                else:
                    new_shape.append(sum(padding) + dim)
            result.set_shape(new_shape)
    return result

Pads a tensor.

This operation pads a `tensor` according to the `paddings` you specify.
`paddings` is an integer tensor with shape `[n, 2]`, where n is the rank of
`tensor`. For each dimension D of `input`, `paddings[D, 0]` indicates how
many values to add before the contents of `tensor` in that dimension, and
`paddings[D, 1]` indicates how many values to add after the contents of
`tensor` in that dimension. If `mode` is "REFLECT" then both `paddings[D, 0]`
and `paddings[D, 1]` must be no greater than `tensor.dim_size(D) - 1`. If
`mode` is "SYMMETRIC" then both `paddings[D, 0]` and `paddings[D, 1]` must be
no greater than `tensor.dim_size(D)`.

The padded size of each dimension D of the output is:

`paddings[D, 0] + tensor.dim_size(D) + paddings[D, 1]`

For example:

```python
t = tf.constant([[1, 2, 3], [4, 5, 6]])
paddings = tf.constant([[1, 1,], [2, 2]])
# 'constant_values' is 0.
# rank of 't' is 2.
tf.pad(t, paddings, "CONSTANT")  # [[0, 0, 0, 0, 0, 0, 0],
                                 #  [0, 0, 1, 2, 3, 0, 0],
                                 #  [0, 0, 4, 5, 6, 0, 0],
                                 #  [0, 0, 0, 0, 0, 0, 0]]

tf.pad(t, paddings, "REFLECT")  # [[6, 5, 4, 5, 6, 5, 4],
                                #  [3, 2, 1, 2, 3, 2, 1],
                                #  [6, 5, 4, 5, 6, 5, 4],
                                #  [3, 2, 1, 2, 3, 2, 1]]

tf.pad(t, paddings, "SYMMETRIC")  # [[2, 1, 1, 2, 3, 3, 2],
                                  #  [2, 1, 1, 2, 3, 3, 2],
                                  #  [5, 4, 4, 5, 6, 6, 5],
                                  #  [5, 4, 4, 5, 6, 6, 5]]
```

Args:
  tensor: A `Tensor`.
  paddings: A `Tensor` of type `int32`.
  mode: One of "CONSTANT", "REFLECT", or "SYMMETRIC" (case-insensitive)
  name: A name for the operation (optional).
  constant_values: In "CONSTANT" mode, the scalar pad value to use. Must be
    same type as `tensor`.

Returns:
  A `Tensor`. Has the same type as `tensor`.

Raises:
  ValueError: When mode is not one of "CONSTANT", "REFLECT", or "SYMMETRIC".


def build_string_from_input(prompt, bos_token, image_seq_len, image_token, num_images):
    return f'{image_token * image_seq_len * num_images}{bos_token}{prompt}\n'

Builds a string from the input prompt and image tokens.
For example, for the call:
build_string_from_input(
    prompt="Prefix str"
    bos_token="<s>",
    image_seq_len=3,
    image_token="<im>",
)
The output will be:
"<im><im><im><s>Initial str"
Args:
    prompt (`List[Union[str, ImageInput]]`): The input prompt.
    bos_token (`str`): The beginning of sentence token.
    image_seq_len (`int`): The length of the image sequence.
    image_token (`str`): The image token.
    num_images (`int`): Number of images in the prompt.


def __init__(self, features: List[np.ndarray], timestamps: np.ndarray, schema: Optional[Schema]=None) -> None:
    self.features = features
    self.timestamps = timestamps
    if schema is not None:
        self.check_schema(schema)

Initializes the IndexData object by checking and setting the features
and timestamps.

Raises:
    ValueError: If features are not one-dimensional arrays.
    ValueError: If the number of elements in features and timestamps
        do not match.


def __init__(self, cluster_resolver=None, communication_options=None):
    if communication_options is None:
        communication_options = collective_util.Options()
    super(CollectiveAllReduceStrategy, self).__init__(CollectiveAllReduceExtended(self, cluster_resolver=cluster_resolver, communication_options=communication_options))
    distribute_lib.distribution_strategy_gauge.get_cell('V2').set('MultiWorkerMirroredStrategy')
    distribute_lib.distribution_strategy_replica_gauge.get_cell('num_workers').set(self.extended._num_workers)
    distribute_lib.distribution_strategy_replica_gauge.get_cell('num_replicas_per_worker').set(self.extended._num_devices_per_worker)

Creates the strategy.

Args:
  cluster_resolver: optional
    `tf.distribute.cluster_resolver.ClusterResolver`. If `None`,
    `tf.distribute.cluster_resolver.TFConfigClusterResolver` is used.
  communication_options: optional
    `tf.distribute.experimental.CommunicationOptions`. This configures the
    default options for cross device communications. It can be overridden by
    options provided to the communication APIs like
    `tf.distribute.ReplicaContext.all_reduce`. See
    `tf.distribute.experimental.CommunicationOptions` for details.


def add_2d_positional_embeddings(self, grid, interpolate_pos_encoding: bool=False):
    batch_size, height, width, hidden_dim = grid.shape
    row_height = min(self.max_grid_row_position_embeddings, height)
    row_position_ids = torch.arange(row_height, dtype=torch.long, device=grid.device)
    row_position_embeddings = self.row_position_embeddings(row_position_ids)
    row_shape = (1,) * (len(grid.shape) - 3) + (row_height, 1, hidden_dim)
    row_position_embeddings = row_position_embeddings.view(*row_shape)
    row_width = min(self.max_grid_col_position_embeddings, width)
    col_position_ids = torch.arange(row_width, dtype=torch.long, device=grid.device)
    col_position_embeddings = self.col_position_embeddings(col_position_ids)
    col_shape = (batch_size, 1, row_width, hidden_dim)
    col_position_embeddings = col_position_embeddings.view(*col_shape)
    positional_embeddings = row_position_embeddings + col_position_embeddings
    if interpolate_pos_encoding and (height > self.max_grid_row_position_embeddings or width > self.max_grid_col_position_embeddings):
        grid = grid + self.interpolate_pos_encoding(positional_embeddings, height, width)
    else:
        grid = grid + positional_embeddings
    return grid

Args:
    grid: (batch_size, height, width, hidden_dim)
    interpolate_pos_encoding: (`bool`, *optional*, defaults to `False`):
        Whether to interpolate the pre-trained position encodings.
Returns:
    grid + col_position_embeddings.view(*col_shape): (batch_size, *, height, width, hidden_dim)


def choose_from_datasets(datasets, choice_dataset, stop_on_empty_dataset=True) -> 'DatasetV2':
    from tensorflow.python.data.ops import choose_from_datasets_op
    return choose_from_datasets_op._choose_from_datasets(datasets, choice_dataset, stop_on_empty_dataset)

Creates a dataset that deterministically chooses elements from `datasets`.

For example, given the following datasets:

```python
datasets = [tf.data.Dataset.from_tensors("foo").repeat(),
            tf.data.Dataset.from_tensors("bar").repeat(),
            tf.data.Dataset.from_tensors("baz").repeat()]

# Define a dataset containing `[0, 1, 2, 0, 1, 2, 0, 1, 2]`.
choice_dataset = tf.data.Dataset.range(3).repeat(3)

result = tf.data.Dataset.choose_from_datasets(datasets, choice_dataset)
```

The elements of `result` will be:

```
"foo", "bar", "baz", "foo", "bar", "baz", "foo", "bar", "baz"
```

Args:
  datasets: A non-empty list of `tf.data.Dataset` objects with compatible
    structure.
  choice_dataset: A `tf.data.Dataset` of scalar `tf.int64` tensors between
    `0` and `len(datasets) - 1`.
  stop_on_empty_dataset: If `True`, selection stops if it encounters an
    empty dataset. If `False`, it skips empty datasets. It is recommended to
    set it to `True`. Otherwise, the selected elements start off as the user
    intends, but may change as input datasets become empty. This can be
    difficult to detect since the dataset starts off looking correct.
    Defaults to `True`.

Returns:
  A new `Dataset` with the transformation applied as described above.

Raises:
  TypeError: If `datasets` or `choice_dataset` has the wrong type.
  ValueError: If `datasets` is empty.


def invoke_process_element(self, sdf_invoker, output_processor, element, restriction, watermark_estimator_state, *args, **kwargs):
    assert isinstance(sdf_invoker, DoFnInvoker)

    class CheckpointState(object):

        def __init__(self):
            self.checkpointed = None
            self.residual_restriction = None
    checkpoint_state = CheckpointState()

    def initiate_checkpoint():
        with self._checkpoint_lock:
            if checkpoint_state.checkpointed:
                return
            checkpoint_state.checkpointed = object()
        split = sdf_invoker.try_split(0)
        if split:
            _, checkpoint_state.residual_restriction = split
        else:
            checkpoint_state.checkpointed = None
    output_processor.reset()
    Timer(self._max_duration, initiate_checkpoint).start()
    sdf_invoker.invoke_process(element, additional_args=args, restriction=restriction, watermark_estimator_state=watermark_estimator_state)
    assert output_processor.output_iter is not None
    output_count = 0
    process_continuation = None
    for output in output_processor.output_iter:
        assert not process_continuation
        if isinstance(output, ProcessContinuation):
            initiate_checkpoint()
            process_continuation = output
            continue
        yield output
        output_count += 1
        if self._max_num_outputs and output_count >= self._max_num_outputs:
            initiate_checkpoint()
    result = SDFProcessElementInvoker.Result(residual_restriction=checkpoint_state.residual_restriction) if checkpoint_state.residual_restriction else SDFProcessElementInvoker.Result()
    yield result

Invokes `process()` method of a Splittable `DoFn` for a given element.

Args:
  sdf_invoker: a `DoFnInvoker` for the Splittable `DoFn`.
  element: the element to process
Returns:
  a `SDFProcessElementInvoker.Result` object.


class TvltConfig(PretrainedConfig):
    model_type = 'tvlt'

    def __init__(self, image_size=224, spectrogram_length=2048, frequency_length=128, image_patch_size=[16, 16], audio_patch_size=[16, 16], num_image_channels=3, num_audio_channels=1, num_frames=8, hidden_size=768, num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072, hidden_act='gelu', hidden_dropout_prob=0.0, attention_probs_dropout_prob=0.0, initializer_range=0.02, layer_norm_eps=1e-06, qkv_bias=True, use_mean_pooling=False, decoder_num_attention_heads=16, decoder_hidden_size=512, decoder_num_hidden_layers=8, decoder_intermediate_size=2048, pixel_mask_ratio=0.75, audio_mask_ratio=0.15, audio_mask_type='frame-level', task_matching=True, task_mae=True, loss_type='classification', **kwargs):
        super().__init__(**kwargs)
        if audio_mask_type not in ('frame-level', 'patch_level'):
            raise ValueError(f"audio_mask_type must be one of two acceptable strategies - {{'frame_level', 'patch-level') got {audio_mask_type}")
        self.image_size = image_size
        self.spectrogram_length = spectrogram_length
        self.frequency_length = frequency_length
        self.image_patch_size = image_patch_size
        self.audio_patch_size = audio_patch_size
        self.num_image_channels = num_image_channels
        self.num_audio_channels = num_audio_channels
        self.num_frames = num_frames
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.intermediate_size = intermediate_size
        self.hidden_act = hidden_act
        self.hidden_dropout_prob = hidden_dropout_prob
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.initializer_range = initializer_range
        self.layer_norm_eps = layer_norm_eps
        self.qkv_bias = qkv_bias
        self.use_mean_pooling = use_mean_pooling
        self.decoder_num_attention_heads = decoder_num_attention_heads
        self.decoder_hidden_size = decoder_hidden_size
        self.decoder_num_hidden_layers = decoder_num_hidden_layers
        self.decoder_intermediate_size = decoder_intermediate_size
        self.pixel_mask_ratio = pixel_mask_ratio
        self.audio_mask_ratio = audio_mask_ratio
        self.audio_mask_type = audio_mask_type
        self.task_matching = task_matching
        self.task_mae = task_mae
        self.loss_type = loss_type

This is the configuration class to store the configuration of a [`TvltModel`]. It is used to instantiate a TVLT
model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
defaults will yield a similar configuration to that of the TVLT
[ZinengTang/tvlt-base](https://huggingface.co/ZinengTang/tvlt-base) architecture.

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

Args:
    image_size (`int`, *optional*, defaults to 224):
        The size (resolution) of each image.
    spectrogram_length (`int`, *optional*, defaults to 2048):
        The time length of each audio spectrogram.
    frequency_length (`int`, *optional*, defaults to 128):
        The frequency length of audio spectrogram.
    image_patch_size (`List[int]`, *optional*, defaults to `[16, 16]`):
        The size (resolution) of each image patch.
    audio_patch_size (`List[int]`, *optional*, defaults to `[16, 16]`):
        The size (resolution) of each audio patch.
    num_image_channels (`int`, *optional*, defaults to 3):
        The number of input image channels.
    num_audio_channels (`int`, *optional*, defaults to 1):
        The number of input audio channels.
    num_frames (`int`, *optional*, defaults to 8):
        The maximum number of frames for an input video.
    hidden_size (`int`, *optional*, defaults to 768):
        Dimensionality of the encoder layers and the pooler layer.
    num_hidden_layers (`int`, *optional*, defaults to 12):
        Number of hidden layers in the Transformer encoder.
    num_attention_heads (`int`, *optional*, defaults to 12):
        Number of attention heads for each attention layer in the Transformer encoder.
    intermediate_size (`int`, *optional*, defaults to 3072):
        Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
    hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
        The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
        `"relu"`, `"selu"` and `"gelu_new"` are supported.
    hidden_dropout_prob (`float`, *optional*, defaults to 0.0):
        The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
    attention_probs_dropout_prob (`float`, *optional*, defaults to 0.0):
        The dropout ratio for the attention probabilities.
    initializer_range (`float`, *optional*, defaults to 0.02):
        The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
    layer_norm_eps (`float`, *optional*, defaults to 1e-06):
        The epsilon used by the layer normalization layers.
    qkv_bias (`bool`, *optional*, defaults to `True`):
        Whether to add a bias to the queries, keys and values.
    use_mean_pooling (`bool`, *optional*, defaults to `False`):
        Whether to mean pool the final hidden states instead of using the final hidden state of the [CLS] token.
    decoder_num_attention_heads (`int`, *optional*, defaults to 16):
        Number of attention heads for each attention layer in the decoder.
    decoder_hidden_size (`int`, *optional*, defaults to 512):
        Dimensionality of the decoder.
    decoder_num_hidden_layers (`int`, *optional*, defaults to 8):
        Number of hidden layers in the decoder.
    decoder_intermediate_size (`int`, *optional*, defaults to 2048):
        Dimensionality of the "intermediate" (i.e., feed-forward) layer in the decoder.
    pixel_mask_ratio (`float`, *optional*, defaults to 0.75):
        Image patch masking ratio.
    audio_mask_ratio (`float`, *optional*, defaults to 0.15):
        Audio patch masking ratio.
    audio_mask_type (`str`, *optional*, defaults to `"frame-level"`):
        Audio patch masking type, choose between "frame-level" and "patch-level".
    task_matching (`bool`, *optional*, defaults to `True`):
        Whether to use vision audio matching task in pretraining.
    task_mae (`bool`, *optional*, defaults to `True`):
        Whether to use the masked auto-encoder (MAE) in pretraining.
    loss_type (`str`, *optional*, defaults to `"classification"`):
        Loss types including regression and classification.

Example:

```python
>>> from transformers import TvltConfig, TvltModel

>>> # # Initializing a TVLT ZinengTang/tvlt-base style configuration
>>> configuration = TvltConfig()

>>> # # Initializing a model (with random weights) from the ZinengTang/tvlt-base style configuration
>>> model = TvltModel(configuration)

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


def try_listify_dict_with_int_keys(src: Dict[Any, Any], convert_when_sparse: bool=False) -> Tuple[Union[List[Any], Dict[Any, Any]], bool]:
    if not src:
        return (src, False)
    min_key = None
    max_key = None
    for key in src.keys():
        if not isinstance(key, int):
            return (src, False)
        if min_key is None or min_key > key:
            min_key = key
        if max_key is None or max_key < key:
            max_key = key
    if convert_when_sparse or (min_key == 0 and max_key == len(src) - 1):
        return ([src[key] for key in sorted(src.keys())], True)
    return (src, False)

Try to convert a dictionary with consequentive integer keys to a list.

Args:
  src: A dict whose keys may be int type and their range form a perfect
    range(0, N) list unless convert_when_sparse is set to True.
  convert_when_sparse: When src is a int-key dict, force convert
    it to a list ordered by key, even it's sparse.

Returns:
  converted list or src unchanged.


def forward(self, input_ids: Optional[torch.LongTensor]=None, bbox: Optional[torch.Tensor]=None, attention_mask: Optional[torch.FloatTensor]=None, token_type_ids: Optional[torch.LongTensor]=None, position_ids: Optional[torch.LongTensor]=None, head_mask: Optional[torch.FloatTensor]=None, inputs_embeds: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    outputs = self.lilt(input_ids, bbox=bbox, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    sequence_output = outputs[0]
    logits = self.classifier(sequence_output)
    loss = None
    if labels is not None:
        labels = labels.to(logits.device)
        if self.config.problem_type is None:
            if self.num_labels == 1:
                self.config.problem_type = 'regression'
            elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
                self.config.problem_type = 'single_label_classification'
            else:
                self.config.problem_type = 'multi_label_classification'
        if self.config.problem_type == 'regression':
            loss_fct = MSELoss()
            if self.num_labels == 1:
                loss = loss_fct(logits.squeeze(), labels.squeeze())
            else:
                loss = loss_fct(logits, labels)
        elif self.config.problem_type == 'single_label_classification':
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
        elif self.config.problem_type == 'multi_label_classification':
            loss_fct = BCEWithLogitsLoss()
            loss = loss_fct(logits, labels)
    if not return_dict:
        output = (logits,) + outputs[2:]
        return (loss,) + output if loss is not None else output
    return SequenceClassifierOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions)

bbox (`torch.LongTensor` of shape `(batch_size, sequence_length, 4)`, *optional*):
    Bounding boxes of each input sequence tokens. Selected in the range `[0,
    config.max_2d_position_embeddings-1]`. Each bounding box should be a normalized version in (x0, y0, x1, y1)
    format, where (x0, y0) corresponds to the position of the upper left corner in the bounding box, and (x1,
    y1) represents the position of the lower right corner. See [Overview](#Overview) for normalization.
labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
    Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
    config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
    `config.num_labels > 1` a classification loss is computed (Cross-Entropy).

Examples:

```python
>>> from transformers import AutoTokenizer, AutoModelForSequenceClassification
>>> from datasets import load_dataset

>>> tokenizer = AutoTokenizer.from_pretrained("SCUT-DLVCLab/lilt-roberta-en-base")
>>> model = AutoModelForSequenceClassification.from_pretrained("SCUT-DLVCLab/lilt-roberta-en-base")

>>> dataset = load_dataset("nielsr/funsd-layoutlmv3", split="train", trust_remote_code=True)
>>> example = dataset[0]
>>> words = example["tokens"]
>>> boxes = example["bboxes"]

>>> encoding = tokenizer(words, boxes=boxes, return_tensors="pt")

>>> outputs = model(**encoding)
>>> predicted_class_idx = outputs.logits.argmax(-1).item()
>>> predicted_class = model.config.id2label[predicted_class_idx]
```


def extract_distribution(monitoring_info_proto):
    if not is_distribution(monitoring_info_proto):
        raise ValueError('Unsupported type %s' % monitoring_info_proto.type)
    return _decode_distribution(coders.VarIntCoder(), monitoring_info_proto.payload)

Returns a tuple of (count, sum, min, max).

Args:
  proto: The monitoring info for the distribution.


def _distance_graph(cls, inputs, clusters, distance_metric):
    assert isinstance(inputs, list)
    if distance_metric == SQUARED_EUCLIDEAN_DISTANCE:
        return cls._compute_euclidean_distance(inputs, clusters)
    elif distance_metric == COSINE_DISTANCE:
        return cls._compute_cosine_distance(inputs, clusters, inputs_normalized=True)
    else:
        assert False, str(distance_metric)

Computes distance between each input and each cluster center.

Args:
  inputs: list of input Tensors.
  clusters: cluster Tensor.
  distance_metric: distance metric used for clustering

Returns:
  list of Tensors, where each element corresponds to each element in inputs.
  The value is the distance of each row to all the cluster centers.
  Currently only Euclidean distance and cosine distance are supported.


def in_place_subclassed_model_state_restoration(model):
    assert not model._is_graph_network
    if hasattr(model, '_original_attributes_cache') and model._original_attributes_cache is not None:
        setattr_tracking = model._setattr_tracking
        model._setattr_tracking = False
        model._self_tracked_trackables = []
        for name, value in model._original_attributes_cache.items():
            setattr(model, name, value)
            if isinstance(value, Layer):
                model._self_tracked_trackables.append(value)
        model._original_attributes_cache = None
        model._setattr_tracking = setattr_tracking
    else:
        _reset_build_compile_trackers(model)

Restores the original state of a model after it was "reset".

This undoes this action of `_in_place_subclassed_model_reset`, which is called
in `clone_and_build_model` if `in_place_reset` is set to True.

Args:
  model: Instance of a Keras model created via subclassing, on which
    `_in_place_subclassed_model_reset` was previously called.


def count_up_to(self, limit):
    return gen_state_ops.resource_count_up_to(self.handle, limit=limit, T=self.dtype)

Increments this variable until it reaches `limit`.

When that Op is run it tries to increment the variable by `1`. If
incrementing the variable would bring it above `limit` then the Op raises
the exception `OutOfRangeError`.

If no error is raised, the Op outputs the value of the variable before
the increment.

This is essentially a shortcut for `count_up_to(self, limit)`.

Args:
  limit: value at which incrementing the variable raises an error.

Returns:
  A `Tensor` that will hold the variable value before the increment. If no
  other Op modifies this variable, the values produced will all be
  distinct.