def gen_outputs(self, ds_fn, break_points, num_outputs, ckpt_saved=False, sparse_tensors=False, verify_exhausted=True, save_checkpoint_at_end=True):
    outputs = []
    if context.executing_eagerly():
        for i in range(len(break_points) + 1):
            iterator = iter(ds_fn())
            ckpt = tracking_util.Checkpoint(iterator=iterator)
            if ckpt_saved:
                ckpt_path = self._latest_ckpt()
                ckpt.restore(ckpt_path)
            start = break_points[i - 1] if i > 0 else 0
            end = break_points[i] if i < len(break_points) else num_outputs
            num_iters = end - start
            for _ in range(num_iters):
                outputs.append(self.evaluate(next(iterator)))
            if i < len(break_points) and end == num_outputs and verify_exhausted:
                with self.assertRaises(StopIteration):
                    next(iterator)
            if i == len(break_points) and verify_exhausted:
                with self.assertRaises(StopIteration):
                    next(iterator)
            if save_checkpoint_at_end or i < len(break_points):
                ckpt_options = checkpoint_options.CheckpointOptions()
                ckpt_options.experimental_enable_async_checkpoint = False
                ckpt_options.enable_async = False
                ckpt_path = ckpt.save(self._ckpt_path(), options=ckpt_options)
                ckpt_saved = True
    else:

        def get_ops():
            if ckpt_saved:
                saver = self._import_meta_graph()
                init_op, get_next_op = self._get_iterator_ops_from_collection(ds_fn, sparse_tensors=sparse_tensors)
            else:
                init_op, get_next_op, saver = self._build_graph(ds_fn, sparse_tensors=sparse_tensors)
            return (init_op, get_next_op, saver)
        for i in range(len(break_points) + 1):
            with ops.Graph().as_default() as g:
                init_op, get_next_op, saver = get_ops()
                get_next_op = remove_variants(get_next_op)
                with self.session(graph=g) as sess:
                    if ckpt_saved:
                        self._initialize(init_op, sess)
                        self._restore(saver, sess)
                    else:
                        self._initialize(init_op, sess)
                    start = break_points[i - 1] if i > 0 else 0
                    end = break_points[i] if i < len(break_points) else num_outputs
                    num_iters = end - start
                    for _ in range(num_iters):
                        outputs.append(sess.run(get_next_op))
                    if i < len(break_points) and end == num_outputs and verify_exhausted:
                        with self.assertRaises(errors.OutOfRangeError):
                            sess.run(get_next_op)
                    if i == len(break_points) and verify_exhausted:
                        with self.assertRaises(errors.OutOfRangeError):
                            sess.run(get_next_op)
                    if save_checkpoint_at_end or i < len(break_points):
                        self._save(sess, saver)
                        ckpt_saved = True
    return outputs

Generates elements from input dataset while stopping at break points.

Produces `num_outputs` outputs and saves the state of the iterator in the
Saver checkpoint.

Args:
  ds_fn: 0-argument function that returns the dataset.
  break_points: A list of integers. For each `break_point` in
    `break_points`, we produce outputs till `break_point` number of items
    have been produced and then checkpoint the state. The current graph and
    session are destroyed and a new graph and session are used to produce
    outputs till next checkpoint or till `num_outputs` elements have been
    produced. `break_point` must be <= `num_outputs`.
  num_outputs: The total number of outputs to produce from the iterator.
  ckpt_saved: Whether a checkpoint already exists.
  sparse_tensors:  Whether dataset is built from SparseTensor(s).
  verify_exhausted: Whether to verify that the iterator has been exhausted
    after producing `num_outputs` elements.
  save_checkpoint_at_end: Whether to save a checkpoint after producing all
    outputs. If False, checkpoints are saved each break point but not at the
    end. Note that checkpoints overwrite each other so there is always only
    a single checkpoint available. Defaults to True.

Returns:
  A list of `num_outputs` items.


def _calculate_scores(self, query, key):
    return NotImplementedError

Calculates attention scores.

Args:
  query: Query tensor of shape `[batch_size, Tq, dim]`.
  key: Key tensor of shape `[batch_size, Tv, dim]`.

Returns:
  Tensor of shape `[batch_size, Tq, Tv]`.


def _broadcast_dynamic_shape_next_layer_both_uniform(ac_0: _LayerBroadcaster, bc_0: _LayerBroadcaster, a_1: RowPartition, b_1: RowPartition) -> Tuple[RowPartition, _LayerBroadcaster, _LayerBroadcaster]:
    if not isinstance(ac_0, _LayerBroadcaster):
        raise TypeError('ac_0 should be a _LayerBroadcaster')
    if not isinstance(bc_0, _LayerBroadcaster):
        raise TypeError('bc_0 should be a _LayerBroadcaster')
    if not isinstance(a_1, RowPartition):
        raise TypeError('a_1 should be a RowPartition')
    if not isinstance(b_1, RowPartition):
        raise TypeError('b_1 should be a RowPartition')
    assert a_1.is_uniform()
    assert b_1.is_uniform()
    static_a_1 = tensor_util.constant_value(a_1.uniform_row_length())
    static_b_1 = tensor_util.constant_value(b_1.uniform_row_length())
    if static_a_1 is not None:
        if static_a_1 == static_b_1:
            [ac_1, _] = _broadcast_half(ac_0, a_1)
            [bc_1, _] = _broadcast_half(bc_0, b_1)
            c_1 = RowPartition.from_uniform_row_length(static_a_1, nrows=ac_0.dest_nrows())
            return [c_1, ac_1, bc_1]
        elif static_a_1 == 1:
            [bc_1, c_1b] = _broadcast_half(bc_0, b_1)
            ac_1 = _LayerBroadcaster.from_gather_index(array_ops.gather(ac_0.gather_index, c_1b.value_rowids()))
            c_1 = RowPartition.from_uniform_row_length(b_1.uniform_row_length(), nrows=bc_0.dest_nrows())
            return [c_1, ac_1, bc_1]
    if static_b_1 == 1:
        [ac_1, c_1a] = _broadcast_half(ac_0, a_1)
        bc_1 = _LayerBroadcaster.from_gather_index(array_ops.gather(bc_0.gather_index, c_1a.value_rowids()))
        c_1 = RowPartition.from_uniform_row_length(a_1.uniform_row_length(), nrows=ac_0.dest_nrows())
        return [c_1, ac_1, bc_1]

    def broadcast_noop():
        [ac_1, _] = _broadcast_half(ac_0, a_1)
        [bc_1, _] = _broadcast_half(bc_0, b_1)
        return [a_1.uniform_row_length(), ac_1.gather_index, bc_1.gather_index]

    def broadcast_a():
        [bc_1, c_1b] = _broadcast_half(bc_0, b_1)
        ac_1_gather_index = array_ops.gather(ac_0.gather_index, c_1b.value_rowids())
        return [b_1.uniform_row_length(), ac_1_gather_index, bc_1.gather_index]

    def broadcast_b():
        [ac_1, c_1a] = _broadcast_half(ac_0, a_1)
        bc_1_gather_index = array_ops.gather(bc_0.gather_index, c_1a.value_rowids())
        return [a_1.uniform_row_length(), ac_1.gather_index, bc_1_gather_index]
    can_broadcast_b = math_ops.equal(b_1.uniform_row_length(), 1)

    def no_broadcast_a():
        return cond.cond(can_broadcast_b, true_fn=broadcast_b, false_fn=broadcast_noop)
    can_broadcast_a = math_ops.equal(a_1.uniform_row_length(), 1)
    broadcast_asserts = [check_ops.assert_equal(math_ops.logical_or(math_ops.logical_or(can_broadcast_a, can_broadcast_b), math_ops.equal(a_1.uniform_row_length(), b_1.uniform_row_length())), True)]
    result = cond.cond(can_broadcast_a, true_fn=broadcast_a, false_fn=no_broadcast_a)
    [c_1_uniform_row_length, ac_1_gather_index, bc_1_gather_index] = [control_flow_ops.with_dependencies(broadcast_asserts, x) for x in result]
    c_1 = RowPartition.from_uniform_row_length(c_1_uniform_row_length, nvals=c_1_uniform_row_length * ac_0.dest_nrows(), nrows=ac_0.dest_nrows())
    ac_1 = _LayerBroadcaster.from_gather_index(ac_1_gather_index)
    bc_1 = _LayerBroadcaster.from_gather_index(bc_1_gather_index)
    return [c_1, ac_1, bc_1]

Broadcast target and next layer broadcaster of two uniform dynamic shapes.

   *--ac_0-->*<--bc_0--*
   |         |         |
  a_1       c_1       b_1
   |         |         |
   V         V         V
   *--ac_1-->*<--bc_1--*

Args:
  ac_0: _LayerBroadcaster from a to c in the previous layer.
  bc_0: _LayerBroadcaster from b to c in the previous layer.
  a_1: a RowPartition for the next layer of a.
  b_1: a RowPartition for the next layer of b.

Returns:
  (c_1, ac_1, bc_1)
  c_1: a RowPartition for the next layer of the dynamic shape.
  ac_1: _LayerBroadcaster from a to c in the next layer.
  bc_1: _LayerBroadcaster from b to c in the next layer.


class Emu3Config(PretrainedConfig):
    model_type = 'emu3'
    keys_to_ignore_at_inference = ['past_key_values']
    sub_configs = {'text_config': Emu3TextConfig, 'vq_config': Emu3VQVAEConfig}

    def __init__(self, vq_config: Union[Dict, Emu3VQVAEConfig]=None, text_config: Union[Dict, Emu3TextConfig]=None, vocabulary_map: Optional[Dict[int, int]]=None, **kwargs):
        if vq_config is None:
            vq_config = Emu3VQVAEConfig()
        elif isinstance(vq_config, dict):
            vq_config = Emu3VQVAEConfig(**vq_config)
        if text_config is None:
            text_config = Emu3TextConfig()
        elif isinstance(text_config, dict):
            text_config = Emu3TextConfig(**text_config)
        self.vq_config = vq_config
        self.text_config = text_config
        self.vocabulary_map = vocabulary_map
        self.image_token_id = vocabulary_map.get('<image>') if vocabulary_map is not None else None
        super().__init__(**kwargs)

