def preprocess(self, videos: ImageInput, do_resize: Optional[bool]=None, size: Optional[Dict[str, int]]=None, patch_size: Optional[List[int]]=None, num_frames: Optional[int]=None, resample: PILImageResampling=None, do_center_crop: Optional[bool]=None, crop_size: Optional[Dict[str, int]]=None, do_rescale: Optional[bool]=None, rescale_factor: Optional[float]=None, do_normalize: Optional[bool]=None, image_mean: Optional[Union[float, List[float]]]=None, image_std: Optional[Union[float, List[float]]]=None, is_mixed: bool=False, return_tensors: Optional[Union[str, TensorType]]=None, data_format: ChannelDimension=ChannelDimension.FIRST, input_data_format: Optional[Union[str, ChannelDimension]]=None, **kwargs) -> BatchFeature:
    do_resize = do_resize if do_resize is not None else self.do_resize
    resample = resample if resample is not None else self.resample
    do_center_crop = do_center_crop if do_center_crop is not None else self.do_center_crop
    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
    size = size if size is not None else self.size
    size = get_size_dict(size, default_to_square=False)
    crop_size = crop_size if crop_size is not None else self.crop_size
    crop_size = get_size_dict(crop_size, param_name='crop_size')
    patch_size = patch_size if patch_size is not None else self.patch_size
    num_frames = num_frames if patch_size is not None else self.num_frames
    validate_kwargs(captured_kwargs=kwargs.keys(), valid_processor_keys=self._valid_processor_keys)
    if not valid_images(videos):
        raise ValueError('Invalid image or video type. Must be of type PIL.Image.Image, numpy.ndarray, torch.Tensor, tf.Tensor or jax.ndarray.')
    videos = make_batched(videos)
    for video in videos:
        if len(video) > self.num_frames:
            raise ValueError(f'number of frames must not be greater than the maximum frames of the model {self.num_frames}.')
    max_num_frames = max([len(video) for video in videos])
    num_patches_per_image = (size['shortest_edge'] // patch_size[0]) ** 2
    video_masks = np.array([len(video) * num_patches_per_image * [1] + (max_num_frames - len(video)) * num_patches_per_image * [0] for video in videos])
    videos = [[self._preprocess_image(image=img, do_resize=do_resize, size=size, resample=resample, do_center_crop=do_center_crop, crop_size=crop_size, 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 img in video] for video in videos]
    if is_mixed:
        data = {'pixel_values_mixed': videos, 'pixel_mask_mixed': video_masks}
    else:
        data = {'pixel_values': videos, 'pixel_mask': video_masks}
    return BatchFeature(data=data, tensor_type=return_tensors)

Preprocess an videos or image or batch of videos or images.

Args:
    videos (`ImageInput`):
        Images or videos to preprocess. Expects a single or batch of frames with pixel values ranging from 0 to
        255. If passing in frames with pixel values between 0 and 1, set `do_rescale=False`.
    do_resize (`bool`, *optional*, defaults to `self.do_resize`):
        Whether to resize the image.
    size (`Dict[str, int]`, *optional*, defaults to `self.size`):
        Size of the image after applying resize.
    patch_size (`List[int]` *optional*, defaults to self.patch_size):
        The patch size of image patch embedding.
    num_frames (`int` *optional*, defaults to self.num_frames):
        The maximum number of video frames.
    resample (`PILImageResampling`, *optional*, defaults to `self.resample`):
        Resampling filter to use if resizing the image. This can be one of the enum `PILImageResampling`, Only
        has an effect if `do_resize` is set to `True`.
    do_center_crop (`bool`, *optional*, defaults to `self.do_centre_crop`):
        Whether to centre crop the image.
    crop_size (`Dict[str, int]`, *optional*, defaults to `self.crop_size`):
        Size of the image after applying the centre crop.
    do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
        Whether to rescale the image values between [0 - 1].
    rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
        Rescale factor to rescale the image by if `do_rescale` is set to `True`.
    do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
        Whether to normalize the image.
    image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
        Image mean.
    image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
        Image standard deviation.
    is_mixed (`bool`, *optional*):
        If the input video has negative samples.
    return_tensors (`str` or `TensorType`, *optional*):
        The type of tensors to return. Can be one of:
            - Unset: Return a list of `np.ndarray`.
            - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
            - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
            - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
            - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
    data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
        The channel dimension format for the output image. Can be one of:
            - `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
            - `ChannelDimension.LAST`: image in (height, width, num_channels) format.
            - Unset: Use the inferred channel dimension format of the input image.
    input_data_format (`ChannelDimension` or `str`, *optional*):
        The channel dimension format for the input image. If unset, the channel dimension format is inferred
        from the input image. Can be one of:
        - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
        - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
        - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.

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

    - **pixel_values** -- Pixel values to be fed to a model, of shape (batch_size, num_channels, height,
      width).

    - **pixel_mask** -- Pixel masks to be fed to a model, of shape (batch_size, num_pixel_patches).

    - **pixel_values_mixed** -- Pixel values with both positive or negative to be fed to a model, of shape
      (batch_size, num_channels, height, width).

    - **pixel_mask_mixed** -- Pixel masks with both positive or negative to be fed to a model, of shape
      (batch_size, num_pixel_patches).


def tick_calendar(self: EventSetOrNode, second: Optional[Union[int, Literal['*']]]=None, minute: Optional[Union[int, Literal['*']]]=None, hour: Optional[Union[int, Literal['*']]]=None, mday: Optional[Union[int, Literal['*']]]=None, month: Optional[Union[int, Literal['*']]]=None, wday: Optional[Union[int, Literal['*']]]=None, after_last: bool=True, before_first: bool=False) -> EventSetOrNode:
    from temporian.core.operators.tick_calendar import tick_calendar
    return tick_calendar(self, second=second, minute=minute, hour=hour, mday=mday, month=month, wday=wday, after_last=after_last, before_first=before_first)

Generates events periodically at fixed times or dates e.g. each month.

Events are generated in the range of the input
[`EventSet`][temporian.EventSet] independently for each index.

The usability is inspired in the crontab format, where arguments can
take a value of `'*'` to tick at all values, or a fixed integer to
tick only at that precise value.

Non-specified values (`None`), are set to `'*'` if a finer
resolution argument is specified, or fixed to the first valid value if
a lower resolution is specified. For example, setting only
`tick_calendar(hour='*')`
is equivalent to:
`tick_calendar(second=0, minute=0, hour='*', mday='*', month='*')`
, resulting in one tick at every exact hour of every day/month/year in
the input guide range.

The datetime timezone is always assumed to be UTC.

Examples:
    ```python
    >>> # Every day (at 00:00:00) in the period (exactly one year)
    >>> a = tp.event_set(timestamps=["2021-01-01", "2021-12-31 23:59:59"])
    >>> b = a.tick_calendar(hour=0)
    >>> b
    indexes: ...
    events:
        (366 events):
            timestamps: [...]
    ...

    >>> # Every day at 2:30am
    >>> b = a.tick_calendar(hour=2, minute=30)
    >>> tp.glue(b.calendar_hour(), b.calendar_minute())
    indexes: ...
    events:
        (366 events):
            timestamps: [...]
            'calendar_hour': [2 2 2 ... 2 2 2]
            'calendar_minute': [30 30 30 ... 30 30 30]
    ...

    >>> # Day 5 of every month (at 00:00)
    >>> b = a.tick_calendar(mday=5)
    >>> b.calendar_day_of_month()
    indexes: ...
    events:
        (13 events):
            timestamps: [...]
            'calendar_day_of_month': [5 5 5 ... 5 5 5]
    ...

    >>> # 1st of February of every year
    >>> a = tp.event_set(timestamps=["2020-01-01", "2021-12-31"])
    >>> b = a.tick_calendar(month=2)
    >>> tp.glue(b.calendar_day_of_month(), b.calendar_month())
    indexes: ...
    events:
        (3 events):
            timestamps: [...]
            'calendar_day_of_month': [1 1 1]
            'calendar_month': [2 2 2]
    ...

    >>> # Every second in the period  (2 hours -> 7200 seconds)
    >>> a = tp.event_set(timestamps=["2020-01-01 00:00:00",
    ...                              "2020-01-01 01:59:59"])
    >>> b = a.tick_calendar(second='*')
    >>> b
    indexes: ...
    events:
        (7200 events):
            timestamps: [...]
    ...

    >>> # Every second of the minute 30 of every hour (00:30 and 01:30)
    >>> a = tp.event_set(timestamps=["2020-01-01 00:00",
    ...                              "2020-01-01 02:00"])
    >>> b = a.tick_calendar(second='*', minute=30)
    >>> b
    indexes: ...
    events:
        (121 events):
            timestamps: [...]
    ...

    >>> # Not allowed: intermediate arguments (minute, hour) not specified
    >>> b = a.tick_calendar(second=1, mday=1)  # ambiguous meaning
    Traceback (most recent call last):
        ...
    ValueError: Can't set argument to None because previous and
    following arguments were specified. Set to '*' or an integer ...

    >>> # not after_last
    >>> a = tp.event_set(timestamps=["2020-02-01", "2020-04-01"])
    >>> b = a.tick_calendar(mday=10, after_last=False)
    >>> tp.glue(b.calendar_day_of_month(), b.calendar_month())
    indexes: ...
    events:
        (2 events):
            timestamps: [...]
            'calendar_day_of_month': [10 10]
            'calendar_month': [2 3]
    ...

    >>> # before_first
    >>> a = tp.event_set(timestamps=["2020-02-01", "2020-04-01"])
    >>> b = a.tick_calendar(mday=10, before_first=True)
    >>> tp.glue(b.calendar_day_of_month(), b.calendar_month())
    indexes: ...
    events:
        (4 events):
            timestamps: [...]
            'calendar_day_of_month': [10 10 10 10]
            'calendar_month': [1 2 3 4]
    ...

    ```

Args:
    second: '*' (any second), None (auto) or number in range `[0-59]`
            to tick at specific second of each minute.
    minute: '*' (any minute), None (auto) or number in range `[0-59]`
            to tick at specific minute of each hour.
    hour: '*' (any hour), None (auto), or number in range `[0-23]` to
            tick at specific hour of each day.
    mday: '*' (any day), None (auto) or number in range `[1-31]`
                to tick at specific day of each month. Note that months
                without some particular day may not have any tick
                (e.g: day 31 on February).
    month: '*' (any month), None (auto) or number in range `[1-12]` to
            tick at one particular month of each year.
    wday: '*' (any day), None (auto) or number in range `[0-6]`
            (Sun-Sat) to tick at particular day of week. Can only be
            specified if `day_of_month` is `None`.
    after_last: If True, a tick after the last timestamp is included.
            Useful for window operations where you want the timestamps
            to be included in the range of the ticks.
    before_first: If True, a tick before the first timestamp is
            included.
            Useful for window operations where you want the timestamps
            to be included in the range of the ticks.

Returns:
    A feature-less EventSet with timestamps at specified interval.


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

DonutSwin outputs for image classification.

Args:
    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
        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 (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
        shape `(batch_size, sequence_length, hidden_size)`.

        Hidden-states of the model at the output of each layer plus the initial embedding outputs.
    attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
        Tuple of `torch.FloatTensor` (one for each stage) 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.
    reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
        shape `(batch_size, hidden_size, height, width)`.

        Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
        include the spatial dimensions.


class DataCollatorForSeq2Seq:
    tokenizer: PreTrainedTokenizerBase
    model: Optional[Any] = None
    padding: Union[bool, str, PaddingStrategy] = True
    max_length: Optional[int] = None
    pad_to_multiple_of: Optional[int] = None
    label_pad_token_id: int = -100
    return_tensors: str = 'pt'

    def __call__(self, features, return_tensors=None):
        if return_tensors is None:
            return_tensors = self.return_tensors
        label_name = 'label' if 'label' in features[0].keys() else 'labels'
        labels = [feature[label_name] for feature in features] if label_name in features[0].keys() else None
        if labels is not None and all((label is None for label in labels)):
            labels = None
        non_labels_features = [{k: v for k, v in feature.items() if k != label_name} for feature in features]
        batch = pad_without_fast_tokenizer_warning(self.tokenizer, non_labels_features, padding=self.padding, max_length=self.max_length, pad_to_multiple_of=self.pad_to_multiple_of, return_tensors=return_tensors)
        no_padding = self.padding is False or self.padding == PaddingStrategy.DO_NOT_PAD
        if labels is not None:
            if no_padding:
                if isinstance(features[0][label_name], list):
                    batch['labels'] = list(labels)
                else:
                    batch['labels'] = [np.concatenate([label, []]) for label in labels]
            else:
                max_padding = self.padding == PaddingStrategy.MAX_LENGTH and self.max_length is not None
                max_label_length = max((len(l) for l in labels)) if not max_padding else self.max_length
                if self.pad_to_multiple_of is not None:
                    max_label_length = (max_label_length + self.pad_to_multiple_of - 1) // self.pad_to_multiple_of * self.pad_to_multiple_of
                padding_side = self.tokenizer.padding_side
                if isinstance(features[0][label_name], list):
                    batch['labels'] = [label + [self.label_pad_token_id] * (max_label_length - len(label)) if padding_side == 'right' else [self.label_pad_token_id] * (max_label_length - len(label)) + label for label in labels]
                else:
                    batch['labels'] = [np.concatenate([label, np.array([self.label_pad_token_id] * (max_label_length - len(label)), dtype=np.int64)]) if padding_side == 'right' else np.concatenate([np.array([self.label_pad_token_id] * (max_label_length - len(label)), dtype=np.int64), label]) for label in labels]
        if batch.get('labels', None) is not None:
            if return_tensors == 'pt':
                import torch
                batch['labels'] = torch.tensor(batch['labels'], dtype=torch.int64)
            elif return_tensors == 'tf':
                import tensorflow as tf
                batch['labels'] = tf.constant(batch['labels'], dtype=tf.int64)
            else:
                batch['labels'] = np.array(batch['labels'], dtype=np.int64)
        else:
            batch['labels'] = None
        if labels is not None and self.model is not None and hasattr(self.model, 'prepare_decoder_input_ids_from_labels'):
            decoder_input_ids = self.model.prepare_decoder_input_ids_from_labels(labels=batch['labels'])
            batch['decoder_input_ids'] = decoder_input_ids
        return batch

Data collator that will dynamically pad the inputs received, as well as the labels.

Args:
    tokenizer ([`PreTrainedTokenizer`] or [`PreTrainedTokenizerFast`]):
        The tokenizer used for encoding the data.
    model ([`PreTrainedModel`], *optional*):
        The model that is being trained. If set and has the *prepare_decoder_input_ids_from_labels*, use it to
        prepare the *decoder_input_ids*

        This is useful when using *label_smoothing* to avoid calculating loss twice.
    padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`):
        Select a strategy to pad the returned sequences (according to the model's padding side and padding index)
        among:

        - `True` or `'longest'` (default): Pad to the longest sequence in the batch (or no padding if only a single
          sequence is provided).
        - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
          acceptable input length for the model if that argument is not provided.
        - `False` or `'do_not_pad'`: No padding (i.e., can output a batch with sequences of different lengths).
    max_length (`int`, *optional*):
        Maximum length of the returned list and optionally padding length (see above).
    pad_to_multiple_of (`int`, *optional*):
        If set will pad the sequence to a multiple of the provided value.

        This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=
        7.0 (Volta).
    label_pad_token_id (`int`, *optional*, defaults to -100):
        The id to use when padding the labels (-100 will be automatically ignored by PyTorch loss functions).
    return_tensors (`str`, *optional*, defaults to `"pt"`):
        The type of Tensor to return. Allowable values are "np", "pt" and "tf".


def recipe_cm360_to_dv360(config, auth_dv, auth_cm, auth_sheet, auth_bigquery, recipe_name, recipe_slug, command):
    dataset(config, {'__comment__': 'Ensure dataset exists.', 'auth': auth_bigquery, 'dataset': recipe_slug})
    drive(config, {'__comment__': 'Copy the default template to sheet with the recipe name', 'auth': auth_sheet, 'copy': {'source': 'https://docs.google.com/spreadsheets/d/1XjEHq-nEFMW8RVmCNJ-TVGvVcVBEADzjbvhmAvF04iQ/edit#gid=594912061', 'destination': recipe_name}})
    cm_to_dv(config, {'__comment': 'Depending on users choice, execute a different part of the solution.', 'auth_dv': auth_dv, 'auth_cm': auth_cm, 'auth_sheets': auth_sheet, 'auth_bigquery': auth_bigquery, 'sheet': recipe_name, 'dataset': recipe_slug, 'command': command})

Allows bulk creating DV360 Insertion Orders and Line Items from CM360.

Args:
  auth_dv (authentication) - Credentials used for dv.
  auth_cm (authentication) - Credentials used for dv.
  auth_sheet (authentication) - Credentials used for sheet.
  auth_bigquery (authentication) - Credentials used for bigquery.
  recipe_name (string) - Name of Google Sheet to create.
  recipe_slug (string) - Name of Google BigQuery dataset to create.
  command (choice) - Action to take.


def _source_file_paths_outside_tensorflow_py_library(code_defs, id_to_string):
    file_ids = set()
    for code_def in code_defs:
        for trace in code_def.traces:
            file_ids.add(trace.file_id)
    non_tf_files = (id_to_string[file_id] for file_id in file_ids)
    non_tf_files = (f for f in non_tf_files if not source_utils.guess_is_tensorflow_py_library(f) and gfile.Exists(f))
    return non_tf_files

Extract source file paths outside TensorFlow Python library.

Args:
  code_defs: An iterable of `CodeDef` protos, i.e., an iterable of stack
    traces.
  id_to_string: A proto map from integer ids to strings.

Returns:
  An iterable of source file paths outside the TensorFlow Python library.


def Count3(pcoll, factor=1):
    return pcoll | 'PairWithOne' >> beam.Map(lambda v: (v, factor)) | beam.CombinePerKey(sum)

Count as a decorated function with a side input.

Args:
  pcoll: the PCollection passed in from the previous transform
  factor: the amount by which to count

Returns:
  A PCollection counting the number of times each unique element occurs.


def _assertOpOutputMatchesExpected(self, op, axis, output_type, op_input, expected):
    with self.session() as session:
        with self.test_scope():
            pinp = array_ops.placeholder(dtypes.as_dtype(op_input.dtype), op_input.shape, name='a')
            output = op(pinp, axis=axis, output_type=output_type)
        result = session.run(output, {pinp: op_input})
        self.assertAllEqual(result, expected)

Verifies that 'op' produces 'expected' when fed input 'op_input' .

Args:
  op: argmin or argmax operator to test.
  axis: integer axis to reduce across.
  output_type: numpy datatype of the output to produce.
  op_input: numpy input array to use as input to 'op'.
  expected: numpy array representing the expected output of 'op'.


def call(self, input_values: tf.Tensor, attention_mask: tf.Tensor | None=None, token_type_ids: tf.Tensor | None=None, position_ids: tf.Tensor | None=None, head_mask: tf.Tensor | None=None, inputs_embeds: tf.Tensor | None=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, training: bool=False) -> Union[TFBaseModelOutput, Tuple[tf.Tensor]]:
    output_hidden_states = output_hidden_states if output_hidden_states else self.config.output_hidden_states
    output_attentions = output_attentions if output_attentions else self.config.output_attentions
    return_dict = return_dict if return_dict else self.config.return_dict
    outputs = self.hubert(input_values=input_values, 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, training=training)
    return outputs

Returns:

Example:

```python
>>> from transformers import AutoProcessor, TFHubertModel
>>> from datasets import load_dataset
>>> import soundfile as sf

>>> processor = AutoProcessor.from_pretrained("facebook/hubert-large-ls960-ft")
>>> model = TFHubertModel.from_pretrained("facebook/hubert-large-ls960-ft")

>>> def map_to_array(batch):
...     speech, _ = sf.read(batch["file"])
...     batch["speech"] = speech
...     return batch

>>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
>>> ds = ds.map(map_to_array)

>>> input_values = processor(ds["speech"][0], return_tensors="tf").input_values  # Batch size 1
>>> hidden_states = model(input_values).last_hidden_state
```


class DebugUnderflowOverflow:

    def __init__(self, model, max_frames_to_save=21, trace_batch_nums=[], abort_after_batch_num=None):
        self.model = model
        self.trace_batch_nums = trace_batch_nums
        self.abort_after_batch_num = abort_after_batch_num
        self.frames = collections.deque([], max_frames_to_save)
        self.frame = []
        self.batch_number = 0
        self.total_calls = 0
        self.detected_overflow = False
        self.prefix = '                 '
        self.analyse_model()
        self.register_forward_hook()

    def save_frame(self, frame=None):
        if frame is not None:
            self.expand_frame(frame)
        self.frames.append('\n'.join(self.frame))
        self.frame = []

    def expand_frame(self, line):
        self.frame.append(line)

    def trace_frames(self):
        print('\n'.join(self.frames))
        self.frames = []

    def reset_saved_frames(self):
        self.frames = []

    def dump_saved_frames(self):
        print(f'\nDetected inf/nan during batch_number={self.batch_number}')
        print(f'Last {len(self.frames)} forward frames:')
        print(f'{'abs min':8} {'abs max':8} metadata')
        print('\n'.join(self.frames))
        print('\n\n')
        self.frames = []

    def analyse_model(self):
        self.module_names = {m: name for name, m in self.model.named_modules()}

    def analyse_variable(self, var, ctx):
        if torch.is_tensor(var):
            self.expand_frame(get_abs_min_max(var, ctx))
            if detect_overflow(var, ctx):
                self.detected_overflow = True
        elif var is None:
            self.expand_frame(f'{'None':>17} {ctx}')
        else:
            self.expand_frame(f'{'not a tensor':>17} {ctx}')

    def batch_start_frame(self):
        self.expand_frame(f'\n\n{self.prefix} *** Starting batch number={self.batch_number} ***')
        self.expand_frame(f'{'abs min':8} {'abs max':8} metadata')

    def batch_end_frame(self):
        self.expand_frame(f'{self.prefix} *** Finished batch number={self.batch_number - 1} ***\n\n')

    def create_frame(self, module, input, output):
        self.expand_frame(f'{self.prefix} {self.module_names[module]} {module.__class__.__name__}')
        for name, p in module.named_parameters(recurse=False):
            self.analyse_variable(p, name)
        if isinstance(input, tuple):
            for i, x in enumerate(input):
                self.analyse_variable(x, f'input[{i}]')
        else:
            self.analyse_variable(input, 'input')
        if isinstance(output, tuple):
            for i, x in enumerate(output):
                if isinstance(x, tuple):
                    for j, y in enumerate(x):
                        self.analyse_variable(y, f'output[{i}][{j}]')
                else:
                    self.analyse_variable(x, f'output[{i}]')
        else:
            self.analyse_variable(output, 'output')
        self.save_frame()

    def register_forward_hook(self):
        self.model.apply(self._register_forward_hook)

    def _register_forward_hook(self, module):
        module.register_forward_hook(self.forward_hook)

    def forward_hook(self, module, input, output):
        last_frame_of_batch = False
        trace_mode = True if self.batch_number in self.trace_batch_nums else False
        if trace_mode:
            self.reset_saved_frames()
        if self.total_calls == 0:
            self.batch_start_frame()
        self.total_calls += 1
        if module == self.model:
            self.batch_number += 1
            last_frame_of_batch = True
        self.create_frame(module, input, output)
        if trace_mode:
            self.trace_frames()
        if last_frame_of_batch:
            self.batch_start_frame()
        if self.detected_overflow and (not trace_mode):
            self.dump_saved_frames()
            raise ValueError('DebugUnderflowOverflow: inf/nan detected, aborting as there is no point running further. Please scroll up above this traceback to see the activation values prior to this event.')
        if self.abort_after_batch_num is not None and self.batch_number > self.abort_after_batch_num:
            raise ValueError(f'DebugUnderflowOverflow: aborting after {self.batch_number} batches due to `abort_after_batch_num={self.abort_after_batch_num}` arg')

This debug class helps detect and understand where the model starts getting very large or very small, and more
importantly `nan` or `inf` weight and activation elements.

There are 2 working modes:

1. Underflow/overflow detection (default)
2. Specific batch absolute min/max tracing without detection

Mode 1: Underflow/overflow detection

To activate the underflow/overflow detection, initialize the object with the model :

```python
debug_overflow = DebugUnderflowOverflow(model)
```

then run the training as normal and if `nan` or `inf` gets detected in at least one of the weight, input or output
elements this module will throw an exception and will print `max_frames_to_save` frames that lead to this event,
each frame reporting

1. the fully qualified module name plus the class name whose `forward` was run
2. the absolute min and max value of all elements for each module weights, and the inputs and output

For example, here is the header and the last few frames in detection report for `google/mt5-small` run in fp16
mixed precision :

```
Detected inf/nan during batch_number=0
Last 21 forward frames:
abs min  abs max  metadata
[...]
                  encoder.block.2.layer.1.DenseReluDense.wi_0 Linear
2.17e-07 4.50e+00 weight
1.79e-06 4.65e+00 input[0]
2.68e-06 3.70e+01 output
                  encoder.block.2.layer.1.DenseReluDense.wi_1 Linear
8.08e-07 2.66e+01 weight
1.79e-06 4.65e+00 input[0]
1.27e-04 2.37e+02 output
                  encoder.block.2.layer.1.DenseReluDense.wo Linear
1.01e-06 6.44e+00 weight
0.00e+00 9.74e+03 input[0]
3.18e-04 6.27e+04 output
                  encoder.block.2.layer.1.DenseReluDense T5DenseGatedGeluDense
1.79e-06 4.65e+00 input[0]
3.18e-04 6.27e+04 output
                  encoder.block.2.layer.1.dropout Dropout
3.18e-04 6.27e+04 input[0]
0.00e+00      inf output
```

You can see here, that `T5DenseGatedGeluDense.forward` resulted in output activations, whose absolute max value was
around 62.7K, which is very close to fp16's top limit of 64K. In the next frame we have `Dropout` which
renormalizes the weights, after it zeroed some of the elements, which pushes the absolute max value to more than
64K, and we get an overflow.

As you can see it's the previous frames that we need to look into when the numbers start going into very large for
fp16 numbers.

The tracking is done in a forward hook, which gets invoked immediately after `forward` has completed.

By default the last 21 frames are printed. You can change the default to adjust for your needs. For example :

```python
debug_overflow = DebugUnderflowOverflow(model, max_frames_to_save=100)
```

    To validate that you have set up this debugging feature correctly, and you intend to use it in a training that
    may take hours to complete, first run it with normal tracing enabled for one of a few batches as explained in
    the next section.

    Mode 2. Specific batch absolute min/max tracing without detection

    The second work mode is per-batch tracing with the underflow/overflow detection feature turned off.

    Let's say you want to watch the absolute min and max values for all the ingredients of each `forward` call of a
given batch, and only do that for batches 1 and 3. Then you instantiate this class as :

```python
debug_overflow = DebugUnderflowOverflow(model, trace_batch_nums=[1, 3])
```

And now full batches 1 and 3 will be traced using the same format as explained above. Batches are 0-indexed.

This is helpful if you know that the program starts misbehaving after a certain batch number, so you can
fast-forward right to that area.

Early stopping:

You can also specify the batch number after which to stop the training, with :

```python
debug_overflow = DebugUnderflowOverflow(model, trace_batch_nums=[1, 3], abort_after_batch_num=3)
```

This feature is mainly useful in the tracing mode, but you can use it for any mode.

**Performance**:

As this module measures absolute `min`/``max` of each weight of the model on every forward it'll slow the training
down. Therefore remember to turn it off once the debugging needs have been met.

Args:
    model (`nn.Module`):
        The model to debug.
    max_frames_to_save (`int`, *optional*, defaults to 21):
        How many frames back to record
    trace_batch_nums(`List[int]`, *optional*, defaults to `[]`):
        Which batch numbers to trace (turns detection off)
    abort_after_batch_num  (`int``, *optional*):
        Whether to abort after a certain batch number has finished


def join(self) -> None:
    self._is_thread_joined = True
    self._thread.join()
    if self._exception is not None:
        self._testcase.fail('Error in checkedThread: %s' % str(self._exception))

Blocks until the thread terminates.

Raises:
  self._testcase.failureException: If the thread terminates with due to
    an exception.


def parse_args(test: typing.Optional[typing.List[str]]=None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument('train_data', help='File path for the encoded training data.')
    parser.add_argument('base_model', help='File path for the base model file.')
    parser.add_argument('-o', '--output', help=f'File path for the output weights. (default: {DEFAULT_OUTPUT_NAME})', type=str, default=DEFAULT_OUTPUT_NAME)
    parser.add_argument('--val-data', help='File path for the encoded validation data.', type=str)
    parser.add_argument('--iters', help=f'Number of iterations for training. (default: {DEFAULT_NUM_ITERS})', type=int, default=DEFAULT_NUM_ITERS)
    parser.add_argument('--log-span', help=f'Iteration span to print metrics. (default: {DEFAULT_LOG_SPAN})', type=int, default=DEFAULT_LOG_SPAN)
    parser.add_argument('--learning-rate', help=f'Learning rate. (default: {DEFAULT_LEARNING_RATE})', type=float, default=DEFAULT_LEARNING_RATE)
    if test is None:
        return parser.parse_args()
    else:
        return parser.parse_args(test)

Parses commandline arguments.

Args:
  test (typing.Optional[typing.List[str]], optional): Commandline args for
    testing. Defaults to None.

Returns:
  Parsed arguments (argparse.Namespace).


def compute_stats(array):
    q1 = np.percentile(array, 25)
    q3 = np.percentile(array, 75)
    low = q1 - 1.5 * (q3 - q1)
    high = q3 + 1.5 * (q3 - q1)
    filtered_array = list(filter(lambda x: low <= x and x <= high, array))
    mean = np.mean(filtered_array)
    min_val = np.min(filtered_array)
    max_val = np.max(filtered_array)
    max_diff = max(max_val - mean, mean - min_val)
    diff = max_diff / mean * 100.0
    return (mean, diff)

Reports mean and ± range for the given array.

The range computation follows benchstat's.

Args:
  array: The array to compute stats for.

Returns:
  mean and ± %diff range.


def correct_tables(self, generation: str) -> str:
    for l in generation.split('\n'):
        if l.count('\\begin{tabular}') > 15 or l.count('\\multicolumn') > 60 or l.count('&') > 400:
            generation = generation.replace(l, '')
    generation = generation.replace('\\begin{table} \\begin{tabular}', '\\begin{table}\n\\begin{tabular}')
    generation = generation.replace('\\end{tabular} \\end{table}', '\\end{tabular}\n\\end{table}')
    generation = generation.replace('\\end{table} Tab', '\\end{table}\nTab')
    generation = re.sub('(^.+)\\\\begin{tab', '\\1\\n\\\\begin{tab', generation, flags=re.M)
    generation = generation.replace('\\begin{tabular}{l l}  & \\\\ \\end{tabular}', '')
    generation = generation.replace('\\begin{tabular}{}\n\n\\end{tabular}', '')
    return generation

        Takes a generated string and fixes tables/tabulars to make them match the markdown format needed.

        Args:
            generation (str): The generated text to be postprocessed.

        Returns:
            str: The postprocessed text.

        Example:

        ```python
        correct_tables("\begin{table} \begin{tabular}{l l} & \ \end{tabular} \end{table}")
        "\begin{table}
\begin{tabular}{l l} & \ \end{tabular}
\end{table}"
        ```


class Conv1D(BaseConv):

    def __init__(self, filters, kernel_size, strides=1, padding='valid', data_format=None, dilation_rate=1, groups=1, activation=None, use_bias=True, kernel_initializer='glorot_uniform', bias_initializer='zeros', kernel_regularizer=None, bias_regularizer=None, activity_regularizer=None, kernel_constraint=None, bias_constraint=None, **kwargs):
        super().__init__(rank=1, filters=filters, kernel_size=kernel_size, strides=strides, padding=padding, data_format=data_format, dilation_rate=dilation_rate, groups=groups, activation=activation, use_bias=use_bias, kernel_initializer=kernel_initializer, bias_initializer=bias_initializer, kernel_regularizer=kernel_regularizer, bias_regularizer=bias_regularizer, activity_regularizer=activity_regularizer, kernel_constraint=kernel_constraint, bias_constraint=bias_constraint, **kwargs)

    def _compute_causal_padding(self):
        left_pad = self.dilation_rate[0] * (self.kernel_size[0] - 1)
        if self.data_format == 'channels_last':
            causal_padding = [[0, 0], [left_pad, 0], [0, 0]]
        else:
            causal_padding = [[0, 0], [0, 0], [left_pad, 0]]
        return causal_padding

    def call(self, inputs):
        padding = self.padding
        if self.padding == 'causal':
            inputs = ops.pad(inputs, self._compute_causal_padding())
            padding = 'valid'
        outputs = ops.conv(inputs, self.kernel, strides=list(self.strides), padding=padding, dilation_rate=self.dilation_rate, data_format=self.data_format)
        if self.use_bias:
            if self.data_format == 'channels_last':
                bias_shape = (1,) * (self.rank + 1) + (self.filters,)
            else:
                bias_shape = (1, self.filters) + (1,) * self.rank
            bias = ops.reshape(self.bias, bias_shape)
            outputs = ops.add(outputs, bias)
        if self.activation is not None:
            return self.activation(outputs)
        return outputs

1D convolution layer (e.g. temporal convolution).

This layer creates a convolution kernel that is convolved with the layer
input over a single spatial (or temporal) dimension to produce a tensor of
outputs. If `use_bias` is True, a bias vector is created and added to the
outputs. Finally, if `activation` is not `None`, it is applied to the
outputs as well.

Args:
    filters: int, the dimension of the output space (the number of filters
        in the convolution).
    kernel_size: int or tuple/list of 1 integer, specifying the size of the
        convolution window.
    strides: int or tuple/list of 1 integer, specifying the stride length
        of the convolution. `strides > 1` is incompatible with
        `dilation_rate > 1`.
    padding: string, `"valid"`, `"same"` or `"causal"`(case-insensitive).
        `"valid"` means no padding. `"same"` results in padding evenly to
        the left/right or up/down of the input. When `padding="same"` and
        `strides=1`, the output has the same size as the input.
        `"causal"` results in causal(dilated) convolutions, e.g. `output[t]`
        does not depend on`input[t+1:]`. Useful when modeling temporal data
        where the model should not violate the temporal order.
        See [WaveNet: A Generative Model for Raw Audio, section2.1](
        https://arxiv.org/abs/1609.03499).
    data_format: string, either `"channels_last"` or `"channels_first"`.
        The ordering of the dimensions in the inputs. `"channels_last"`
        corresponds to inputs with shape `(batch, steps, features)`
        while `"channels_first"` corresponds to inputs with shape
        `(batch, features, steps)`. It defaults to the `image_data_format`
        value found in your Keras config file at `~/.keras/keras.json`.
        If you never set it, then it will be `"channels_last"`.
    dilation_rate: int or tuple/list of 1 integers, specifying the dilation
        rate to use for dilated convolution.
    groups: A positive int specifying the number of groups in which the
        input is split along the channel axis. Each group is convolved
        separately with `filters // groups` filters. The output is the
        concatenation of all the `groups` results along the channel axis.
        Input channels and `filters` must both be divisible by `groups`.
    activation: Activation function. If `None`, no activation is applied.
    use_bias: bool, if `True`, bias will be added to the output.
    kernel_initializer: Initializer for the convolution kernel. If `None`,
        the default initializer (`"glorot_uniform"`) will be used.
    bias_initializer: Initializer for the bias vector. If `None`, the
        default initializer (`"zeros"`) will be used.
    kernel_regularizer: Optional regularizer for the convolution kernel.
    bias_regularizer: Optional regularizer for the bias vector.
    activity_regularizer: Optional regularizer function for the output.
    kernel_constraint: Optional projection function to be applied to the
        kernel after being updated by an `Optimizer` (e.g. used to implement
        norm constraints or value constraints for layer weights). The
        function must take as input the unprojected variable and must return
        the projected variable (which must have the same shape). Constraints
        are not safe to use when doing asynchronous distributed training.
    bias_constraint: Optional projection function to be applied to the
        bias after being updated by an `Optimizer`.

Input shape:

- If `data_format="channels_last"`:
    A 3D tensor with shape: `(batch_shape, steps, channels)`
- If `data_format="channels_first"`:
    A 3D tensor with shape: `(batch_shape, channels, steps)`

Output shape:

- If `data_format="channels_last"`:
    A 3D tensor with shape: `(batch_shape, new_steps, filters)`
- If `data_format="channels_first"`:
    A 3D tensor with shape: `(batch_shape, filters, new_steps)`

Returns:
    A 3D tensor representing `activation(conv1d(inputs, kernel) + bias)`.

Raises:
    ValueError: when both `strides > 1` and `dilation_rate > 1`.

Example:

>>> # The inputs are 128-length vectors with 10 timesteps, and the
>>> # batch size is 4.
>>> x = np.random.rand(4, 10, 128)
>>> y = keras.layers.Conv1D(32, 3, activation='relu')(x)
>>> print(y.shape)
(4, 8, 32)


def stringify_default(default: Any) -> str:
    if isinstance(default, bool):
        return f'`{default}`'
    elif isinstance(default, enum.Enum):
        return f'`{str(default)}`'
    elif isinstance(default, int):
        return str(default)
    elif isinstance(default, float):
        result = str(default)
        return str(round(default, 2)) if len(result) > 6 else result
    elif isinstance(default, str):
        return str(default) if default.isnumeric() else f'`"{default}"`'
    elif isinstance(default, type):
        return f'`{default.__name__}`'
    else:
        return f'`{default}`'

Returns the string representation of a default value, as used in docstring: numbers are left as is, all other
objects are in backtiks.

Args:
    default (`Any`): The default value to process

Returns:
    `str`: The string representation of that default.


def get_image_features(self, pixel_values: torch.FloatTensor, image_sizes: torch.LongTensor):
    image_tokens = self.get_image_tokens(pixel_values, image_sizes)
    split_sizes = [height // self.vqmodel.vision_spatial_factor * (width // self.vqmodel.vision_spatial_factor + 1) for height, width in image_sizes]
    image_features = self.get_input_embeddings()(image_tokens)
    image_features = torch.split(image_features, split_sizes)
    return image_features

Tokenizes images into discrete tokens with VQGAN module and embeds
them with text embeddings layer

Args:
    pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)):
        The tensors corresponding to the input images.


def to_csv(pipe: BeamEventSet, file_path_prefix: str, schema: Schema, timestamp_key: str='timestamp', **wargs):
    header_values = [timestamp_key] + schema.index_names() + schema.feature_names()
    header_string = io.StringIO()
    header_writer = csv.writer(header_string)
    header_writer.writerow(header_values)
    return add_feature_idx_and_flatten(pipe) | 'Group by features' >> beam.GroupByKey() | 'Convert to csv' >> beam.Map(_convert_to_csv) | 'Write csv' >> beam.io.textio.WriteToText(file_path_prefix=file_path_prefix, header=header_string.getvalue(), append_trailing_newlines=False, **wargs)

Writes a Beam EventSet to a file or set of csv files.

Limitation: Timestamps are always stored as numerical values.
TODO: Support datetime timestamps.

Usage example:

```
input_node: tp.EventSetNode = ...
( p
  | tpb.from_csv("/input.csv", input_node.schema)
  | ... # processing
  | tpb.to_csv("/output.csv", output_node.schema)
)
```

Args:
    pipe: Beam pipe containing an EventSet.
    file_path_prefix: Path or path matching expression compatible with
        WriteToText.
    schema: Schema of the data. If you have a Temporian node, the schema is
        available with `node.schema`.
    timestamp_key: Key containing the timestamps.
    **wargs: Arguments passed to `beam.io.textio.WriteToText`.


class BigQueryVectorSearchParameters:
    project: str
    table_name: str
    embedding_column: str
    columns: List[str]
    neighbor_count: int
    metadata_restriction_template: Optional[Union[str, Callable[[Chunk], str]]] = None
    distance_type: Optional[str] = None
    options: Optional[Dict[str, Any]] = None
    include_distance: bool = False

    def _format_restrict(self, chunk: Chunk) -> str:
        assert self.metadata_restriction_template is not None, 'metadata_restriction_template cannot be None when formatting. This indicates a logical error in the code.'
        if callable(self.metadata_restriction_template):
            return self.metadata_restriction_template(chunk)
        return self.metadata_restriction_template.format(**chunk.metadata)

    def format_query(self, chunks: List[Chunk]) -> str:
        """Format the vector search query template."""
        base_columns_str = ', '.join((f'base.{col}' for col in self.columns))
        columns_str = ', '.join(self.columns)
        distance_clause = f", distance_type => '{self.distance_type}'" if self.distance_type else ''
        options_clause = f', options => {self.options}' if self.options else ''
        metadata_fn = '\n        CREATE TEMP FUNCTION check_metadata(\n          metadata ARRAY<STRUCT<key STRING, value STRING>>, \n          search_key STRING, \n          search_value STRING\n        ) \n        AS ((\n            SELECT COUNT(*) > 0 \n            FROM UNNEST(metadata) \n            WHERE key = search_key AND value = search_value\n        ));\n    ' if self.metadata_restriction_template else ''
        condition_groups = defaultdict(list)
        if self.metadata_restriction_template:
            for chunk in chunks:
                condition = self._format_restrict(chunk)
                condition_groups[condition].append(chunk)
        else:
            condition_groups[''] = chunks
        vector_searches = []
        for condition, group_chunks in condition_groups.items():
            embedding_unions = []
            for chunk in group_chunks:
                if chunk.embedding is None or chunk.embedding.dense_embedding is None:
                    raise ValueError(f'Chunk {chunk.id} missing embedding')
                embedding_str = f"SELECT '{chunk.id}' as id, {[float(x) for x in chunk.embedding.dense_embedding]} as embedding"
                embedding_unions.append(embedding_str)
            group_embeddings = ' UNION ALL '.join(embedding_unions)
            where_clause = f'WHERE {condition}' if condition else ''
            vector_search = f"\n            SELECT \n                query.id,\n                ARRAY_AGG(\n                    STRUCT({('distance, ' if self.include_distance else '')} {base_columns_str})\n                ) as chunks\n            FROM VECTOR_SEARCH(\n                (SELECT {columns_str}, {self.embedding_column} \n                 FROM `{self.table_name}`\n                 {where_clause}),\n                '{self.embedding_column}',\n                (SELECT * FROM ({group_embeddings})),\n                top_k => {self.neighbor_count}\n                {distance_clause}\n                {options_clause}\n            )\n            GROUP BY query.id\n        "
            vector_searches.append(vector_search)
        combined_searches = ' UNION ALL '.join(vector_searches)
        return f'\n        {metadata_fn}\n\n        {combined_searches}\n    '

Parameters for configuring BigQuery vector similarity search.

This class is used by BigQueryVectorSearchEnrichmentHandler to perform
vector similarity search using BigQuery's VECTOR_SEARCH function. It
processes :class:`~apache_beam.ml.rag.types.Chunk` objects that contain
:class:`~apache_beam.ml.rag.types.Embedding` and returns similar vectors
from a BigQuery table.

BigQueryVectorSearchEnrichmentHandler is used with
:class:`~apache_beam.transforms.enrichment.Enrichment` transform to enrich
Chunks with similar content from a vector database. For example:

>>> # Create search parameters
>>> params = BigQueryVectorSearchParameters(
...     table_name='project.dataset.embeddings',
...     embedding_column='embedding',
...     columns=['content'],
...     neighbor_count=5
... )
>>> # Use in pipeline
>>> enriched = (
...     chunks
...     | "Generate Embeddings" >> MLTransform(...)
...     | "Find Similar" >> Enrichment(
...         BigQueryVectorSearchEnrichmentHandler(
...             project='my-project',
...             vector_search_parameters=params
...         )
...     )
... )

BigQueryVectorSearchParameters encapsulates the configuration needed to
perform vector similarity search using BigQuery's VECTOR_SEARCH function.
It handles formatting the query with proper embedding vectors and metadata
restrictions.

Example with flattened metadata column:

Table schema::

    embedding: ARRAY<FLOAT64>  # Vector embedding
    content: STRING           # Document content
    language: STRING          # Direct metadata column

Code::

    >>> params = BigQueryVectorSearchParameters(
    ...     table_name='project.dataset.embeddings',
    ...     embedding_column='embedding',
    ...     columns=['content', 'language'],
    ...     neighbor_count=5,
    ...     # For column 'language', value comes from
    ...     # chunk.metadata['language']
    ...     metadata_restriction_template="language = '{language}'"
    ... )
    >>> # When processing a chunk with metadata={'language': 'en'},
    >>> # generates: WHERE language = 'en'

Example with nested repeated metadata:

Table schema::

    embedding: ARRAY<FLOAT64>  # Vector embedding
    content: STRING           # Document content
    metadata: ARRAY<STRUCT>   # Nested repeated metadata
      key: STRING,
      value: STRING
    >>

Code::

    >>> params = BigQueryVectorSearchParameters(
    ...     table_name='project.dataset.embeddings',
    ...     embedding_column='embedding',
    ...     columns=['content', 'metadata'],
    ...     neighbor_count=5,
    ...     # check_metadata(field_name, key_to_search, value_from_chunk)
    ...     metadata_restriction_template=(
    ...         "check_metadata(metadata, 'language', '{language}')"
    ...     )
    ... )
    >>> # When processing a chunk with metadata={'language': 'en'},
    >>> # generates: WHERE check_metadata(metadata, 'language', 'en')
    >>> # Searches for {key: 'language', value: 'en'} in metadata array

Args:
    project: GCP project ID containing the BigQuery dataset
    table_name: Fully qualified BigQuery table name containing vectors.
    embedding_column: Column name containing the embedding vectors.
    columns: List of columns to retrieve from matched vectors.
    neighbor_count: Number of similar vectors to return (top-k).
    metadata_restriction_template: Template string or callable for filtering
        vectors. Template string supports two formats:

        1. For flattened metadata columns:
           ``column_name = '{metadata_key}'`` where column_name is the
           BigQuery column and metadata_key is used to get the value from
           chunk.metadata[metadata_key].
        2. For nested repeated metadata (ARRAY<STRUCT<key,value>>):
           ``check_metadata(field_name, 'key_to_match', '{metadata_key}')``
           where field_name is the ARRAY<STRUCT> column in BigQuery,
           key_to_match is the literal key to search for in the array, and
           metadata_key is used to get value from
           chunk.metadata[metadata_key].

        Multiple conditions can be combined using AND/OR operators. For
        example::

            >>> # Combine metadata check with column filter
            >>> template = (
            ...     "check_metadata(metadata, 'language', '{language}') "
            ...     "AND source = '{source}'"
            ... )
            >>> # When chunk.metadata = {'language': 'en', 'source': 'web'}
            >>> # Generates: WHERE
            >>> #             check_metadata(metadata, 'language', 'en')
            >>> #           AND source = 'web'

    distance_type: Optional distance metric to use. Supported values:
        COSINE (default), EUCLIDEAN, or DOT_PRODUCT.
    options: Optional dictionary of additional VECTOR_SEARCH options.
    include_distance: Reurns the vector search similarity score if True.


def make_basket_put_payoff(strike_price, dtype=None, name=None):
    strike_price = tf.convert_to_tensor(strike_price, dtype=dtype, name='strike_price')
    put_valuer = functools.partial(_put_valuer, strike_price=strike_price, dtype=dtype, name=name)
    return put_valuer

Produces a callable from samples to payoff of a simple basket put option.

Args:
  strike_price: A `Tensor` of `dtype` consistent with `samples` and shape
    `[num_samples, num_strikes]`.
  dtype: Optional `dtype`. Either `tf.float32` or `tf.float64`. The `dtype`
    If supplied, represents the `dtype` for the 'strike_price' as well as
    for the input argument of the output payoff callable.
    Default value: `None`, which means that the `dtype` inferred by TensorFlow
    is used.
  name: Python `str` name prefixed to Ops created by the callable created
    by this function.
    Default value: `None` which is mapped to the default name 'put_valuer'

Returns:
  A callable from `Tensor` of shape `[num_samples, num_exercise_times, dim]`
  and a scalar `Tensor` representing current time to a `Tensor` of shape
  `[num_samples, num_strikes]`.


def get_image_features(self, pixel_values: torch.FloatTensor, image_sizes: torch.Tensor, vision_feature_layer: Optional[Union[int, List[int]]]=None, **kwargs):
    vision_feature_layer = vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer
    kwargs = {k: v for k, v in kwargs.items() if v is not None}
    image_outputs = self.vision_tower(pixel_values, image_sizes=image_sizes, output_hidden_states=True, **kwargs)
    if isinstance(vision_feature_layer, int):
        selected_image_feature = image_outputs.hidden_states[vision_feature_layer]
    else:
        hs_pool = [image_outputs.hidden_states[layer_idx] for layer_idx in vision_feature_layer]
        selected_image_feature = torch.cat(hs_pool, dim=-1)
    image_features = self.multi_modal_projector(selected_image_feature.squeeze(0), image_sizes)
    downsample_ratio = self.vision_tower.patch_size * self.config.spatial_merge_size
    split_sizes = [height // downsample_ratio * (width // downsample_ratio) for height, width in image_sizes]
    image_features = torch.split(image_features.squeeze(0), split_sizes)
    return image_features

Obtains image last hidden states from the vision tower and apply multimodal projection.

Args:
    pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`):
       The tensors corresponding to the input images.
    vision_feature_layer (`Union[int, List[int]]`, *optional*):
        The index of the layer to select the vision feature. If multiple indices are provided,
        the vision feature of the corresponding indices will be concatenated to form the
        vision features.
    image_sizes (`torch.Tensor`, *optional*):
        Tensor containing the image sizes as returned by the processor.
Returns:
    image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`).


class AcceleratorConfig:
    split_batches: bool = field(default=False, metadata={'help': 'Whether or not the accelerator should split the batches yielded by the dataloaders across the devices. If `True` the actual batch size used will be the same on any kind of distributed processes, but it must be a round multiple of the `num_processes` you are using. If `False`, actual batch size used will be the one set in your script multiplied by the number of processes.'})
    dispatch_batches: Optional[bool] = field(default=None, metadata={'help': 'If set to `True`, the dataloader prepared by the Accelerator is only iterated through on the main process and then the batches are split and broadcast to each process. Will default to `True` for `DataLoader` whose underlying dataset is an `IterableDataslet`, `False` otherwise.'})
    even_batches: bool = field(default=True, metadata={'help': 'If set to `True`, in cases where the total batch size across all processes does not exactly divide the dataset, samples at the start of the dataset will be duplicated so the batch can be divided equally among all workers.'})
    use_seedable_sampler: bool = field(default=True, metadata={'help': 'Whether or not use a fully seedable random sampler ([`accelerate.data_loader.SeedableRandomSampler`]).Ensures training results are fully reproducible using a different sampling technique. While seed-to-seed results may differ, on average the differences are negligible when usingmultiple different seeds to compare. Should also be ran with [`~utils.set_seed`] for the best results.'})
    non_blocking: Optional[bool] = field(default=False, metadata={'help': 'Whether to use non-blocking CUDA calls to help minimize synchronization during distributed training with prepared `DataLoader` inputs being moved to device. Best if used with `pin_memory=True` in the `TrainingArguments`. Requires accelerate v0.30.0.'})
    gradient_accumulation_kwargs: Optional[dict] = field(default=None, metadata={'help': 'Additional kwargs to configure gradient accumulation, see [`accelerate.utils.GradientAccumulationPlugin`]. Any of the following (optional) keys are acceptable:   num_steps (`int`): Will take precedence over [`~.TrainingArguments.gradient_accumulation_steps`] if     the latter is set to 1, otherwise an exception will be raised.   adjust_scheduler (`bool`): Whether to adjust the scheduler steps to account for [`~.TrainingArguments.gradient_accumulation_steps`].     The [`accelerate.utils.GradientAccumulationPlugin`] default is `True`.   sync_each_batch (`bool`): Whether to synchronize the gradients at each data batch.     The [`accelerate.utils.GradientAccumulationPlugin`] default is `False`.'})
    use_configured_state: bool = field(default=False, metadata={'help': 'Whether or not to use a pre-configured `AcceleratorState` or `PartialState` defined before calling `TrainingArguments`.If `True`, an `Accelerator` or `PartialState` must be initialized. May lead to issues using sweeps or hyperparameter tuning.'})

    @classmethod
    def from_json_file(cls, json_file):
        open_file = io.open if os.path.exists(json_file) else open
        with open_file(json_file, 'r', encoding='utf-8') as f:
            config_dict = json.load(f)
        extra_keys = sorted((key for key in config_dict.keys() if key not in cls.__dataclass_fields__.keys()))
        if len(extra_keys) > 0:
            raise ValueError(f'The config file at {json_file} had unknown keys ({extra_keys}), please try upgrading your `transformers` version or fix (and potentially remove these keys) from your config file.')
        return cls(**config_dict)

    def to_dict(self):
        return copy.deepcopy(self.__dict__)

    def pop(self, key, default=None):
        return self.__dict__.pop(key, default)

A subset of arguments relating to the underlying [`accelerate.Accelerator`]
implementation utilized in the `Trainer` that can be customized.
Mostly relating to data.

Parameters:
    split_batches (`bool`, *optional*, defaults to `False`):
        Whether or not the accelerator should split the batches yielded by the dataloaders across the devices. If
        `True` the actual batch size used will be the same on any kind of distributed processes, but it must be a
        round multiple of the `num_processes` you are using. If `False`, actual batch size used will be the one set
        in your script multiplied by the number of processes.
    dispatch_batches (`bool`, *optional*):
        If set to `True`, the dataloader prepared by the Accelerator is only iterated through on the main process
        and then the batches are split and broadcast to each process. Will default to `True` for `DataLoader` whose
        underlying dataset is an `IterableDataset`, `False` otherwise.
    even_batches (`bool`, *optional*, defaults to `True`):
        If set to `True`, in cases where the total batch size across all processes does not exactly divide the
        dataset, samples at the start of the dataset will be duplicated so the batch can be divided equally among
        all workers.
    use_seedable_sampler (`bool`, *optional*, defaults to `True`):
        Whether or not use a fully seedable random sampler ([`accelerate.data_loader.SeedableRandomSampler`]). Ensures
        training results are fully reproducible using a different sampling technique. While seed-to-seed results
        may differ, on average the differences are negligible when using multiple different seeds to compare. Should
        also be ran with [`~utils.set_seed`] for the best results.
    gradient_accumulation_kwargs (`dict`, *optional*):
        Additional kwargs to configure gradient accumulation, see [`accelerate.utils.GradientAccumulationPlugin`].
        Any of the following (optional) keys are acceptable:
          num_steps (`int`): Will take precedence over [`~.TrainingArguments.gradient_accumulation_steps`] if
            the latter is set to 1, otherwise an exception will be raised.
          adjust_scheduler (`bool`): Whether to adjust the scheduler steps to account for [`~.TrainingArguments.gradient_accumulation_steps`].
            The [`accelerate.utils.GradientAccumulationPlugin`] default is `True`.
          sync_each_batch (`bool`): Whether to synchronize the gradients at each data batch.
            The [`accelerate.utils.GradientAccumulationPlugin`] default is `False`.
    non_blocking (`bool`, *optional*, defaults to `False`):
        Whether to use non-blocking CUDA calls to help minimize synchronization during
        distributed training with prepared `DataLoader` inputs being moved to device.
        Best if used with `pin_memory=True` in the `TrainingArguments`.
    use_configured_state (`bool*, *optional*, defaults to `False`):
        Whether or not to use a pre-configured `AcceleratorState` or `PartialState` defined
        before calling `TrainingArguments`. If `True`, an `Accelerator` or `PartialState`
        must be initialized. May lead to issues using sweeps or hyperparameter tuning.


def decode(self, decoder_input_ids, encoder_outputs, encoder_attention_mask: Optional[jnp.ndarray]=None, decoder_attention_mask: Optional[jnp.ndarray]=None, decoder_position_ids: Optional[jnp.ndarray]=None, past_key_values: Optional[dict]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, train: bool=False, params: Optional[dict]=None, dropout_rng: PRNGKey=None):
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    return_dict = return_dict if return_dict is not None else self.config.return_dict
    encoder_hidden_states = encoder_outputs[0]
    batch_size, sequence_length = decoder_input_ids.shape
    if decoder_position_ids is None:
        if past_key_values is not None:
            raise ValueError('Make sure to provide `decoder_position_ids` when passing `past_key_values`.')
        if decoder_attention_mask is not None:
            decoder_position_ids = decoder_attention_mask.cumsum(-1) * decoder_attention_mask - 1
        else:
            decoder_position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))
    if decoder_attention_mask is None:
        decoder_attention_mask = jnp.ones((batch_size, sequence_length), dtype='i4')
    rngs = {}
    if dropout_rng is not None:
        rngs['dropout'] = dropout_rng
    inputs = {'params': params or self.params}
    if past_key_values:
        inputs['cache'] = past_key_values
        mutable = ['cache']
    else:
        mutable = False

    def _decoder_forward(module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs):
        decoder_module = module._get_decoder_module()
        outputs = decoder_module(input_ids=decoder_input_ids, attention_mask=decoder_attention_mask, position_ids=decoder_position_ids, **kwargs)
        hidden_states = outputs[0]
        if self.config.tie_word_embeddings:
            shared_embedding = module.model.decoder.embed_tokens.variables['params']['embedding']
            lm_logits = module.lm_head.apply({'params': {'kernel': shared_embedding.T}}, hidden_states)
        else:
            lm_logits = module.lm_head(hidden_states)
        return (lm_logits, outputs)
    outputs = self.module.apply(inputs, decoder_input_ids=jnp.array(decoder_input_ids, dtype='i4'), decoder_attention_mask=jnp.array(decoder_attention_mask, dtype='i4'), decoder_position_ids=jnp.array(decoder_position_ids, dtype='i4'), encoder_hidden_states=encoder_hidden_states, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, deterministic=not train, rngs=rngs, mutable=mutable, method=_decoder_forward)
    if past_key_values is None:
        lm_logits, decoder_outputs = outputs
    else:
        (lm_logits, decoder_outputs), past = outputs
    if return_dict:
        outputs = FlaxCausalLMOutputWithCrossAttentions(logits=lm_logits, hidden_states=decoder_outputs.hidden_states, attentions=decoder_outputs.attentions, cross_attentions=decoder_outputs.cross_attentions)
    else:
        outputs = (lm_logits,) + decoder_outputs[1:]
    if past_key_values is not None and return_dict:
        outputs['past_key_values'] = unfreeze(past['cache'])
        return outputs
    elif past_key_values is not None and (not return_dict):
        outputs = outputs[:1] + (unfreeze(past['cache']),) + outputs[1:]
    return outputs

Returns:

Example:

```python
>>> from transformers import WhisperProcessor, FlaxWhisperForConditionalGeneration
>>> from datasets import load_dataset

>>> processor = WhisperProcessor.from_pretrained("openai/whisper-tiny.en")
>>> model = FlaxWhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en", from_pt=True)
>>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
>>> inputs = processor(ds[0]["audio"]["array"], return_tensors="np")
>>> input_features = inputs.input_features
>>> encoder_outputs = model.encode(input_features=input_features)
>>> decoder_start_token_id = model.config.decoder_start_token_id

>>> decoder_input_ids = jnp.ones((inputs.input_ids.shape[0], 1), dtype="i4") * decoder_start_token_id

>>> outputs = model.decode(decoder_input_ids, encoder_outputs)
>>> last_decoder_hidden_states = outputs.last_hidden_state
```


def matrices_to_flat_transforms(transform_matrices):
    with ops.name_scope('matrices_to_flat_transforms'):
        transform_matrices = ops.convert_to_tensor(transform_matrices, name='transform_matrices')
        if transform_matrices.shape.ndims not in (2, 3):
            raise ValueError('Matrices should be 2D or 3D, got: %s' % transform_matrices)
        transforms = array_ops.reshape(transform_matrices, constant_op.constant([-1, 9]))
        transforms /= transforms[:, 8:9]
        return transforms[:, :8]

Converts affine matrices to `tf.contrib.image` projective transforms.

Note that we expect matrices that map output coordinates to input coordinates.
To convert forward transformation matrices, call `tf.linalg.inv` on the
matrices and use the result here.

Args:
  transform_matrices: One or more affine transformation matrices, for the
    reverse transformation in homogeneous coordinates. Shape `(3, 3)` or `(N,
    3, 3)`.

Returns:
  2D tensor of flat transforms with shape `(N, 8)`, which may be passed into
    `tf.contrib.image.transform`.

Raises:
  ValueError: If `transform_matrices` have an invalid shape.


Holds the output of the image processor specific `__call__` methods.

This class is derived from a python dictionary and can be used as a dictionary.

Args:
    data (`dict`):
        Dictionary of lists/arrays/tensors returned by the __call__ method ('pixel_values', etc.).
    tensor_type (`Union[None, str, TensorType]`, *optional*):
        You can give a tensor_type here to convert the lists of integers in PyTorch/TensorFlow/Numpy Tensors at
        initialization.


def copy_code(source: message.Message, target: message.Message) -> None:
    if not fhir_types.is_type_or_profile_of_code(source.DESCRIPTOR):
        raise fhir_errors.InvalidFhirError(f'Source: {source.DESCRIPTOR.full_name} is not type or profile of Code.')
    if not fhir_types.is_type_or_profile_of_code(target.DESCRIPTOR):
        raise fhir_errors.InvalidFhirError(f'Target: {target.DESCRIPTOR.full_name} is not type or profile of Code.')
    if proto_utils.are_same_message_type(source.DESCRIPTOR, target.DESCRIPTOR):
        target.CopyFrom(source)
        return
    source_value_field = source.DESCRIPTOR.fields_by_name.get('value')
    target_value_field = target.DESCRIPTOR.fields_by_name.get('value')
    if source_value_field is None or target_value_field is None:
        raise fhir_errors.InvalidFhirError(f'Unable to copy code from {source.DESCRIPTOR.full_name} to {target.DESCRIPTOR.full_name}.')
    proto_utils.copy_common_field(source, target, 'id')
    proto_utils.copy_common_field(source, target, 'extension')
    if source_value_field.type not in _CODE_TYPES or target_value_field.type not in _CODE_TYPES:
        raise ValueError(f'Unable to copy from {source.DESCRIPTOR.full_name} to {target.DESCRIPTOR.full_name}. Must have a field of TYPE_ENUM or TYPE_STRING.')
    source_value = proto_utils.get_value_at_field(source, source_value_field)
    if source_value_field.type == target_value_field.type:
        proto_utils.set_value_at_field(target, target_value_field, source_value)
    elif source_value_field.type == descriptor.FieldDescriptor.TYPE_STRING:
        source_enum_value = code_string_to_enum_value_descriptor(source_value, target_value_field.enum_type)
        proto_utils.set_value_at_field(target, target_value_field, source_enum_value.number)
    elif source_value_field.type == descriptor.FieldDescriptor.TYPE_ENUM:
        source_string_value = enum_value_descriptor_to_code_string(source_value_field.enum_type.values_by_number[source_value])
        proto_utils.set_value_at_field(target, target_value_field, source_string_value)
    else:
        raise ValueError(f'Unexpected generic value field type: {source_value_field.type}. Must be a field of TYPE_ENUM or TYPE_STRING in order to copy.')

Adds all fields from source to target.

Args:
  source: The FHIR Code instance to copy from.
  target: The target FHIR Code instance to copy to.


def __call__(self, inputs: List[Any], global_state: Optional[pg.geno.AttributeDict]=None, step: int=0) -> List[Any]:
    if self.input_element_type is not None:
        elem_type = self.input_element_type
        for i, elem in enumerate(inputs):
            if not isinstance(elem, elem_type):
                raise TypeError(f'The input is expected to be a list of {elem_type!r} but {elem!r} is encountered at position {i}.')
    if global_state is None:
        global_state = pg.geno.AttributeDict()
    self._on_input(inputs)
    outputs = self._operate(inputs, global_state=global_state, step=step)
    if self.output_element_type is not None:
        elem_type = self.output_element_type
        for i, elem in enumerate(outputs):
            if not isinstance(elem, elem_type):
                raise TypeError(f'The output is expected to be a list of {elem_type!r} but {elem!r} is encountered at position {i}.')
    return outputs

Transform a list of input values to a list of output values.

Args:
  inputs: A list of values as inputs.
  global_state: An `AttributeDict` object (dictionary that provides
    attribute access) as the global state container, which is
    readable/writable during the operation.
  step: Number of examples historically proposed, which can be used for
    determining a cross over schedule.

Returns:
  A list of values as output of current operation.


def image_guided_detection(self, pixel_values: torch.FloatTensor, query_pixel_values: Optional[torch.FloatTensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, interpolate_pos_encoding: bool=False, return_dict: Optional[bool]=None) -> Owlv2ImageGuidedObjectDetectionOutput:
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    return_dict = return_dict if return_dict is not None else self.config.return_dict
    query_feature_map = self.image_embedder(pixel_values=query_pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)[0]
    feature_map, vision_outputs = self.image_embedder(pixel_values=pixel_values, output_attentions=output_attentions, output_hidden_states=output_hidden_states, interpolate_pos_encoding=interpolate_pos_encoding)
    batch_size, num_patches_height, num_patches_width, hidden_dim = feature_map.shape
    image_feats = torch.reshape(feature_map, (batch_size, num_patches_height * num_patches_width, hidden_dim))
    batch_size, num_patches_height, num_patches_width, hidden_dim = query_feature_map.shape
    query_image_feats = torch.reshape(query_feature_map, (batch_size, num_patches_height * num_patches_width, hidden_dim))
    query_embeds, best_box_indices, query_pred_boxes = self.embed_image_query(query_image_feats, query_feature_map, interpolate_pos_encoding)
    pred_logits, class_embeds = self.class_predictor(image_feats=image_feats, query_embeds=query_embeds)
    target_pred_boxes = self.box_predictor(image_feats, feature_map, interpolate_pos_encoding)
    if not return_dict:
        output = (feature_map, query_feature_map, target_pred_boxes, query_pred_boxes, pred_logits, class_embeds, vision_outputs.to_tuple())
        output = tuple((x for x in output if x is not None))
        return output
    return Owlv2ImageGuidedObjectDetectionOutput(image_embeds=feature_map, query_image_embeds=query_feature_map, target_pred_boxes=target_pred_boxes, query_pred_boxes=query_pred_boxes, logits=pred_logits, class_embeds=class_embeds, text_model_output=None, vision_model_output=vision_outputs)

query_pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
    Pixel values of query image(s) to be detected. Pass in one query image per target image.

Examples:
```python
>>> import requests
>>> from PIL import Image
>>> import torch
>>> from transformers import AutoProcessor, Owlv2ForObjectDetection

>>> processor = AutoProcessor.from_pretrained("google/owlv2-base-patch16-ensemble")
>>> model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16-ensemble")

>>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
>>> image = Image.open(requests.get(url, stream=True).raw)
>>> query_url = "http://images.cocodataset.org/val2017/000000001675.jpg"
>>> query_image = Image.open(requests.get(query_url, stream=True).raw)
>>> inputs = processor(images=image, query_images=query_image, return_tensors="pt")

>>> # forward pass
>>> with torch.no_grad():
...     outputs = model.image_guided_detection(**inputs)

>>> target_sizes = torch.Tensor([image.size[::-1]])

>>> # Convert outputs (bounding boxes and class logits) to Pascal VOC format (xmin, ymin, xmax, ymax)
>>> results = processor.post_process_image_guided_detection(
...     outputs=outputs, threshold=0.9, nms_threshold=0.3, target_sizes=target_sizes
... )
>>> i = 0  # Retrieve predictions for the first image
>>> boxes, scores = results[i]["boxes"], results[i]["scores"]
>>> for box, score in zip(boxes, scores):
...     box = [round(i, 2) for i in box.tolist()]
...     print(f"Detected similar object with confidence {round(score.item(), 3)} at location {box}")
Detected similar object with confidence 0.938 at location [327.31, 54.94, 547.39, 268.06]
Detected similar object with confidence 0.959 at location [5.78, 360.65, 619.12, 366.39]
Detected similar object with confidence 0.902 at location [2.85, 360.01, 627.63, 380.8]
Detected similar object with confidence 0.985 at location [176.98, -29.45, 672.69, 182.83]
Detected similar object with confidence 1.0 at location [6.53, 14.35, 624.87, 470.82]
Detected similar object with confidence 0.998 at location [579.98, 29.14, 615.49, 489.05]
Detected similar object with confidence 0.985 at location [206.15, 10.53, 247.74, 466.01]
Detected similar object with confidence 0.947 at location [18.62, 429.72, 646.5, 457.72]
Detected similar object with confidence 0.996 at location [523.88, 20.69, 586.84, 483.18]
Detected similar object with confidence 0.998 at location [3.39, 360.59, 617.29, 499.21]
Detected similar object with confidence 0.969 at location [4.47, 449.05, 614.5, 474.76]
Detected similar object with confidence 0.966 at location [31.44, 463.65, 654.66, 471.07]
Detected similar object with confidence 0.924 at location [30.93, 468.07, 635.35, 475.39]
```


def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, output_attentions: Optional[bool]=False) -> Tuple[torch.FloatTensor]:
    residual = hidden_states
    hidden_states = self.layer_norm1(hidden_states)
    hidden_states, attn_weights = self.self_attn(hidden_states=hidden_states, attention_mask=attention_mask, output_attentions=output_attentions)
    hidden_states = residual + hidden_states
    residual = hidden_states
    hidden_states = self.layer_norm2(hidden_states)
    hidden_states = self.mlp(hidden_states)
    hidden_states = residual + hidden_states
    outputs = (hidden_states,)
    if output_attentions:
        outputs += (attn_weights,)
    return outputs

Args:
    hidden_states (`torch.FloatTensor`):
        Input to the layer of shape `(batch, seq_len, embed_dim)`.
    attention_mask (`torch.FloatTensor`):
        Attention mask of shape `(batch, 1, q_len, k_v_seq_len)` where padding elements are indicated by very large negative values.
    output_attentions (`bool`, *optional*, defaults to `False`):
        Whether or not to return the attentions tensors of all attention layers. See `attentions` under
        returned tensors for more detail.


def get(self):
    raise NotImplementedError()

Get the current tracking value.

Returns:
  The current tracked value, the type of which depends on the specific
  tracker implementation.


def create_base_for_fuse_batchnorm(self, pattern_match_mode='MATCH_ALL'):
    with self.cached_session() as sess:
        data_format = 'NHWC'
        if pattern_match_mode == 'MISMATCH_FORMAT':
            data_format = 'NCHW'
        inputs = [1, 4, 2, 5, 3, 6, -1, -4, -2, -5, -3, -6]
        input_op = constant_op.constant(np.array(inputs), shape=[1, 1, 6, 2] if data_format == 'NHWC' else [1, 2, 1, 6], dtype=dtypes.float32)
        weights = [1, 2, 3, 4, 0.1, 0.2, 0.3, 0.4]
        weights_op = constant_op.constant(np.array(weights), shape=[1, 2, 2, 2], dtype=dtypes.float32)
        conv_op = nn_ops.conv2d(input_op, weights_op, [1, 1, 1, 1], data_format=data_format, padding='SAME', name='conv_op')
        const_op_1 = None
        const_op_2 = constant_op.constant(1e-05, dtype=dtypes.float32)
        const_op_3 = None
        const_op_4 = None
        const_op_5 = None
        const_op_6 = None
        if data_format == 'NHWC':
            const_op_1 = constant_op.constant(np.array([0.25, 0.5]), shape=[2], dtype=dtypes.float32)
            const_op_3 = constant_op.constant(np.array([10, 20]), shape=[2], dtype=dtypes.float32)
            const_op_4 = constant_op.constant(np.array([0.1, 0.6]), shape=[2], dtype=dtypes.float32)
            const_op_5 = constant_op.constant(np.array([1.0, 2.0]), shape=[2], dtype=dtypes.float32)
            const_op_6 = constant_op.constant(np.array([0.2, 0.5]), shape=[2], dtype=dtypes.float32)
        else:
            const_op_1 = constant_op.constant(np.array([0.25, 0.5, 0.6, 0.7, 0.8, 0.9]), shape=[6], dtype=dtypes.float32)
            const_op_3 = constant_op.constant(np.array([10, 20, 30, 40, 50, 60]), shape=[6], dtype=dtypes.float32)
            const_op_4 = constant_op.constant(np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6]), shape=[6], dtype=dtypes.float32)
            const_op_5 = constant_op.constant(np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]), shape=[6], dtype=dtypes.float32)
            const_op_6 = constant_op.constant(np.array([0.2, 0.4, 0.5, 0.6, 0.7, 0.8]), shape=[6], dtype=dtypes.float32)
        add_op_1 = gen_math_ops.add(const_op_1, const_op_2)
        rsqrt_op = math_ops.rsqrt(add_op_1)
        variable_op = None
        if pattern_match_mode == 'MATCH_NO_GAMMA':
            variable_op = rsqrt_op
        else:
            variable_op = math_ops.multiply(rsqrt_op, const_op_5)
        mul_op_1 = math_ops.multiply(conv_op, variable_op)
        mul_op_2 = None
        if pattern_match_mode == 'MISMATCH_PATTERN':
            mul_op_2 = math_ops.multiply(const_op_3, const_op_6)
        else:
            mul_op_2 = math_ops.multiply(const_op_3, variable_op)
        sub_op = math_ops.subtract(const_op_4, mul_op_2)
        if pattern_match_mode == 'MATCH_SWITCH_ORDER':
            gen_math_ops.add(sub_op, mul_op_1, name='output')
        else:
            gen_math_ops.add(mul_op_1, sub_op, name='output')
        test_util.set_producer_version(ops.get_default_graph(), 8)
        original_graph = sess.graph_def
        original_result = sess.run(['output:0'])
        return (original_graph, original_result)

Create testing graph and compute the result from original graph.

Args:
  pattern_match_mode: A label string to indicate which batchnorm composition
    pattern to create in the resulting graph. "MATCH_ALL" - Create a graph
    matching the decomposed batchnorm pattern with full set of primitive
    ops. "MATCH_NO_GAMMA" - Create a graph matching the decomposed batchnorm
    pattern when gamma factor is 1 and multiplication with gamma is omitted.
    "MATCH_SWITCH_ORDER" - Create a graph matching the decomposed batchnorm
    pattern with a different order of inputs to the root Add node.
    "MISMATCH_PATTERN" - Create a graph with same set of primitive ops which
    makes up the decomposed batchnorm, but not matching the pattern.
    "MISMATCH_FORMAT" - Create a graph with NCHW format as input.

Returns:
  A GraphDef as original graph to run the decomposed batchnorm test cases.
  Computation result from executing the original graph defined by GraphDef.


def resize(self, image: np.ndarray, size: Dict[str, int], resample: PILImageResampling=PILImageResampling.BICUBIC, data_format: Optional[Union[str, ChannelDimension]]=None, input_data_format: Optional[Union[str, ChannelDimension]]=None, **kwargs) -> np.ndarray:
    default_to_square = True
    if 'shortest_edge' in size:
        size = size['shortest_edge']
        default_to_square = False
    elif 'height' in size and 'width' in size:
        size = (size['height'], size['width'])
    else:
        raise ValueError("Size must contain either 'shortest_edge' or 'height' and 'width'.")
    output_size = get_resize_output_image_size(image, size=size, default_to_square=default_to_square, input_data_format=input_data_format)
    return resize(image, size=output_size, resample=resample, data_format=data_format, input_data_format=input_data_format, **kwargs)

Resize an image. The shortest edge of the image is resized to size["shortest_edge"], with the longest edge
resized to keep the input aspect ratio.

Args:
    image (`np.ndarray`):
        Image to resize.
    size (`Dict[str, int]`):
        Size of the output image.
    resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BICUBIC`):
        Resampling filter to use when resiizing the image.
    data_format (`str` or `ChannelDimension`, *optional*):
        The channel dimension format of the image. If not provided, it will be the same as the input image.
    input_data_format (`ChannelDimension` or `str`, *optional*):
        The channel dimension format of the input image. If not provided, it will be inferred.


def run_test_in_subprocess(test_case, target_func, inputs=None, timeout=None):
    if timeout is None:
        timeout = int(os.environ.get('PYTEST_TIMEOUT', 600))
    start_methohd = 'spawn'
    ctx = multiprocessing.get_context(start_methohd)
    input_queue = ctx.Queue(1)
    output_queue = ctx.JoinableQueue(1)
    input_queue.put(inputs, timeout=timeout)
    process = ctx.Process(target=target_func, args=(input_queue, output_queue, timeout))
    process.start()
    try:
        results = output_queue.get(timeout=timeout)
        output_queue.task_done()
    except Exception as e:
        process.terminate()
        test_case.fail(e)
    process.join(timeout=timeout)
    if results['error'] is not None:
        test_case.fail(f'{results['error']}')

To run a test in a subprocess. In particular, this can avoid (GPU) memory issue.

Args:
    test_case (`unittest.TestCase`):
        The test that will run `target_func`.
    target_func (`Callable`):
        The function implementing the actual testing logic.
    inputs (`dict`, *optional*, defaults to `None`):
        The inputs that will be passed to `target_func` through an (input) queue.
    timeout (`int`, *optional*, defaults to `None`):
        The timeout (in seconds) that will be passed to the input and output queues. If not specified, the env.
        variable `PYTEST_TIMEOUT` will be checked. If still `None`, its value will be set to `600`.


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

    def __init__(self, offline_detector: OfflineDetector):
        self._offline_detector = offline_detector

    def _restore_and_convert(self, elem: tuple[tuple[Any, Any, beam.Row], Any]) -> NestedKeyedOutputT:
        """Converts the model output to AnomalyResult.

    Args:
      elem: A tuple containing the combined key (original key, temp key, row)
        and the output from RunInference.

    Returns:
      A tuple containing the keyed AnomalyResult.
    """
        (orig_key, temp_key, row), prediction = elem
        assert isinstance(prediction, AnomalyPrediction), 'Wrong model handler output type.' + f"Expected: 'AnomalyPrediction', but got '{type(prediction).__name__}'. " + 'Consider adding a post-processing function via `with_postprocess_fn` ' + f"to convert from '{type(prediction).__name__}' to 'AnomalyPrediction', " + 'or use `score_prediction_adapter` or `label_prediction_adapter` to ' + 'perform the conversion.'
        result = AnomalyResult(example=row, predictions=[dataclasses.replace(prediction, model_id=self._offline_detector._model_id)])
        return (orig_key, (temp_key, result))

    def _select_features(self, elem: tuple[Any, beam.Row]) -> tuple[Any, beam.Row]:
        assert self._offline_detector._features is not None
        k, v = elem
        row_dict = v._asdict()
        return (k, beam.Row(**{k: row_dict[k] for k in self._offline_detector._features}))

    def expand(self, input: beam.PCollection[NestedKeyedInputT]) -> beam.PCollection[NestedKeyedOutputT]:
        model_uuid = f'{self._offline_detector._model_id}:{uuid.uuid4().hex[:6]}'
        run_inference = RunInference(self._offline_detector._keyed_model_handler, **self._offline_detector._run_inference_args)
        rekeyed_model_input = input | 'Rekey' >> beam.Map(lambda x: ((x[0], x[1][0], x[1][1]), x[1][1]))
        if self._offline_detector._features is not None:
            rekeyed_model_input = rekeyed_model_input | 'Select Features' >> beam.Map(self._select_features)
        rekeyed_model_output = rekeyed_model_input | f'Call RunInference ({model_uuid})' >> run_inference
        ret = rekeyed_model_output | 'Restore keys and convert model output' >> beam.Map(self._restore_and_convert)
        if self._offline_detector._threshold_criterion:
            ret = ret | f'Run Threshold Criterion ({model_uuid})' >> RunThresholdCriterion(self._offline_detector._threshold_criterion)
        return ret

Runs a offline anomaly detector on a PCollection of data.

This PTransform applies a `OfflineDetector` to the input data, handling
custom input/output conversion and inference.

Args:
  offline_detector: The `OfflineDetector` to run.


class Manifest(abc.ABC):
    location: Location
    prev_stage_precise_start_time: t.Optional[str] = None
    status: t.Optional[DownloadStatus] = None

    def __post_init__(self):
        """Initialize the manifest."""
        pass

    def schedule(self, config_name: str, dataset: str, selection: t.Dict, location: str, user: str) -> None:
        """Indicate that a job has been scheduled for download.

        'scheduled' jobs occur before 'in-progress', 'success' or 'finished'.
        """
        scheduled_time = datetime.datetime.utcnow().replace(tzinfo=datetime.timezone.utc).isoformat(timespec='seconds')
        self.status = DownloadStatus(config_name=config_name, dataset=dataset if dataset else None, selection=selection, location=location, area=fetch_geo_polygon(selection.get('area', GLOBAL_COVERAGE_AREA)), username=user, stage=None, status=Status.SCHEDULED, error=None, size=None, scheduled_time=scheduled_time, retrieve_start_time=None, retrieve_end_time=None, fetch_start_time=None, fetch_end_time=None, download_start_time=None, download_end_time=None, upload_start_time=None, upload_end_time=None)
        self._update(self.status)

    def skip(self, config_name: str, dataset: str, selection: t.Dict, location: str, user: str) -> None:
        """Updates the manifest to mark the shards that were skipped in the current job
        as 'upload' stage and 'success' status, indicating that they have already been downloaded.
        """
        old_status = self._read(location)
        if old_status.location != location or old_status.stage != Stage.UPLOAD or old_status.status != Status.SUCCESS:
            current_utc_time = datetime.datetime.utcnow().replace(tzinfo=datetime.timezone.utc).isoformat(timespec='seconds')
            size = get_file_size(location)
            status = DownloadStatus(config_name=config_name, dataset=dataset if dataset else None, selection=selection, location=location, area=fetch_geo_polygon(selection.get('area', GLOBAL_COVERAGE_AREA)), username=user, stage=Stage.UPLOAD, status=Status.SUCCESS, error=None, size=size, scheduled_time=None, retrieve_start_time=None, retrieve_end_time=None, fetch_start_time=None, fetch_end_time=None, download_start_time=None, download_end_time=None, upload_start_time=current_utc_time, upload_end_time=current_utc_time)
            self._update(status)
            logger.info(f'Manifest updated for skipped shard: {location!r} -- {DownloadStatus.to_dict(status)!r}.')

    def _set_for_transaction(self, config_name: str, dataset: str, selection: t.Dict, location: str, user: str) -> None:
        """Reset Manifest state in preparation for a new transaction."""
        self.status = dataclasses.replace(self._read(location))
        self.status.config_name = config_name
        self.status.dataset = dataset if dataset else None
        self.status.selection = selection
        self.status.location = location
        self.status.username = user

    def __enter__(self) -> None:
        pass

    def __exit__(self, exc_type, exc_inst, exc_tb) -> None:
        """Record end status of a transaction as either 'success' or 'failure'."""
        if exc_type is None:
            status = Status.SUCCESS
            error = None
        else:
            status = Status.FAILURE
            error = '\n'.join(traceback.format_exception(exc_type, exc_inst, exc_tb))
        new_status = dataclasses.replace(self.status)
        new_status.error = error
        new_status.status = status
        current_utc_time = datetime.datetime.utcnow().replace(tzinfo=datetime.timezone.utc).isoformat(timespec='seconds')
        if new_status.stage == Stage.FETCH:
            new_status.fetch_start_time = self.prev_stage_precise_start_time
            new_status.fetch_end_time = current_utc_time
        elif new_status.stage == Stage.RETRIEVE:
            new_status.retrieve_start_time = self.prev_stage_precise_start_time
            new_status.retrieve_end_time = current_utc_time
        elif new_status.stage == Stage.DOWNLOAD:
            new_status.download_start_time = self.prev_stage_precise_start_time
            new_status.download_end_time = current_utc_time
        else:
            new_status.upload_start_time = self.prev_stage_precise_start_time
            new_status.upload_end_time = current_utc_time
        new_status.size = get_file_size(new_status.location)
        self.status = new_status
        self._update(self.status)

    def transact(self, config_name: str, dataset: str, selection: t.Dict, location: str, user: str) -> 'Manifest':
        """Create a download transaction."""
        self._set_for_transaction(config_name, dataset, selection, location, user)
        return self

    def set_stage(self, stage: Stage) -> None:
        """Sets the current stage in manifest."""
        prev_stage = self.status.stage
        new_status = dataclasses.replace(self.status)
        new_status.stage = stage
        new_status.status = Status.IN_PROGRESS
        current_utc_time = datetime.datetime.utcnow().replace(tzinfo=datetime.timezone.utc).isoformat(timespec='seconds')
        if stage == Stage.FETCH:
            new_status.fetch_start_time = current_utc_time
            new_status.fetch_end_time = None
            new_status.download_start_time = None
            new_status.download_end_time = None
        elif stage == Stage.RETRIEVE:
            new_status.retrieve_start_time = current_utc_time
        elif stage == Stage.DOWNLOAD:
            new_status.fetch_start_time = self.prev_stage_precise_start_time
            new_status.fetch_end_time = current_utc_time
            new_status.download_start_time = current_utc_time
        else:
            if prev_stage == Stage.DOWNLOAD:
                new_status.download_start_time = self.prev_stage_precise_start_time
                new_status.download_end_time = current_utc_time
            else:
                new_status.retrieve_start_time = self.prev_stage_precise_start_time
                new_status.retrieve_end_time = current_utc_time
            new_status.upload_start_time = current_utc_time
        self.status = new_status
        self._update(self.status)

    @abc.abstractmethod
    def _read(self, location: str) -> DownloadStatus:
        pass

    @abc.abstractmethod
    def _update(self, download_status: DownloadStatus) -> None:
        pass

Abstract manifest of download statuses.

Update download statuses to some storage medium.

This class lets one indicate that a download is `scheduled` or in a transaction process.
In the event of a transaction, a download will be updated with an `in-progress`, `success`
or `failure` status (with accompanying metadata).

Example:
    ```
    my_manifest = parse_manifest_location(Location('fs://some-firestore-collection'))

    # Schedule data for download
    my_manifest.schedule({'some': 'metadata'}, 'path/to/downloaded/file', 'my-username')

    # ...

    # Initiate a transaction – it will record that the download is `in-progess`
    with my_manifest.transact({'some': 'metadata'}, 'path/to/downloaded/file', 'my-username') as tx:
        # download logic here
        pass

        # ...

        # on error, will record the download as a `failure` before propagating the error.  By default, it will
        # record download as a `success`.
    ```

Attributes:
    location: An implementation-specific manifest URI.
    status: The current `DownloadStatus` of the Manifest.


class AriaSharedExpertsMLP(LlamaMLP):

    def __init__(self, config: AriaTextConfig):
        super().__init__(self)
        self.intermediate_size = config.intermediate_size * config.moe_num_shared_experts

Shared Expert MLP for shared experts.

Unlike routed experts, shared experts process all tokens without routing.
This class reconfigures the intermediate size in comparison to the LlamaMLP.

Args:
    config (`AriaTextConfig`): Configuration object for the Aria language model.