This is the configuration class to store the configuration of a [`Emu3Model`]. It is used to instantiate a
emu3 model according to the specified arguments, defining the model architecture. Instantiating a
configuration with the defaults will yield a similar configuration to that of the
[Emu3-community/Emu3-Chat-hf](https://huggingface.co/Emu3-community/Emu3-Chat-hf).

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

Args:
    vq_config (`Union[Dict, Emu3VQVAEConfig]`, *optional*):
        Emu3VQVAEConfig instance containing the configuration for the VQ-VAE model.
    text_config (`Union[Dict, Emu3TextConfig]``, *optional*):
        Emu3TextConfig instance containing the configuration for the language model.
    vocabulary_map (`dict`, *optional*):
        A dictionary containing the vocabulary map from the tokenizer. Used to obtain tokens from the image inputs.


def forward(self, input_ids: Optional[torch.LongTensor]=None, pixel_values: Optional[torch.FloatTensor]=None, attention_mask: Optional[torch.Tensor]=None, position_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Union[List[torch.FloatTensor], Cache]]=None, token_type_ids: Optional[torch.LongTensor]=None, cache_position: Optional[torch.LongTensor]=None, inputs_embeds: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, logits_to_keep: Union[int, torch.Tensor]=0, **lm_kwargs) -> ShieldGemma2ImageClassifierOutputWithNoAttention:
    outputs = self.model(input_ids=input_ids, pixel_values=pixel_values, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, token_type_ids=token_type_ids, cache_position=cache_position, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, logits_to_keep=logits_to_keep, **lm_kwargs)
    logits = outputs.logits
    selected_logits = logits[:, -1, [self.yes_token_index, self.no_token_index]]
    probabilities = torch.softmax(selected_logits, dim=-1)
    return ShieldGemma2ImageClassifierOutputWithNoAttention(logits=selected_logits, probabilities=probabilities)

Returns:
    A `ShieldGemma2ImageClassifierOutputWithNoAttention` instance containing the logits and probabilities
    associated with the model predicting the `Yes` or `No` token as the response to that prompt, captured in the
    following properties.

        *   `logits` (`torch.Tensor` of shape `(batch_size, 2)`):
            The first position along dim=1 is the logits for the `Yes` token and the second position along dim=1 is
            the logits for the `No` token.
        *   `probabilities` (`torch.Tensor` of shape `(batch_size, 2)`):
            The first position along dim=1 is the probability of predicting the `Yes` token and the second position
            along dim=1 is the probability of predicting the `No` token.

    ShieldGemma prompts are constructed such that predicting the `Yes` token means the content *does violate* the
    policy as described. If you are only interested in the violative condition, use
    `violated = outputs.probabilities[:, 1]` to extract that slice from the output tensors.

    When used with the `ShieldGemma2Processor`, the `batch_size` will be equal to `len(images) * len(policies)`,
    and the order within the batch will be img1_policy1, ... img1_policyN, ... imgM_policyN.


class LEDSeq2SeqSequenceClassifierOutput(ModelOutput):
    loss: Optional[torch.FloatTensor] = None
    logits: Optional[torch.FloatTensor] = None
    past_key_values: Optional[List[torch.FloatTensor]] = None
    decoder_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
    decoder_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
    cross_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
    encoder_last_hidden_state: Optional[torch.FloatTensor] = None
    encoder_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
    encoder_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
    encoder_global_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None

Base class for outputs of sequence-to-sequence sentence classification models.

Args:
    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `label` 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).
    past_key_values (`List[torch.FloatTensor]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
        List of `torch.FloatTensor` of length `config.n_layers`, with each tensor of shape `(2, batch_size,
        num_heads, sequence_length, embed_size_per_head)`).

        Contains pre-computed hidden-states (key and values in the attention blocks) of the decoder that can be
        used (see `past_key_values` input) to speed up sequential decoding.
    decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
        shape `(batch_size, sequence_length, hidden_size)`.

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

        Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the
        self-attention heads.
    cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
        Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
        sequence_length)`.

        Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the
        weighted average in the cross-attention heads.
    encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
        Sequence of hidden-states at the output of the last layer of the encoder of the model.
    encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
        shape `(batch_size, sequence_length, hidden_size)`.

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

        Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the
        self-attention heads.
    encoder_global_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
        Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, x)`,
        where `x` is the number of tokens with global attention mask.

        Global attentions weights after the attention softmax, used to compute the weighted average in the
        self-attention heads. Those are the attention weights from every token with global attention to every token
        in the sequence.


def __init__(self, session, logdir, max_queue=10, flush_secs=120, filename_suffix=''):
    self._session = session
    self._logdir = logdir
    self._closed = False
    gfile.MakeDirs(self._logdir)
    with self._session.graph.as_default():
        with ops.name_scope('filewriter'):
            file_writer = summary_ops_v2.create_file_writer(logdir=self._logdir, max_queue=max_queue, flush_millis=flush_secs * 1000, filename_suffix=filename_suffix)
            with summary_ops_v2.always_record_summaries(), file_writer.as_default():
                self._event_placeholder = array_ops.placeholder_with_default(constant_op.constant('unused', dtypes.string), shape=[])
                self._add_event_op = summary_ops_v2.import_event(self._event_placeholder)
            self._init_op = file_writer.init()
            self._flush_op = file_writer.flush()
            self._close_op = file_writer.close()
        self._session.run(self._init_op)

Creates an `EventFileWriterV2` and an event file to write to.

On construction, this calls `tf.contrib.summary.create_file_writer` within
the graph from `session.graph` to look up a shared summary writer resource
for `logdir` if one exists, and create one if not. Creating the summary
writer resource in turn creates a new event file in `logdir` to be filled
with `Event` protocol buffers passed to `add_event`. Graph ops to control
this writer resource are added to `session.graph` during this init call;
stateful methods on this class will call `session.run()` on these ops.

Note that because the underlying resource is shared, it is possible that
other parts of the code using the same session may interact independently
with the resource, e.g. by flushing or even closing it. It is the caller's
responsibility to avoid any undesirable sharing in this regard.

The remaining arguments to the constructor (`flush_secs`, `max_queue`, and
`filename_suffix`) control the construction of the shared writer resource
if one is created. If an existing resource is reused, these arguments have
no effect.  See `tf.contrib.summary.create_file_writer` for details.

Args:
  session: A `tf.compat.v1.Session`. Session that will hold shared writer
    resource. The writer ops will be added to session.graph during this
    init call.
  logdir: A string. Directory where event file will be written.
  max_queue: Integer. Size of the queue for pending events and summaries.
  flush_secs: Number. How often, in seconds, to flush the
    pending events and summaries to disk.
  filename_suffix: A string. Every event file's name is suffixed with
    `filename_suffix`.


def make_rhs(self, operator, adjoint, with_batch=True):
    raise NotImplementedError('make_rhs is not defined.')

Make a rhs appropriate for calling operator.solve(rhs).

Args:
  operator:  A `LinearOperator`
  adjoint:  Python `bool`.  If `True`, we are making a 'rhs' value for the
    adjoint operator.
  with_batch: Python `bool`. If `True`, create `rhs` with the same batch
    shape as operator, and otherwise create a matrix without any batch
    shape.

Returns:
  A `Tensor`


def assign_add(self, delta, use_locking=False, name=None, read_value=True):
    raise NotImplementedError

Adds a value to this variable.

 This is essentially a shortcut for `assign_add(self, delta)`.

Args:
  delta: A `Tensor`. The value to add to this variable.
  use_locking: If `True`, use locking during the operation.
  name: The name of the operation to be created
  read_value: if True, will return something which evaluates to the new
    value of the variable; if False will return the assign op.

Returns:
  The updated variable. If `read_value` is false, instead returns None in
  Eager mode and the assign op in graph mode.


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

Computes Fresnel's sine integral of `x` element-wise.

The Fresnel sine integral is defined as the integral of `sin(t^2)` from
`0` to `x`, with the domain of definition all real numbers.

>>> tf.math.special.fresnel_sin([-1., -0.1, 0.1, 1.]).numpy()
array([-0.43825912, -0.00052359,  0.00052359,  0.43825912], dtype=float32)

This implementation is based off of the Cephes math library.

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

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

@compatibility(scipy)
Equivalent to scipy.special.fresnel first output.
@end_compatibility


def GetTermSize():
    xy = None
    for get_terminal_size in (_GetTermSizePosix, _GetTermSizeWindows, _GetTermSizeEnvironment, _GetTermSizeTput):
        try:
            xy = get_terminal_size()
            if xy:
                break
        except:
            pass
    return xy or (80, 24)

Gets the terminal x and y dimensions in characters.

_GetTermSize*() helper functions taken from:
  http://stackoverflow.com/questions/263890/

Returns:
  (columns, lines): A tuple containing the terminal x and y dimensions.


def run(argv=None, model_class=None, model_params=None, save_main_session=True, test_pipeline=None) -> PipelineResult:
    known_args, pipeline_args = parse_known_args(argv)
    pipeline_options = PipelineOptions(pipeline_args)
    pipeline_options.view_as(SetupOptions).save_main_session = save_main_session
    if not model_class:
        model_class = maskrcnn_resnet50_fpn
        model_params = {'num_classes': 91}
    model_handler = PytorchModelHandlerTensor(state_dict_path=known_args.model_state_dict_path, model_class=model_class, model_params=model_params)
    pipeline = test_pipeline
    if not test_pipeline:
        pipeline = beam.Pipeline(options=pipeline_options)
    filename_value_pair = pipeline | 'ReadImageNames' >> beam.io.ReadFromText(known_args.input) | 'FilterEmptyLines' >> beam.ParDo(filter_empty_lines) | 'ReadImageData' >> beam.Map(lambda image_name: read_image(image_file_name=image_name, path_to_dir=known_args.images_dir)) | 'PreprocessImages' >> beam.MapTuple(lambda file_name, data: (file_name, preprocess_image(data)))
    predictions = filename_value_pair | 'PyTorchRunInference' >> RunInference(KeyedModelHandler(model_handler)) | 'ProcessOutput' >> beam.ParDo(PostProcessor())
    _ = predictions | 'WriteOutput' >> beam.io.WriteToText(known_args.output, shard_name_template='', append_trailing_newlines=True)
    result = pipeline.run()
    result.wait_until_finish()
    return result

Args:
  argv: Command line arguments defined for this example.
  model_class: Reference to the class definition of the model.
              If None, maskrcnn_resnet50_fpn will be used as default .
  model_params: Parameters passed to the constructor of the model_class.
                These will be used to instantiate the model object in the
                RunInference API.
  save_main_session: Used for internal testing.
  test_pipeline: Used for internal testing.


def __init__(self, num_workers, *unused_args, **unused_kwargs):
    super().__init__(*unused_args, **unused_kwargs)
    self._num_workers = num_workers
    self._successful_ops = util.MovingSum(window_ms=1000, bucket_ms=1000)
    self._first_instant = datetime.datetime.now()
    self._throttled_secs = Metrics.counter(RampupThrottlingFn, 'cumulativeThrottlingSeconds')

Initializes a ramp-up throttler transform.

Args:
  num_workers: A hint for the expected number of workers, used to derive
               the local rate limit.


class Owlv2Processor(ProcessorMixin):
    attributes = ['image_processor', 'tokenizer']
    image_processor_class = 'Owlv2ImageProcessor'
    tokenizer_class = ('CLIPTokenizer', 'CLIPTokenizerFast')
    optional_call_args = ['query_images']

    def __init__(self, image_processor, tokenizer, **kwargs):
        super().__init__(image_processor, tokenizer)

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

        Args:
            images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`,
            `List[torch.Tensor]`):
                The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
                tensor. Both channels-first and channels-last formats are supported.
            text (`str`, `List[str]`, `List[List[str]]`):
                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
            query_images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`):
                The query image to be prepared, one query image is expected per target image to be queried. Each image
                can be a PIL image, NumPy array or PyTorch tensor. In case of a NumPy array/PyTorch tensor, each image
                should be of shape (C, H, W), where C is a number of channels, H and W are image height and width.
            return_tensors (`str` or [`~utils.TensorType`], *optional*):
                If set, will return tensors of a particular framework. Acceptable values are:
                - `'tf'`: Return TensorFlow `tf.constant` objects.
                - `'pt'`: Return PyTorch `torch.Tensor` objects.
                - `'np'`: Return NumPy `np.ndarray` objects.
                - `'jax'`: Return JAX `jnp.ndarray` objects.

        Returns:
            [`BatchFeature`]: A [`BatchFeature`] with the following fields:
            - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
              `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
              `None`).
            - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
            - **query_pixel_values** -- Pixel values of the query images to be fed to a model. Returned when `query_images` is not `None`.
        """
        output_kwargs = self._merge_kwargs(Owlv2ProcessorKwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, **kwargs, **self.prepare_and_validate_optional_call_args(*args))
        query_images = output_kwargs['images_kwargs'].pop('query_images', None)
        return_tensors = output_kwargs['common_kwargs']['return_tensors']
        if text is None and query_images is None and (images is None):
            raise ValueError('You have to specify at least one text or query image or image. All three cannot be none.')
        images, text = _validate_images_text_input_order(images, text)
        data = {}
        if text is not None:
            if isinstance(text, str) or (isinstance(text, List) and (not isinstance(text[0], List))):
                encodings = [self.tokenizer(text, **output_kwargs['text_kwargs'])]
            elif isinstance(text, List) and isinstance(text[0], List):
                encodings = []
                max_num_queries = max([len(text_single) for text_single in text])
                for text_single in text:
                    if len(text_single) != max_num_queries:
                        text_single = text_single + [' '] * (max_num_queries - len(text_single))
                    encoding = self.tokenizer(text_single, **output_kwargs['text_kwargs'])
                    encodings.append(encoding)
            else:
                raise TypeError('Input text should be a string, a list of strings or a nested list of strings')
            if return_tensors == 'np':
                input_ids = np.concatenate([encoding['input_ids'] for encoding in encodings], axis=0)
                attention_mask = np.concatenate([encoding['attention_mask'] for encoding in encodings], axis=0)
            elif return_tensors == 'jax' and is_flax_available():
                import jax.numpy as jnp
                input_ids = jnp.concatenate([encoding['input_ids'] for encoding in encodings], axis=0)
                attention_mask = jnp.concatenate([encoding['attention_mask'] for encoding in encodings], axis=0)
            elif return_tensors == 'pt' and is_torch_available():
                import torch
                input_ids = torch.cat([encoding['input_ids'] for encoding in encodings], dim=0)
                attention_mask = torch.cat([encoding['attention_mask'] for encoding in encodings], dim=0)
            elif return_tensors == 'tf' and is_tf_available():
                import tensorflow as tf
                input_ids = tf.stack([encoding['input_ids'] for encoding in encodings], axis=0)
                attention_mask = tf.stack([encoding['attention_mask'] for encoding in encodings], axis=0)
            else:
                raise ValueError('Target return tensor type could not be returned')
            data['input_ids'] = input_ids
            data['attention_mask'] = attention_mask
        if query_images is not None:
            query_pixel_values = self.image_processor(query_images, **output_kwargs['images_kwargs']).pixel_values
            data = {'query_pixel_values': query_pixel_values}
        if images is not None:
            image_features = self.image_processor(images, **output_kwargs['images_kwargs'])
            data['pixel_values'] = image_features.pixel_values
        return BatchFeature(data=data, tensor_type=return_tensors)

    def post_process_object_detection(self, *args, **kwargs):
        """
        This method forwards all its arguments to [`Owlv2ImageProcessor.post_process_object_detection`]. Please refer
        to the docstring of this method for more information.
        """
        warnings.warn('`post_process_object_detection` method is deprecated for OwlVitProcessor and will be removed in v5. Use `post_process_grounded_object_detection` instead.', FutureWarning)
        return self.image_processor.post_process_object_detection(*args, **kwargs)

    def post_process_grounded_object_detection(self, outputs: 'Owlv2ObjectDetectionOutput', threshold: float=0.1, target_sizes: Optional[Union[TensorType, List[Tuple]]]=None, text_labels: Optional[List[List[str]]]=None):
        """
        Converts the raw output of [`Owlv2ForObjectDetection`] into final bounding boxes in (top_left_x, top_left_y,
        bottom_right_x, bottom_right_y) format.

        Args:
            outputs ([`Owlv2ObjectDetectionOutput`]):
                Raw outputs of the model.
            threshold (`float`, *optional*, defaults to 0.1):
                Score threshold to keep object detection predictions.
            target_sizes (`torch.Tensor` or `List[Tuple[int, int]]`, *optional*):
                Tensor of shape `(batch_size, 2)` or list of tuples (`Tuple[int, int]`) containing the target size
                `(height, width)` of each image in the batch. If unset, predictions will not be resized.
            text_labels (`List[List[str]]`, *optional*):
                List of lists of text labels for each image in the batch. If unset, "text_labels" in output will be
                set to `None`.

        Returns:
            `List[Dict]`: A list of dictionaries, each dictionary containing the following keys:
            - "scores": The confidence scores for each predicted box on the image.
            - "labels": Indexes of the classes predicted by the model on the image.
            - "boxes": Image bounding boxes in (top_left_x, top_left_y, bottom_right_x, bottom_right_y) format.
            - "text_labels": The text labels for each predicted bounding box on the image.
        """
        output = self.image_processor.post_process_object_detection(outputs=outputs, threshold=threshold, target_sizes=target_sizes)
        if text_labels is not None and len(text_labels) != len(output):
            raise ValueError('Make sure that you pass in as many lists of text labels as images')
        if text_labels is not None:
            for image_output, image_text_labels in zip(output, text_labels):
                object_text_labels = [image_text_labels[i] for i in image_output['labels']]
                image_output['text_labels'] = object_text_labels
        else:
            for image_output in output:
                image_output['text_labels'] = None
        return output

    def post_process_image_guided_detection(self, outputs: 'Owlv2ImageGuidedObjectDetectionOutput', threshold: float=0.0, nms_threshold: float=0.3, target_sizes: Optional[Union[TensorType, List[Tuple]]]=None):
        """
        Converts the output of [`Owlv2ForObjectDetection.image_guided_detection`] into the format expected by the COCO
        api.

        Args:
            outputs ([`Owlv2ImageGuidedObjectDetectionOutput`]):
                Raw outputs of the model.
            threshold (`float`, *optional*, defaults to 0.0):
                Minimum confidence threshold to use to filter out predicted boxes.
            nms_threshold (`float`, *optional*, defaults to 0.3):
                IoU threshold for non-maximum suppression of overlapping boxes.
            target_sizes (`torch.Tensor`, *optional*):
                Tensor of shape (batch_size, 2) where each entry is the (height, width) of the corresponding image in
                the batch. If set, predicted normalized bounding boxes are rescaled to the target sizes. If left to
                None, predictions will not be unnormalized.

        Returns:
            `List[Dict]`: A list of dictionaries, each dictionary containing the following keys:
            - "scores": The confidence scores for each predicted box on the image.
            - "boxes": Image bounding boxes in (top_left_x, top_left_y, bottom_right_x, bottom_right_y) format.
            - "labels": Set to `None`.
        """
        return self.image_processor.post_process_image_guided_detection(outputs=outputs, threshold=threshold, nms_threshold=nms_threshold, target_sizes=target_sizes)

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

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

Constructs an Owlv2 processor which wraps [`Owlv2ImageProcessor`] and [`CLIPTokenizer`]/[`CLIPTokenizerFast`] into
a single processor that inherits both the image processor and tokenizer functionalities. See the
[`~OwlViTProcessor.__call__`] and [`~OwlViTProcessor.decode`] for more information.

Args:
    image_processor ([`Owlv2ImageProcessor`]):
        The image processor is a required input.
    tokenizer ([`CLIPTokenizer`, `CLIPTokenizerFast`]):
        The tokenizer is a required input.


class ColQwen2ForRetrievalOutput(ModelOutput):
    loss: Optional[torch.FloatTensor] = None
    embeddings: Optional[torch.Tensor] = None
    past_key_values: Optional[Union[List[torch.FloatTensor], Cache]] = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None

Base class for ColQwen2 embeddings output.

Args:
    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
        Language modeling loss (for next-token prediction).
    embeddings (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
        The embeddings of the model.
    past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
        Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
        `(batch_size, num_heads, sequence_length, embed_size_per_head)`)

        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
        `past_key_values` input) to speed up sequential decoding.
    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, if the model has an embedding layer, +
        one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.

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

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


def initialize_tpu_system_impl(cluster_resolver, tpu_cluster_resolver_cls):
    if tpu_cluster_resolver_cls is None or not issubclass(tpu_cluster_resolver_cls, cluster_resolver_lib.ClusterResolver) or (not hasattr(tpu_cluster_resolver_cls, 'tpu_hardware_feature')):
        raise TypeError('tpu_cluster_resolver_cls is not tf.distribute.cluster_resolver.TPUClusterResolver.')
    logging.info('Deallocate tpu buffers before initializing tpu system.')
    context.context()._clear_caches()
    context.context().clear_kernel_cache()
    gc.collect()
    job = None
    if cluster_resolver is None:
        if context.executing_eagerly():
            curr_device = device.DeviceSpec.from_string(context.context().device_name)
            if curr_device.job is not None:
                job = '{}/replica:0/task:0'.format(curr_device.job)
        cluster_resolver = tpu_cluster_resolver_cls('')
    assert isinstance(cluster_resolver, tpu_cluster_resolver_cls)
    tpu_name = compat.as_text(cluster_resolver._tpu)
    if tpu_name in _INITIALIZED_TPU_SYSTEMS:
        logging.warning('TPU system %s has already been initialized. Reinitializing the TPU can cause previously created variables on TPU to be lost.', tpu_name)
    logging.info('Initializing the TPU system: %s', tpu_name)
    if tpu_name not in _LOCAL_MASTERS:
        job = '{}/replica:0/task:0'.format(cluster_resolver.get_job_name())
    if context.executing_eagerly():

        @def_function.function(autograph=False)
        def _tpu_init_fn():
            return tpu.initialize_system(job=job, compilation_failure_closes_chips=False, tpu_cancellation_closes_chips=False)
        run_eagerly = def_function.functions_run_eagerly()
        if run_eagerly:
            logging.warning('It looks like tf.function behavior was disabled, perhaps using tf.config.run_functions_eagerly. tf.tpu.experimental.initialize_tpu_system requires tf.function to work. This primitive will override the disable.')
            def_function.run_functions_eagerly(False)
        try:
            with ops.device(tpu._tpu_system_device_name(job)):
                output = _tpu_init_fn()
            context.async_wait()
        except errors.InvalidArgumentError as e:
            raise errors.NotFoundError(None, None, 'TPUs not found in the cluster. Failed in initialization: ' + str(e))
        finally:
            if run_eagerly is not None:
                def_function.run_functions_eagerly(run_eagerly)
        context.context()._initialize_logical_devices()
        serialized_topology = output.numpy()
    elif not ops.executing_eagerly_outside_functions():
        master = cluster_resolver.master()
        cluster_spec = cluster_resolver.cluster_spec()
        session_config = config_pb2.ConfigProto(allow_soft_placement=True)
        if cluster_spec:
            session_config.cluster_def.CopyFrom(cluster_spec.as_cluster_def())
        with ops.Graph().as_default():
            with session_lib.Session(config=session_config, target=master) as sess:
                serialized_topology = sess.run(tpu.initialize_system())
    else:
        with ops.device(tpu._tpu_system_device_name(job)):
            serialized_topology = tpu.initialize_system(job=job, compilation_failure_closes_chips=False)
            return serialized_topology
    logging.info('Finished initializing TPU system.')
    tpu_topology = topology.Topology(serialized=serialized_topology)
    cluster_resolver.set_tpu_topology(serialized_topology)
    _INITIALIZED_TPU_SYSTEMS[tpu_name] = tpu_topology
    _tpu_worker_address.get_cell('address').set(cluster_resolver.get_master())
    return tpu_topology

Implementation for tpu.experimental.initialize_tpu_system.

Kept separate to avoid tpu_oss code duplication.

Initialize the TPU devices.

Args:
  cluster_resolver: A tf.distribute.cluster_resolver.TPUClusterResolver,
      which provides information about the TPU cluster.
  tpu_cluster_resolver_cls: a reference to
      tf.distribute.cluster_resolver.TPUClusterResolver so that an instance
      of it can be initialized if cluster_resolver is None.
Returns:
  The tf.tpu.Topology object for the topology of the TPU cluster. If called
  inside tf.function, it returns the serialized topology object instead.

Raises:
  RuntimeError: If running inside a tf.function.
  NotFoundError: If no TPU devices found in eager mode.
  TypeError: If tpu_cluster_resolver_cls is
      not tf.distribute.cluster_resolver.TPUClusterResolver.


def forward(self, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.Tensor]=None, decoder_input_ids: Optional[torch.LongTensor]=None, decoder_attention_mask: Optional[torch.Tensor]=None, head_mask: Optional[torch.Tensor]=None, decoder_head_mask: Optional[torch.Tensor]=None, cross_attn_head_mask: Optional[torch.Tensor]=None, encoder_outputs: Optional[Union[Tuple[torch.Tensor], BaseModelOutput]]=None, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]]=None, inputs_embeds: Optional[torch.FloatTensor]=None, decoder_inputs_embeds: Optional[torch.FloatTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, cache_position: Optional[torch.Tensor]=None) -> Seq2SeqModelOutput:
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    use_cache = use_cache if use_cache is not None else self.config.use_cache
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    if encoder_outputs is None:
        encoder_outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask, head_mask=head_mask, inputs_embeds=inputs_embeds, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    elif return_dict and (not isinstance(encoder_outputs, BaseModelOutput)):
        encoder_outputs = BaseModelOutput(last_hidden_state=encoder_outputs[0], hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None)
    decoder_outputs = self.decoder(input_ids=decoder_input_ids, attention_mask=decoder_attention_mask, encoder_hidden_states=encoder_outputs[0], encoder_attention_mask=attention_mask, head_mask=decoder_head_mask, cross_attn_head_mask=cross_attn_head_mask, past_key_values=past_key_values, inputs_embeds=decoder_inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, cache_position=cache_position)
    if not return_dict:
        return decoder_outputs + encoder_outputs
    return Seq2SeqModelOutput(last_hidden_state=decoder_outputs.last_hidden_state, past_key_values=decoder_outputs.past_key_values, decoder_hidden_states=decoder_outputs.hidden_states, decoder_attentions=decoder_outputs.attentions, cross_attentions=decoder_outputs.cross_attentions, encoder_last_hidden_state=encoder_outputs.last_hidden_state, encoder_hidden_states=encoder_outputs.hidden_states, encoder_attentions=encoder_outputs.attentions)

decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
    Indices of decoder input sequence tokens in the vocabulary.

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

    [What are decoder input IDs?](../glossary#decoder-input-ids)

    Marian uses the `pad_token_id` as the starting token for `decoder_input_ids` generation. If
    `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
    `past_key_values`).
decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
    Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
    be used by default.
cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):
    Mask to nullify selected heads of the cross-attention modules in the decoder. Mask values selected in `[0,
    1]`:

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

Example:

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

>>> tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-de")
>>> model = MarianModel.from_pretrained("Helsinki-NLP/opus-mt-en-de")

>>> inputs = tokenizer("Studies have been shown that owning a dog is good for you", return_tensors="pt")
>>> decoder_inputs = tokenizer(
...     "<pad> Studien haben gezeigt dass es hilfreich ist einen Hund zu besitzen",
...     return_tensors="pt",
...     add_special_tokens=False,
... )
>>> outputs = model(input_ids=inputs.input_ids, decoder_input_ids=decoder_inputs.input_ids)

>>> last_hidden_states = outputs.last_hidden_state
>>> list(last_hidden_states.shape)
[1, 26, 512]
```


def as_str_any(value, encoding='utf-8'):
    if isinstance(value, bytes):
        return as_str(value, encoding=encoding)
    else:
        return str(value)

Converts input to `str` type.

   Uses `str(value)`, except for `bytes` typed inputs, which are converted
   using `as_str`.

Args:
  value: A object that can be converted to `str`.
  encoding: Encoding for `bytes` typed inputs.

Returns:
  A `str` object.


def _useRPCConfig(self):
    return config_pb2.ConfigProto(rpc_options=rpc_options_pb2.RPCOptions(use_rpc_for_inprocess_master=True))

Return a `tf.compat.v1.ConfigProto` that ensures we use the RPC stack for tests.

This configuration ensures that we continue to exercise the gRPC
stack when testing, rather than using the in-process optimization,
which avoids using gRPC as the transport between a client and
master in the same process.

Returns:
  A `tf.compat.v1.ConfigProto`.


def generalized_box_iou(boxes1, boxes2):
    if not (boxes1[:, 2:] >= boxes1[:, :2]).all():
        raise ValueError(f'boxes1 must be in [x0, y0, x1, y1] (corner) format, but got {boxes1}')
    if not (boxes2[:, 2:] >= boxes2[:, :2]).all():
        raise ValueError(f'boxes2 must be in [x0, y0, x1, y1] (corner) format, but got {boxes2}')
    iou, union = box_iou(boxes1, boxes2)
    top_left = torch.min(boxes1[:, None, :2], boxes2[:, :2])
    bottom_right = torch.max(boxes1[:, None, 2:], boxes2[:, 2:])
    width_height = (bottom_right - top_left).clamp(min=0)
    area = width_height[:, :, 0] * width_height[:, :, 1]
    return iou - (area - union) / area

Generalized IoU from https://giou.stanford.edu/. The boxes should be in [x0, y0, x1, y1] (corner) format.

Returns:
    `torch.FloatTensor`: a [N, M] pairwise matrix, where N = len(boxes1) and M = len(boxes2)


class Subtract(_Merge):

    @tf_utils.shape_type_conversion
    def build(self, input_shape):
        super(Subtract, self).build(input_shape)
        if len(input_shape) != 2:
            raise ValueError('A `Subtract` layer should be called on exactly 2 inputs')

    def _merge_function(self, inputs):
        if len(inputs) != 2:
            raise ValueError('A `Subtract` layer should be called on exactly 2 inputs')
        return inputs[0] - inputs[1]

Layer that subtracts two inputs.

It takes as input a list of tensors of size 2,
both of the same shape, and returns a single tensor, (inputs[0] - inputs[1]),
also of the same shape.

Examples:

```python
    import keras

    input1 = keras.layers.Input(shape=(16,))
    x1 = keras.layers.Dense(8, activation='relu')(input1)
    input2 = keras.layers.Input(shape=(32,))
    x2 = keras.layers.Dense(8, activation='relu')(input2)
    # Equivalent to subtracted = keras.layers.subtract([x1, x2])
    subtracted = keras.layers.Subtract()([x1, x2])

    out = keras.layers.Dense(4)(subtracted)
    model = keras.models.Model(inputs=[input1, input2], outputs=out)
```


class RandomElasticTransform(BaseImagePreprocessingLayer):
    _USE_BASE_FACTOR = False
    _FACTOR_BOUNDS = (0, 1)
    _SUPPORTED_INTERPOLATION = ('nearest', 'bilinear')
    _SUPPORTED_FILL_MODES = {'constant', 'nearest', 'wrap', 'mirror', 'reflect'}

    def __init__(self, factor=1.0, scale=1.0, interpolation='bilinear', fill_mode='reflect', fill_value=0.0, value_range=(0, 255), seed=None, data_format=None, **kwargs):
        super().__init__(data_format=data_format, **kwargs)
        self._set_factor(factor)
        self.scale = self._set_factor_by_name(scale, 'scale')
        self.interpolation = interpolation
        self.fill_mode = fill_mode
        self.fill_value = fill_value
        self.value_range = value_range
        self.seed = seed
        self.generator = SeedGenerator(seed)
        if interpolation not in self._SUPPORTED_INTERPOLATION:
            raise NotImplementedError(f'Unknown `interpolation` {interpolation}. Expected of one {self._SUPPORTED_INTERPOLATION}.')
        if fill_mode not in self._SUPPORTED_FILL_MODES:
            raise NotImplementedError(f'Unknown `fill_mode` {fill_mode}. Expected of one {self._SUPPORTED_FILL_MODES}.')
        if self.data_format == 'channels_first':
            self.height_axis = -2
            self.width_axis = -1
            self.channel_axis = -3
        else:
            self.height_axis = -3
            self.width_axis = -2
            self.channel_axis = -1

    def _set_factor_by_name(self, factor, name):
        error_msg = f'The `{name}` argument should be a number (or a list of two numbers) in the range [{self._FACTOR_BOUNDS[0]}, {self._FACTOR_BOUNDS[1]}]. Received: factor={factor}'
        if isinstance(factor, (tuple, list)):
            if len(factor) != 2:
                raise ValueError(error_msg)
            if factor[0] > self._FACTOR_BOUNDS[1] or factor[1] < self._FACTOR_BOUNDS[0]:
                raise ValueError(error_msg)
            lower, upper = sorted(factor)
        elif isinstance(factor, (int, float)):
            if factor < self._FACTOR_BOUNDS[0] or factor > self._FACTOR_BOUNDS[1]:
                raise ValueError(error_msg)
            factor = abs(factor)
            lower, upper = [max(-factor, self._FACTOR_BOUNDS[0]), factor]
        else:
            raise ValueError(error_msg)
        return (lower, upper)

    def get_random_transformation(self, data, training=True, seed=None):
        if not training:
            return None
        if self.scale[1] == 0 or self.factor[1] == 0:
            return None
        if isinstance(data, dict):
            images = data['images']
        else:
            images = data
        images_shape = self.backend.shape(images)
        unbatched = len(images_shape) == 3
        if unbatched:
            batch_size = 1
        else:
            batch_size = images_shape[0]
        seed = seed or self._get_seed_generator(self.backend._backend)
        transformation_probability = self.backend.random.uniform(shape=(batch_size,), minval=self.factor[0], maxval=self.factor[1], seed=seed)
        random_threshold = self.backend.random.uniform(shape=(batch_size,), minval=0.0, maxval=1.0, seed=seed)
        apply_transform = random_threshold < transformation_probability
        distortion_factor = self.backend.random.uniform(shape=(), minval=self.scale[0], maxval=self.scale[1], seed=seed, dtype=self.compute_dtype)
        return {'apply_transform': apply_transform, 'distortion_factor': distortion_factor, 'seed': seed}

    def get_elastic_transform_params(self, height, width, factor):
        alpha_scale = 0.1 * factor
        sigma_scale = 0.05 * factor
        alpha = max(height, width) * alpha_scale
        sigma = min(height, width) * sigma_scale
        return (alpha, sigma)

    def transform_images(self, images, transformation, training=True):
        images = self.backend.cast(images, self.compute_dtype)
        if training and transformation is not None:
            apply_transform = transformation['apply_transform']
            distortion_factor = transformation['distortion_factor']
            seed = transformation['seed']
            height, width = (images.shape[self.height_axis], images.shape[self.width_axis])
            alpha, sigma = self.get_elastic_transform_params(height, width, distortion_factor)
            transformed_images = self.backend.image.elastic_transform(images, alpha=alpha, sigma=sigma, interpolation=self.interpolation, fill_mode=self.fill_mode, fill_value=self.fill_value, seed=seed, data_format=self.data_format)
            apply_transform = apply_transform[:, None, None] if len(images.shape) == 3 else apply_transform[:, None, None, None]
            images = self.backend.numpy.where(apply_transform, transformed_images, images)
            images = self.backend.numpy.clip(images, self.value_range[0], self.value_range[1])
            images = self.backend.cast(images, self.compute_dtype)
        return images

    def transform_labels(self, labels, transformation, training=True):
        return labels

    def transform_segmentation_masks(self, segmentation_masks, transformation, training=True):
        return self.transform_images(segmentation_masks, transformation, training=training)

    def compute_output_shape(self, input_shape):
        return input_shape

    def get_config(self):
        base_config = super().get_config()
        config = {'factor': self.factor, 'scale': self.scale, 'interpolation': self.interpolation, 'fill_mode': self.fill_mode, 'fill_value': self.fill_value, 'value_range': self.value_range, 'seed': self.seed}
        return {**base_config, **config}

A preprocessing layer that applies random elastic transformations.

This layer distorts input images by applying elastic deformations,
simulating a physically realistic transformation. The magnitude of the
distortion is controlled by the `scale` parameter, while the `factor`
determines the probability of applying the transformation.

Args:
    factor: A single float or a tuple of two floats.
        `factor` controls the probability of applying the transformation.
        - `factor=0.0` ensures no erasing is applied.
        - `factor=1.0` means erasing is always applied.
        - If a tuple `(min, max)` is provided, a probability value
          is sampled between `min` and `max` for each image.
        - If a single float is provided, a probability is sampled
          between `0.0` and the given float.
        Default is 1.0.
    scale: A float or a tuple of two floats defining the magnitude of
        the distortion applied.
        - If a tuple `(min, max)` is provided, a random scale value is
          sampled within this range.
        - If a single float is provided, a random scale value is sampled
          between `0.0` and the given float.
        Default is 1.0.
    interpolation: Interpolation mode. Supported values: `"nearest"`,
        `"bilinear"`.
    fill_mode: Points outside the boundaries of the input are filled
        according to the given mode. Available methods are `"constant"`,
        `"nearest"`, `"wrap"` and `"reflect"`. Defaults to `"constant"`.
        - `"reflect"`: `(d c b a | a b c d | d c b a)`
            The input is extended by reflecting about the edge of the last
            pixel.
        - `"constant"`: `(k k k k | a b c d | k k k k)`
            The input is extended by filling all values beyond
            the edge with the same constant value k specified by
            `fill_value`.
        - `"wrap"`: `(a b c d | a b c d | a b c d)`
            The input is extended by wrapping around to the opposite edge.
        - `"nearest"`: `(a a a a | a b c d | d d d d)`
            The input is extended by the nearest pixel.
        Note that when using torch backend, `"reflect"` is redirected to
        `"mirror"` `(c d c b | a b c d | c b a b)` because torch does not
        support `"reflect"`.
        Note that torch backend does not support `"wrap"`.
    fill_value: a float represents the value to be filled outside the
        boundaries when `fill_mode="constant"`.
    value_range: the range of values the incoming images will have.
        Represented as a two-number tuple written `[low, high]`. This is
        typically either `[0, 1]` or `[0, 255]` depending on how your
        preprocessing pipeline is set up.
    seed: Integer. Used to create a random seed.


def _parse_cli_args(argv):
    parser = argparse.ArgumentParser(description='Mobly Suite Executable.')
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('-c', '--config', type=str, metavar='<PATH>', help='Path to the test configuration file.')
    group.add_argument('-l', '--list_tests', action='store_true', help='Print the names of the tests defined in a script without executing them.')
    parser.add_argument('--tests', '--test_case', nargs='+', type=str, metavar='[ClassA[_test_suffix][.test_a] ClassB[_test_suffix][.test_b] ...]', help='A list of test classes and optional tests to execute. Note: test_suffix based names are only supported when running by suite class')
    parser.add_argument('-tb', '--test_bed', nargs='+', type=str, metavar='[<TEST BED NAME1> <TEST BED NAME2> ...]', help='Specify which test beds to run tests on.')
    parser.add_argument('-v', '--verbose', action='store_true', help='Set console logger level to DEBUG')
    if not argv:
        argv = sys.argv[1:]
    return parser.parse_known_args(argv)[0]

Parses cli args that are consumed by Mobly.

Args:
  argv: A list that is then parsed as cli args. If None, defaults to cli
    input.

Returns:
  Namespace containing the parsed args.


def ctc(y_true, y_pred):
    if len(ops.shape(y_true)) != 2:
        raise ValueError(f'Targets `y_true` are expected to be a tensor of shape `(batch_size, max_length)` in integer format. Received: y_true.shape={ops.shape(y_true)}')
    if len(ops.shape(y_pred)) != 3:
        raise ValueError(f'Logits `y_pred` are expected to be a tensor of shape `(batch_size, max_length, num_classes)`. Received: y_pred.shape={ops.shape(y_pred)}')
    mask_index = 0
    batch_length = ops.shape(y_pred)[0]
    input_length = ops.shape(y_pred)[1]
    input_length = input_length * ops.ones((batch_length,), dtype='int32')
    label_length = ops.cast(ops.sum(y_true != mask_index, axis=-1), dtype='int32')
    return ops.ctc_loss(y_true, y_pred, label_length, input_length, mask_index=mask_index)

CTC (Connectionist Temporal Classification) loss.

Args:
    y_true: A tensor of shape `(batch_size, max_length)` containing
        the true labels in integer format. `0` always represents
        the blank/mask index and should not be used for classes.
    y_pred: A tensor of shape `(batch_size, max_length, num_classes)`
        containing logits (the output of your model).
        They should *not* be normalized via softmax.


def forward(self, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.Tensor]=None, head_mask: Optional[torch.Tensor]=None, inputs_embeds: Optional[torch.FloatTensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[Tuple, BaseModelOutput]:
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    if input_ids is not None and inputs_embeds is not None:
        raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time')
    elif input_ids is not None:
        input = input_ids
        input_ids = input_ids.view(-1, input_ids.shape[-1])
    elif inputs_embeds is not None:
        input = inputs_embeds[:, :, -1]
    else:
        raise ValueError('You have to specify either input_ids or inputs_embeds')
    if inputs_embeds is None:
        inputs_embeds = self.embed_tokens(input_ids)
    embed_pos = self.embed_positions(input)
    embed_pos = embed_pos.to(inputs_embeds.device)
    hidden_states = inputs_embeds + embed_pos
    hidden_states = self.layernorm_embedding(hidden_states)
    hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
    attention_mask = self._update_full_mask(attention_mask, inputs_embeds)
    encoder_states = () if output_hidden_states else None
    all_attentions = () if output_attentions else None
    if head_mask is not None:
        if head_mask.size()[0] != len(self.layers):
            raise ValueError(f'The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}.')
    for idx, encoder_layer in enumerate(self.layers):
        if output_hidden_states:
            encoder_states = encoder_states + (hidden_states,)
        to_drop = False
        if self.training:
            dropout_probability = torch.rand([])
            if dropout_probability < self.layerdrop:
                to_drop = True
        if to_drop:
            layer_outputs = (None, None)
        else:
            if self.gradient_checkpointing and self.training:
                layer_outputs = self._gradient_checkpointing_func(encoder_layer.__call__, hidden_states, attention_mask, head_mask[idx] if head_mask is not None else None, output_attentions)
            else:
                layer_outputs = encoder_layer(hidden_states, attention_mask, layer_head_mask=head_mask[idx] if head_mask is not None else None, output_attentions=output_attentions)
            hidden_states = layer_outputs[0]
        if output_attentions:
            all_attentions = all_attentions + (layer_outputs[1],)
    if output_hidden_states:
        encoder_states = encoder_states + (hidden_states,)
    if not return_dict:
        return tuple((v for v in [hidden_states, encoder_states, all_attentions] if v is not None))
    return BaseModelOutput(last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions)

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

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

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

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

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

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

    inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
        Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
        This is useful if you want more control over how to convert `input_ids` indices into associated vectors
        than the model's internal embedding lookup matrix.
    output_attentions (`bool`, *optional*):
        Whether or not to return the attentions tensors of all attention layers. See `attentions` under
        returned tensors for more detail.
    output_hidden_states (`bool`, *optional*):
        Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
        for more detail.
    return_dict (`bool`, *optional*):
        Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.


def get_head_mask(self, head_mask: Optional[Tensor], num_hidden_layers: int, is_attention_chunked: bool=False) -> Tensor:
    if head_mask is not None:
        head_mask = self._convert_head_mask_to_5d(head_mask, num_hidden_layers)
        if is_attention_chunked is True:
            head_mask = head_mask.unsqueeze(-1)
    else:
        head_mask = [None] * num_hidden_layers
    return head_mask

Prepare the head mask if needed.

Args:
    head_mask (`torch.Tensor` with shape `[num_heads]` or `[num_hidden_layers x num_heads]`, *optional*):
        The mask indicating if we should keep the heads or not (1.0 for keep, 0.0 for discard).
    num_hidden_layers (`int`):
        The number of hidden layers in the model.
    is_attention_chunked (`bool`, *optional*, defaults to `False`):
        Whether or not the attentions scores are computed by chunks or not.

Returns:
    `torch.Tensor` with shape `[num_hidden_layers x batch x num_heads x seq_length x seq_length]` or list with
    `[None]` for each layer.


class DetrDecoderOutput(BaseModelOutputWithCrossAttentions):
    intermediate_hidden_states: Optional[torch.FloatTensor] = None

Base class for outputs of the DETR decoder. This class adds one attribute to BaseModelOutputWithCrossAttentions,
namely an optional stack of intermediate decoder activations, i.e. the output of each decoder layer, each of them
gone through a layernorm. This is useful when training the model with auxiliary decoding losses.

Args:
    last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
        Sequence of hidden-states at the output of the last layer of the model.
    hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
        shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer
        plus the initial embedding outputs.
    attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
        Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
        sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in
        the self-attention heads.
    cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=True` is passed or when `config.output_attentions=True`):
        Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
        sequence_length)`. Attentions weights of the decoder's cross-attention layer, after the attention softmax,
        used to compute the weighted average in the cross-attention heads.
    intermediate_hidden_states (`torch.FloatTensor` of shape `(config.decoder_layers, batch_size, num_queries, hidden_size)`, *optional*, returned when `config.auxiliary_loss=True`):
        Intermediate decoder activations, i.e. the output of each decoder layer, each of them gone through a
        layernorm.


class Llama4VisionConfig(PretrainedConfig):
    base_model_tp_plan = {'model.layers.*.self_attn.q_proj': 'colwise', 'model.layers.*.self_attn.k_proj': 'colwise', 'model.layers.*.self_attn.v_proj': 'colwise', 'model.layers.*.self_attn.o_proj': 'rowwise', 'vision_adapter.mlp.fc1': 'colwise', 'vision_adapter.mlp.fc2': 'rowwise', 'patch_embedding.linear': 'colwise_rep'}
    model_type = 'llama4_vision_model'
    base_config_key = 'vision_config'

    def __init__(self, hidden_size: int=768, hidden_act: str='gelu', num_hidden_layers: int=34, num_attention_heads: int=16, num_channels: int=3, intermediate_size: int=5632, vision_output_dim: int=7680, image_size: int=448, patch_size: int=14, norm_eps: float=1e-05, vision_feature_layer=-1, vision_feature_select_strategy='default', initializer_range: float=0.02, pixel_shuffle_ratio=0.5, projector_input_dim=4096, projector_output_dim=4096, multi_modal_projector_bias=False, projector_dropout=0.0, attention_dropout=0.0, rope_theta=10000, **kwargs):
        self.hidden_size = hidden_size
        self.hidden_act = hidden_act
        self.num_hidden_layers = num_hidden_layers
        self.num_channels = num_channels
        self.intermediate_size = intermediate_size
        self.image_size = image_size
        self.vision_output_dim = vision_output_dim
        self.patch_size = patch_size
        self.norm_eps = norm_eps
        self.num_attention_heads = num_attention_heads
        self.initializer_range = initializer_range
        self.pixel_shuffle_ratio = pixel_shuffle_ratio
        self.projector_input_dim = projector_input_dim
        self.projector_output_dim = projector_output_dim
        self.multi_modal_projector_bias = multi_modal_projector_bias
        self.projector_dropout = projector_dropout
        self.attention_dropout = attention_dropout
        self.vision_feature_layer = vision_feature_layer
        self.vision_feature_select_strategy = vision_feature_select_strategy
        self.rope_theta = rope_theta
        super().__init__(**kwargs)

This is the configuration class to store the configuration of a [`Llama4VisionModel`]. It is used to instantiate a
Llama4 vision model according to the specified arguments, defining the model architecture. Instantiating a configuration
with the defaults will yield a similar configuration to that of the Llama4 109B.

e.g. [meta-llama/Llama-4-Scout-17B-16E](https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E)

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

Args:
    hidden_size (`int`, *optional*, defaults to 768):
        Dimensionality of the encoder layers and the pooler layer.
    hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
        The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
        `"relu"`, `"selu"` and `"gelu_new"` `"quick_gelu"` are supported.
    num_hidden_layers (`int`, *optional*, defaults to 34):
        Number of hidden layers in the Transformer encoder.
    num_attention_heads (`int`, *optional*, defaults to 16):
        Number of attention heads for each attention layer in the Transformer encoder.
    num_channels (`int`, *optional*, defaults to 3):
        Number of channels in the input image.
    intermediate_size (`int`, *optional*, defaults to 5632):
        Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
    vision_output_dim (`int`, *optional*, defaults to 7680):
        Dimensionality of the vision model output. Includes output of transformer
        encoder with intermediate layers and global transformer encoder.
    image_size (`int`, *optional*, defaults to 448):
        The size (resolution) of each image *tile*.
    patch_size (`int`, *optional*, defaults to 14):
        The size (resolution) of each patch.
    norm_eps (`float`, *optional*, defaults to 1e-05):
        The epsilon used by the layer normalization layers.
    vision_feature_layer (``, *optional*, defaults to -1): TODO
    vision_feature_select_strategy (`int`, *optional*, defaults to `"default"`): TODO
    initializer_range (`float`, *optional*, defaults to 0.02):
        The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
    pixel_shuffle_ratio (`int`, *optional*, defaults to 0.5): TODO
    projector_input_dim (`int`, *optional*, defaults to 4096): TODO
    projector_output_dim (`int`, *optional*, defaults to 4096): TODO
    multi_modal_projector_bias (`int`, *optional*, defaults to `False`): TODO
    projector_dropout (`int`, *optional*, defaults to 0.0): TODO
    attention_dropout (`int`, *optional*, defaults to 0.0): TODO
    rope_theta (`int`, *optional*, defaults to 10000): TODO


def _get_gcc_major_version(path_to_gcc: str) -> int:
    logging.info('Running echo __GNUC__ | %s -E -P -', path_to_gcc)
    gcc_version_proc = subprocess.run([path_to_gcc, '-E', '-P', '-'], input='__GNUC__', check=True, capture_output=True, text=True)
    major_version = int(gcc_version_proc.stdout)
    logging.info('%s reports major version %s.', path_to_gcc, major_version)
    return major_version

Gets the major version of the gcc at `path_to_gcc`.

Args:
  path_to_gcc: Path to a gcc executable

Returns:
  The major version.


def __init__(self, x, name):
    self.x = x
    self.name = name

Construct DivideDelegateWithName.

Args:
  x: Tensor to use as left operand in operator overloads
  name: The name that is preferred for the op created.


def InsertAll(self, request, global_params=None):
    config = self.GetMethodConfig('InsertAll')
    return self._RunMethod(config, request, global_params=global_params)

Streams data into BigQuery one record at a time without needing to run a load job. Requires the WRITER dataset role.

Args:
  request: (BigqueryTabledataInsertAllRequest) input message
  global_params: (StandardQueryParameters, default: None) global arguments
Returns:
  (TableDataInsertAllResponse) The response message.


class DFineConfig(PretrainedConfig):
    model_type = 'd_fine'
    layer_types = ['basic', 'bottleneck']
    attribute_map = {'hidden_size': 'd_model', 'num_attention_heads': 'encoder_attention_heads'}

    def __init__(self, initializer_range=0.01, initializer_bias_prior_prob=None, layer_norm_eps=1e-05, batch_norm_eps=1e-05, backbone_config=None, backbone=None, use_pretrained_backbone=False, use_timm_backbone=False, freeze_backbone_batch_norms=True, backbone_kwargs=None, encoder_hidden_dim=256, encoder_in_channels=[512, 1024, 2048], feat_strides=[8, 16, 32], encoder_layers=1, encoder_ffn_dim=1024, encoder_attention_heads=8, dropout=0.0, activation_dropout=0.0, encode_proj_layers=[2], positional_encoding_temperature=10000, encoder_activation_function='gelu', activation_function='silu', eval_size=None, normalize_before=False, hidden_expansion=1.0, d_model=256, num_queries=300, decoder_in_channels=[256, 256, 256], decoder_ffn_dim=1024, num_feature_levels=3, decoder_n_points=4, decoder_layers=6, decoder_attention_heads=8, decoder_activation_function='relu', attention_dropout=0.0, num_denoising=100, label_noise_ratio=0.5, box_noise_scale=1.0, learn_initial_query=False, anchor_image_size=None, with_box_refine=True, is_encoder_decoder=True, matcher_alpha=0.25, matcher_gamma=2.0, matcher_class_cost=2.0, matcher_bbox_cost=5.0, matcher_giou_cost=2.0, use_focal_loss=True, auxiliary_loss=True, focal_loss_alpha=0.75, focal_loss_gamma=2.0, weight_loss_vfl=1.0, weight_loss_bbox=5.0, weight_loss_giou=2.0, weight_loss_fgl=0.15, weight_loss_ddf=1.5, eos_coefficient=0.0001, eval_idx=-1, layer_scale=1, max_num_bins=32, reg_scale=4.0, depth_mult=1.0, top_prob_values=4, lqe_hidden_dim=64, lqe_layers=2, decoder_offset_scale=0.5, decoder_method='default', up=0.5, **kwargs):
        self.initializer_range = initializer_range
        self.initializer_bias_prior_prob = initializer_bias_prior_prob
        self.layer_norm_eps = layer_norm_eps
        self.batch_norm_eps = batch_norm_eps
        if backbone_config is None and backbone is None:
            logger.info('`backbone_config` and `backbone` are `None`. Initializing the config with the default `HGNet-V2` backbone.')
            backbone_model_type = 'hgnet_v2'
            config_class = CONFIG_MAPPING[backbone_model_type]
            backbone_config = config_class(num_channels=3, embedding_size=64, hidden_sizes=[256, 512, 1024, 2048], depths=[3, 4, 6, 3], layer_type='bottleneck', hidden_act='relu', downsample_in_first_stage=False, downsample_in_bottleneck=False, out_features=None, out_indices=[2, 3, 4])
        elif isinstance(backbone_config, dict):
            backbone_model_type = backbone_config.pop('model_type')
            config_class = CONFIG_MAPPING[backbone_model_type]
            backbone_config = config_class.from_dict(backbone_config)
        verify_backbone_config_arguments(use_timm_backbone=use_timm_backbone, use_pretrained_backbone=use_pretrained_backbone, backbone=backbone, backbone_config=backbone_config, backbone_kwargs=backbone_kwargs)
        self.backbone_config = backbone_config
        self.backbone = backbone
        self.use_pretrained_backbone = use_pretrained_backbone
        self.use_timm_backbone = use_timm_backbone
        self.freeze_backbone_batch_norms = freeze_backbone_batch_norms
        self.backbone_kwargs = backbone_kwargs
        self.encoder_hidden_dim = encoder_hidden_dim
        self.encoder_in_channels = encoder_in_channels
        self.feat_strides = feat_strides
        self.encoder_attention_heads = encoder_attention_heads
        self.encoder_ffn_dim = encoder_ffn_dim
        self.dropout = dropout
        self.activation_dropout = activation_dropout
        self.encode_proj_layers = encode_proj_layers
        self.encoder_layers = encoder_layers
        self.positional_encoding_temperature = positional_encoding_temperature
        self.eval_size = eval_size
        self.normalize_before = normalize_before
        self.encoder_activation_function = encoder_activation_function
        self.activation_function = activation_function
        self.hidden_expansion = hidden_expansion
        self.d_model = d_model
        self.num_queries = num_queries
        self.decoder_ffn_dim = decoder_ffn_dim
        self.decoder_in_channels = decoder_in_channels
        self.num_feature_levels = num_feature_levels
        self.decoder_n_points = decoder_n_points
        self.decoder_layers = decoder_layers
        self.decoder_attention_heads = decoder_attention_heads
        self.decoder_activation_function = decoder_activation_function
        self.attention_dropout = attention_dropout
        self.num_denoising = num_denoising
        self.label_noise_ratio = label_noise_ratio
        self.box_noise_scale = box_noise_scale
        self.learn_initial_query = learn_initial_query
        self.anchor_image_size = anchor_image_size
        self.auxiliary_loss = auxiliary_loss
        self.with_box_refine = with_box_refine
        self.matcher_alpha = matcher_alpha
        self.matcher_gamma = matcher_gamma
        self.matcher_class_cost = matcher_class_cost
        self.matcher_bbox_cost = matcher_bbox_cost
        self.matcher_giou_cost = matcher_giou_cost
        self.use_focal_loss = use_focal_loss
        self.focal_loss_alpha = focal_loss_alpha
        self.focal_loss_gamma = focal_loss_gamma
        self.weight_loss_vfl = weight_loss_vfl
        self.weight_loss_bbox = weight_loss_bbox
        self.weight_loss_giou = weight_loss_giou
        self.weight_loss_fgl = weight_loss_fgl
        self.weight_loss_ddf = weight_loss_ddf
        self.eos_coefficient = eos_coefficient
        self.eval_idx = eval_idx
        self.layer_scale = layer_scale
        self.max_num_bins = max_num_bins
        self.reg_scale = reg_scale
        self.depth_mult = depth_mult
        self.decoder_offset_scale = decoder_offset_scale
        self.decoder_method = decoder_method
        self.top_prob_values = top_prob_values
        self.lqe_hidden_dim = lqe_hidden_dim
        self.lqe_layers = lqe_layers
        self.up = up
        if isinstance(self.decoder_n_points, list):
            if len(self.decoder_n_points) != self.num_feature_levels:
                raise ValueError(f'Length of decoder_n_points list ({len(self.decoder_n_points)}) must match num_feature_levels ({self.num_feature_levels}).')
        head_dim = self.d_model // self.decoder_attention_heads
        if head_dim * self.decoder_attention_heads != self.d_model:
            raise ValueError(f'Embedded dimension {self.d_model} must be divisible by decoder_attention_heads {self.decoder_attention_heads}')
        super().__init__(is_encoder_decoder=is_encoder_decoder, **kwargs)

    @property
    def num_attention_heads(self) -> int:
        return self.encoder_attention_heads

    @property
    def hidden_size(self) -> int:
        return self.d_model

    @classmethod
    def from_backbone_configs(cls, backbone_config: PretrainedConfig, **kwargs):
        """Instantiate a [`DFineConfig`] (or a derived class) from a pre-trained backbone model configuration and DETR model
        configuration.

            Args:
                backbone_config ([`PretrainedConfig`]):
                    The backbone configuration.

            Returns:
                [`DFineConfig`]: An instance of a configuration object
        """
        return cls(backbone_config=backbone_config, **kwargs)

This is the configuration class to store the configuration of a [`DFineModel`]. It is used to instantiate a D-FINE
model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
defaults will yield a similar configuration to that of D-FINE-X-COCO "[ustc-community/dfine-xlarge-coco"](https://huggingface.co/ustc-community/dfine-xlarge-coco").
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
documentation from [`PretrainedConfig`] for more information.

Args:
    initializer_range (`float`, *optional*, defaults to 0.01):
        The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
    initializer_bias_prior_prob (`float`, *optional*):
        The prior probability used by the bias initializer to initialize biases for `enc_score_head` and `class_embed`.
        If `None`, `prior_prob` computed as `prior_prob = 1 / (num_labels + 1)` while initializing model weights.
    layer_norm_eps (`float`, *optional*, defaults to 1e-05):
        The epsilon used by the layer normalization layers.
    batch_norm_eps (`float`, *optional*, defaults to 1e-05):
        The epsilon used by the batch normalization layers.
    backbone_config (`Dict`, *optional*, defaults to `RTDetrResNetConfig()`):
        The configuration of the backbone model.
    backbone (`str`, *optional*):
        Name of backbone to use when `backbone_config` is `None`. If `use_pretrained_backbone` is `True`, this
        will load the corresponding pretrained weights from the timm or transformers library. If `use_pretrained_backbone`
        is `False`, this loads the backbone's config and uses that to initialize the backbone with random weights.
    use_pretrained_backbone (`bool`, *optional*, defaults to `False`):
        Whether to use pretrained weights for the backbone.
    use_timm_backbone (`bool`, *optional*, defaults to `False`):
        Whether to load `backbone` from the timm library. If `False`, the backbone is loaded from the transformers
        library.
    freeze_backbone_batch_norms (`bool`, *optional*, defaults to `True`):
        Whether to freeze the batch normalization layers in the backbone.
    backbone_kwargs (`dict`, *optional*):
        Keyword arguments to be passed to AutoBackbone when loading from a checkpoint
        e.g. `{'out_indices': (0, 1, 2, 3)}`. Cannot be specified if `backbone_config` is set.
    encoder_hidden_dim (`int`, *optional*, defaults to 256):
        Dimension of the layers in hybrid encoder.
    encoder_in_channels (`list`, *optional*, defaults to `[512, 1024, 2048]`):
        Multi level features input for encoder.
    feat_strides (`List[int]`, *optional*, defaults to `[8, 16, 32]`):
        Strides used in each feature map.
    encoder_layers (`int`, *optional*, defaults to 1):
        Total of layers to be used by the encoder.
    encoder_ffn_dim (`int`, *optional*, defaults to 1024):
        Dimension of the "intermediate" (often named feed-forward) layer in decoder.
    encoder_attention_heads (`int`, *optional*, defaults to 8):
        Number of attention heads for each attention layer in the Transformer encoder.
    dropout (`float`, *optional*, defaults to 0.0):
        The ratio for all dropout layers.
    activation_dropout (`float`, *optional*, defaults to 0.0):
        The dropout ratio for activations inside the fully connected layer.
    encode_proj_layers (`List[int]`, *optional*, defaults to `[2]`):
        Indexes of the projected layers to be used in the encoder.
    positional_encoding_temperature (`int`, *optional*, defaults to 10000):
        The temperature parameter used to create the positional encodings.
    encoder_activation_function (`str`, *optional*, defaults to `"gelu"`):
        The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
        `"relu"`, `"silu"` and `"gelu_new"` are supported.
    activation_function (`str`, *optional*, defaults to `"silu"`):
        The non-linear activation function (function or string) in the general layer. If string, `"gelu"`,
        `"relu"`, `"silu"` and `"gelu_new"` are supported.
    eval_size (`Tuple[int, int]`, *optional*):
        Height and width used to computes the effective height and width of the position embeddings after taking
        into account the stride.
    normalize_before (`bool`, *optional*, defaults to `False`):
        Determine whether to apply layer normalization in the transformer encoder layer before self-attention and
        feed-forward modules.
    hidden_expansion (`float`, *optional*, defaults to 1.0):
        Expansion ratio to enlarge the dimension size of RepVGGBlock and CSPRepLayer.
    d_model (`int`, *optional*, defaults to 256):
        Dimension of the layers exclude hybrid encoder.
    num_queries (`int`, *optional*, defaults to 300):
        Number of object queries.
    decoder_in_channels (`list`, *optional*, defaults to `[256, 256, 256]`):
        Multi level features dimension for decoder
    decoder_ffn_dim (`int`, *optional*, defaults to 1024):
        Dimension of the "intermediate" (often named feed-forward) layer in decoder.
    num_feature_levels (`int`, *optional*, defaults to 3):
        The number of input feature levels.
    decoder_n_points (`int`, *optional*, defaults to 4):
        The number of sampled keys in each feature level for each attention head in the decoder.
    decoder_layers (`int`, *optional*, defaults to 6):
        Number of decoder layers.
    decoder_attention_heads (`int`, *optional*, defaults to 8):
        Number of attention heads for each attention layer in the Transformer decoder.
    decoder_activation_function (`str`, *optional*, defaults to `"relu"`):
        The non-linear activation function (function or string) in the decoder. If string, `"gelu"`,
        `"relu"`, `"silu"` and `"gelu_new"` are supported.
    attention_dropout (`float`, *optional*, defaults to 0.0):
        The dropout ratio for the attention probabilities.
    num_denoising (`int`, *optional*, defaults to 100):
        The total number of denoising tasks or queries to be used for contrastive denoising.
    label_noise_ratio (`float`, *optional*, defaults to 0.5):
        The fraction of denoising labels to which random noise should be added.
    box_noise_scale (`float`, *optional*, defaults to 1.0):
        Scale or magnitude of noise to be added to the bounding boxes.
    learn_initial_query (`bool`, *optional*, defaults to `False`):
        Indicates whether the initial query embeddings for the decoder should be learned during training
    anchor_image_size (`Tuple[int, int]`, *optional*):
        Height and width of the input image used during evaluation to generate the bounding box anchors. If None, automatic generate anchor is applied.
    with_box_refine (`bool`, *optional*, defaults to `True`):
        Whether to apply iterative bounding box refinement, where each decoder layer refines the bounding boxes
        based on the predictions from the previous layer.
    is_encoder_decoder (`bool`, *optional*, defaults to `True`):
        Whether the architecture has an encoder decoder structure.
    matcher_alpha (`float`, *optional*, defaults to 0.25):
        Parameter alpha used by the Hungarian Matcher.
    matcher_gamma (`float`, *optional*, defaults to 2.0):
        Parameter gamma used by the Hungarian Matcher.
    matcher_class_cost (`float`, *optional*, defaults to 2.0):
        The relative weight of the class loss used by the Hungarian Matcher.
    matcher_bbox_cost (`float`, *optional*, defaults to 5.0):
        The relative weight of the bounding box loss used by the Hungarian Matcher.
    matcher_giou_cost (`float`, *optional*, defaults to 2.0):
        The relative weight of the giou loss of used by the Hungarian Matcher.
    use_focal_loss (`bool`, *optional*, defaults to `True`):
        Parameter informing if focal focal should be used.
    auxiliary_loss (`bool`, *optional*, defaults to `True`):
        Whether auxiliary decoding losses (loss at each decoder layer) are to be used.
    focal_loss_alpha (`float`, *optional*, defaults to 0.75):
        Parameter alpha used to compute the focal loss.
    focal_loss_gamma (`float`, *optional*, defaults to 2.0):
        Parameter gamma used to compute the focal loss.
    weight_loss_vfl (`float`, *optional*, defaults to 1.0):
        Relative weight of the varifocal loss in the object detection loss.
    weight_loss_bbox (`float`, *optional*, defaults to 5.0):
        Relative weight of the L1 bounding box loss in the object detection loss.
    weight_loss_giou (`float`, *optional*, defaults to 2.0):
        Relative weight of the generalized IoU loss in the object detection loss.
    weight_loss_fgl (`float`, *optional*, defaults to 0.15):
        Relative weight of the fine-grained localization loss in the object detection loss.
    weight_loss_ddf (`float`, *optional*, defaults to 1.5):
        Relative weight of the decoupled distillation focal loss in the object detection loss.
    eos_coefficient (`float`, *optional*, defaults to 0.0001):
        Relative classification weight of the 'no-object' class in the object detection loss.
    eval_idx (`int`, *optional*, defaults to -1):
        Index of the decoder layer to use for evaluation. If negative, counts from the end
        (e.g., -1 means use the last layer). This allows for early prediction in the decoder
        stack while still training later layers.
    layer_scale (`float`, *optional*, defaults to `1.0`):
        Scaling factor for the hidden dimension in later decoder layers. Used to adjust the
        model capacity after the evaluation layer.
    max_num_bins (`int`, *optional*, defaults to 32):
        Maximum number of bins for the distribution-guided bounding box refinement.
        Higher values allow for more fine-grained localization but increase computation.
    reg_scale (`float`, *optional*, defaults to 4.0):
        Scale factor for the regression distribution. Controls the range and granularity
        of the bounding box refinement process.
    depth_mult (`float`, *optional*, defaults to 1.0):
        Multiplier for the number of blocks in RepNCSPELAN4 layers. Used to scale the model's
        depth while maintaining its architecture.
    top_prob_values (`int`, *optional*, defaults to 4):
        Number of top probability values to consider from each corner's distribution.
    lqe_hidden_dim (`int`, *optional*, defaults to 64):
        Hidden dimension size for the Location Quality Estimator (LQE) network.
    lqe_layers (`int`, *optional*, defaults to 2):
        Number of layers in the Location Quality Estimator MLP.
    decoder_offset_scale (`float`, *optional*, defaults to 0.5):
        Offset scale used in deformable attention.
    decoder_method (`str`, *optional*, defaults to `"default"`):
        The method to use for the decoder: `"default"` or `"discrete"`.
    up (`float`, *optional*, defaults to 0.5):
        Controls the upper bounds of the Weighting Function.


def _VerifyGeneratedGradients(grads, op: ops.Operation):
    if op.type == 'While' or op.type == 'StatelessWhile':
        return
    if len(grads) != len(op.inputs):
        raise ValueError(f'Num gradients {len(grads)} generated for op {op.node_def} do not match num inputs {len(op.inputs)}')

Verify that gradients are valid in number and type.

Args:
  grads: List of generated gradients.
  op: Operation for which the gradients where generated.

Raises:
  ValueError: if sizes of gradients and inputs don't match.
  TypeError: if type of any gradient is not valid for its input.


def calendar_month(self: EventSetOrNode, tz: Union[str, float, int]=0) -> EventSetOrNode:
    from temporian.core.operators.calendar.month import calendar_month
    return calendar_month(self, tz)

Obtains the month the timestamps in an
[`EventSet`][temporian.EventSet]'s sampling are in.

Features in the input are ignored, only the timestamps are used and
they must be unix timestamps (`is_unix_timestamp=True`).

Output feature contains numbers between 1 and 12.

By default, the timezone is UTC unless the `tz` argument is specified,
as an offset in hours or a timezone name. See
[`EventSet.calendar_hour()`][temporian.EventSet.calendar_hour] for an
example using timezones.

Usage example:
    ```python
    >>> a = tp.event_set(
    ...    timestamps=["2023-02-04", "2023-02-20", "2023-03-01", "2023-05-07"],
    ...    name='special_events'
    ... )
    >>> b = a.calendar_month()
    >>> b
    indexes: ...
    features: [('calendar_month', int32)]
    events:
        (4 events):
            timestamps: [...]
            'calendar_month': [2 2 3 5]
    ...

    ```

Args:
    tz: timezone name (see `pytz.all_timezones`) or UTC offset in hours.

Returns:
    EventSet with a single feature with the month each timestamp in
        `sampling` belongs to.