def forward(self, input_ids: torch.LongTensor=None, attention_mask: Optional[torch.Tensor]=None, position_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[List[torch.FloatTensor]]=None, inputs_embeds: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, pixel_values: Optional[torch.Tensor]=None, pixel_values_videos: Optional[torch.FloatTensor]=None, image_grid_thw: Optional[torch.LongTensor]=None, video_grid_thw: Optional[torch.LongTensor]=None, rope_deltas: Optional[torch.LongTensor]=None, cache_position: Optional[torch.LongTensor]=None, second_per_grid_ts: Optional[torch.Tensor]=None, **kwargs: Unpack[KwargsForCausalLM]) -> Union[Tuple, Qwen2_5_VLCausalLMOutputWithPast]:
    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
    outputs = self.model(input_ids=input_ids, pixel_values=pixel_values, pixel_values_videos=pixel_values_videos, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, second_per_grid_ts=second_per_grid_ts, position_ids=position_ids, attention_mask=attention_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=True, cache_position=cache_position, **kwargs)
    hidden_states = outputs[0]
    logits = self.lm_head(hidden_states)
    loss = None
    if labels is not None:
        loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size)
    return Qwen2_5_VLCausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, rope_deltas=outputs.rope_deltas)

labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
    Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
    config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
    (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
pixel_values_videos (`torch.FloatTensor` of shape `(seq_length, num_channels * temporal_size * image_size * image_size)):
    The tensors corresponding to the input videos. Pixel values can be obtained using
    [`AutoImageProcessor`]. See [`Qwen2VLImageProcessor.__call__`] for details. [`Qwen2_5_VLProcessor`] uses
    [`Qwen2VLImageProcessor`] for processing videos.
image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
    The temporal, height and width of feature shape of each image in LLM.
video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
    The temporal, height and width of feature shape of each video in LLM.
rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
    The rope index difference between sequence length and multimodal rope.
second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*):
    The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs.

Example:

```python
>>> from PIL import Image
>>> import requests
>>> from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration

>>> model = Qwen2_5_VLForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")
>>> processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")

>>> messages = [
    {
        "role": "user",
        "content": [
            {"type": "image"},
            {"type": "text", "text": "What is shown in this image?"},
        ],
    },
]
>>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
>>> image = Image.open(requests.get(url, stream=True).raw)

>>> text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
>>> inputs = processor(text=[text], images=[image], vision_infos=[vision_infos])

>>> # Generate
>>> generate_ids = model.generate(inputs.input_ids, max_length=30)
>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
"The image shows a street scene with a red stop sign in the foreground. In the background, there is a large red gate with Chinese characters ..."
```


def run(argv=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
    pipeline = test_pipeline
    if not test_pipeline:
        pipeline = beam.Pipeline(options=pipeline_options)
    model_handler = HuggingFacePipelineModelHandler(task=PipelineTask.QuestionAnswering, model=known_args.model_name, load_model_args={'framework': 'pt', 'revision': known_args.revision})
    if not known_args.input:
        text = pipeline | 'CreateSentences' >> beam.Create(['What does Apache Beam do?;Apache Beam enables batch and streaming data processing.', 'What is the capital of France?;The capital of France is Paris .', 'Where was beam summit?;Apache Beam Summit 2023 was in NYC.'])
    else:
        text = pipeline | 'ReadSentences' >> beam.io.ReadFromText(known_args.input)
    processed_text = text | 'PreProcess' >> beam.ParDo(preprocess) | 'SquadExample' >> beam.ParDo(create_squad_example)
    output = processed_text | 'RunInference' >> RunInference(KeyedModelHandler(model_handler)) | 'ProcessOutput' >> beam.ParDo(PostProcessor())
    _ = output | '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.
  save_main_session: Used for internal testing.
  test_pipeline: Used for internal testing.


class DacDecoderOutput(ModelOutput):
    audio_values: Optional[torch.FloatTensor] = None

Args:
    audio_values (`torch.FloatTensor`  of shape `(batch_size, input_length)`, *optional*):
        Decoded audio values, obtained using the decoder part of Dac.


def _find_resource_in_bundle(uri: str, bundle_json: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    for entry in bundle_json.get('entry', ()):
        resource = entry.get('resource', {})
        if resource.get('url') == uri:
            return resource
        elif resource.get('resourceType') == 'Bundle':
            bundled_resource = _find_resource_in_bundle(uri, resource)
            if bundled_resource is not None:
                return bundled_resource
    return None

Finds the JSON object for the resource with `uri` inside `bundle_json`.

Args:
  uri: The resource URI to search for.
  bundle_json: Parsed JSON object for the bundle to search

Returns:
  Parsed JSON object for the resource or `None` if it can not be found.


def inferred_steps(self):
    return self._inferred_steps

The inferred steps per epoch of the created `Dataset`.

This will be `None` in the case where:

(1) A `Dataset` of unknown cardinality was passed to the `DataHandler`, and
(2) `steps_per_epoch` was not provided, and
(3) The first epoch of iteration has not yet completed.

Returns:
  The inferred steps per epoch of the created `Dataset`.


def assert_no_entries_with_modulus_zero(x, message=None, name='assert_no_entries_with_modulus_zero'):
    with ops.name_scope(name, values=[x]):
        x = tensor_conversion.convert_to_tensor_v2_with_dispatch(x, name='x')
        dtype = x.dtype.base_dtype
        should_be_nonzero = math_ops.abs(x)
        zero = tensor_conversion.convert_to_tensor_v2_with_dispatch(0, dtype=dtype.real_dtype)
        return check_ops.assert_less(zero, should_be_nonzero, message=message)

Returns `Op` that asserts Tensor `x` has no entries with modulus zero.

Args:
  x:  Numeric `Tensor`, real, integer, or complex.
  message:  A string message to prepend to failure message.
  name:  A name to give this `Op`.

Returns:
  An `Op` that asserts `x` has no entries with modulus zero.


def reduce(self, reduce_op, value, axis):
    return super(CentralStorageStrategy, self).reduce(reduce_op, value, axis)

Reduce `value` across replicas.

Given a per-replica value returned by `run`, say a
per-example loss, the batch will be divided across all the replicas. This
function allows you to aggregate across replicas and optionally also across
batch elements.  For example, if you have a global batch size of 8 and 2
replicas, values for examples `[0, 1, 2, 3]` will be on replica 0 and
`[4, 5, 6, 7]` will be on replica 1. By default, `reduce` will just
aggregate across replicas, returning `[0+4, 1+5, 2+6, 3+7]`. This is useful
when each replica is computing a scalar or some other value that doesn't
have a "batch" dimension (like a gradient). More often you will want to
aggregate across the global batch, which you can get by specifying the batch
dimension as the `axis`, typically `axis=0`. In this case it would return a
scalar `0+1+2+3+4+5+6+7`.

If there is a last partial batch, you will need to specify an axis so
that the resulting shape is consistent across replicas. So if the last
batch has size 6 and it is divided into [0, 1, 2, 3] and [4, 5], you
would get a shape mismatch unless you specify `axis=0`. If you specify
`tf.distribute.ReduceOp.MEAN`, using `axis=0` will use the correct
denominator of 6. Contrast this with computing `reduce_mean` to get a
scalar value on each replica and this function to average those means,
which will weigh some values `1/8` and others `1/4`.

For Example:
```
strategy = tf.distribute.experimental.CentralStorageStrategy(
    compute_devices=['CPU:0', 'GPU:0'], parameter_device='CPU:0')
ds = tf.data.Dataset.range(10)
# Distribute that dataset
dist_dataset = strategy.experimental_distribute_dataset(ds)

with strategy.scope():
  @tf.function
  def train_step(val):
    # pass through
    return val

  # Iterate over the distributed dataset
  for x in dist_dataset:
    result = strategy.run(train_step, args=(x,))

result = strategy.reduce(tf.distribute.ReduceOp.SUM, result,
                         axis=None).numpy()
# result: array([ 4,  6,  8, 10])

result = strategy.reduce(tf.distribute.ReduceOp.SUM, result, axis=0).numpy()
# result: 28
```

Args:
  reduce_op: A `tf.distribute.ReduceOp` value specifying how values should
    be combined.
  value: A "per replica" value, e.g. returned by `run` to
    be combined into a single tensor.
  axis: Specifies the dimension to reduce along within each
    replica's tensor. Should typically be set to the batch dimension, or
    `None` to only reduce across replicas (e.g. if the tensor has no batch
    dimension).

Returns:
  A `Tensor`.


def __init__(self, topic, with_attributes=False, id_label=None, timestamp_attribute=None, expansion_service=None):
    self.params = WriteToPubsubSchema(topic=topic, id_label=id_label, timestamp_attribute=timestamp_attribute)
    self.expansion_service = expansion_service
    self.with_attributes = with_attributes

Initializes ``WriteToPubSub``.

Args:
  topic: Cloud Pub/Sub topic in the form "/topics/<project>/<topic>".
  with_attributes:
    True - input elements will be
    :class:`~apache_beam.io.gcp.pubsub.PubsubMessage` objects.
    False - input elements will be of type ``bytes`` (message
    data only).
  id_label: If set, will set an attribute for each Cloud Pub/Sub message
    with the given name and a unique value. This attribute can then be used
    in a ReadFromPubSub PTransform to deduplicate messages.
  timestamp_attribute: If set, will set an attribute for each Cloud Pub/Sub
    message with the given name and the message's publish time as the value.


def _get_python_exe_version(python_exe: list[str]):
    try:
        python_exe_version = subprocess.check_output(python_exe + ['-V'], stderr=subprocess.STDOUT).decode()
    except (subprocess.CalledProcessError, FileNotFoundError):
        return None
    return _parse_exe_version_string(python_exe_version)

Determine the major and minor version of given Python executable.

Arguments:
  python_exe: absolute path to the Python executable

Returns:
  Version as (major, minor) tuple, or None if it could not be determined.


def preprocess(self, images, do_resize: Optional[bool]=None, size: Optional[Dict[str, int]]=None, resample: Optional[PILImageResampling]=None, do_pad: Optional[bool]=None, padding_value: Optional[float]=None, padding_mode: Optional[str]=None, do_normalize: Optional[bool]=None, image_mean: Optional[float]=None, image_std: Optional[float]=None, do_rescale: Optional[bool]=None, rescale_factor: Optional[float]=None, patch_size: Optional[Dict[str, int]]=None, data_format: Optional[Union[str, ChannelDimension]]=ChannelDimension.FIRST, input_data_format: Optional[Union[str, ChannelDimension]]=None, return_tensors: Optional[TensorType]=None):
    do_resize = do_resize if do_resize is not None else self.do_resize
    size = size if size is not None else self.size
    resample = resample if resample is not None else self.resample
    do_pad = do_pad if do_pad is not None else self.do_pad
    do_rescale = do_rescale if do_rescale is not None else self.do_rescale
    rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
    do_normalize = do_normalize if do_normalize is not None else self.do_normalize
    image_mean = image_mean if image_mean is not None else self.image_mean
    image_std = image_std if image_std is not None else self.image_std
    padding_value = padding_value if padding_value is not None else self.padding_value
    padding_mode = padding_mode if padding_mode is not None else self.padding_mode
    do_rescale = do_rescale if do_rescale is not None else self.do_rescale
    rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
    patch_size = patch_size if patch_size is not None else self.patch_size
    if isinstance(images, list) and any((isinstance(elem, list) and len(elem) >= 2 for elem in images)):
        raise ValueError('Multiple images for a single sample are not yet supported.')
    batch_images = make_list_of_list_of_images(images)
    validate_preprocess_arguments(do_rescale=do_rescale, rescale_factor=rescale_factor, do_normalize=do_normalize, image_mean=image_mean, image_std=image_std, do_pad=do_pad, size_divisibility=size, do_resize=do_resize, size=size, resample=resample)
    batch_images = [[to_numpy_array(image) for image in images] for images in batch_images]
    if do_rescale and is_scaled_image(batch_images[0][0]):
        logger.warning_once('It looks like you are trying to rescale already rescaled images. If the input images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again.')
    if input_data_format is None:
        input_data_format = infer_channel_dimension_format(batch_images[0][0])
    original_image_sizes = [get_image_size(images[0], channel_dim=input_data_format) for images in batch_images]
    size = get_size_dict(size)
    if do_resize:
        batch_images = [[self.resize(image, size=size, input_data_format=input_data_format) for image in images] for images in batch_images]
    image_sizes = [get_image_size(images[0], channel_dim=input_data_format) for images in batch_images]
    image_unpadded_heights = [[image_size[0]] for image_size in image_sizes]
    image_unpadded_widths = [[image_size[1]] for image_size in image_sizes]
    image_scale_factors = [[resized_size[0] / original_size[0]] for original_size, resized_size in zip(original_image_sizes, image_sizes)]
    if do_pad:
        batch_images = [[self.pad_image(image, size=size, mode=padding_mode, constant_values=padding_value, input_data_format=input_data_format) for image in images] for images in batch_images]
    if do_rescale:
        batch_images = [[self.rescale(image, scale=rescale_factor, input_data_format=input_data_format) for image in images] for images in batch_images]
    if do_normalize:
        batch_images = [[self.normalize(image, mean=image_mean, std=image_std, input_data_format=input_data_format) for image in images] for images in batch_images]
    if data_format is not None:
        batch_images = [[to_channel_dimension_format(image, data_format, input_data_format) for image in images] for images in batch_images]
    data = {'images': batch_images, 'image_unpadded_heights': image_unpadded_heights, 'image_unpadded_widths': image_unpadded_widths, 'image_scale_factors': image_scale_factors}
    return FuyuBatchFeature(data=data, tensor_type=return_tensors)

Utility function to preprocess the images and extract necessary information about original formats.

Args:
    images (`ImageInput`):
        Images to preprocess. Expects a single image, a list or images or a list of lists of images. Pixel
        values range from 0 to 255, or between 0 and 1 if `do_rescale` is `False`.
    do_resize (`bool`, *optional*, defaults to `self.do_resize`):
        Whether to resize the image to `size`.
    size (`Dict[str, int]`, *optional*, defaults to `self.size`):
        Dictionary in the format `{"height": int, "width": int}` specifying the size of the output image.
    resample (`PILImageResampling`, *optional*, defaults to `self.resample`):
        `PILImageResampling` filter to use when resizing the image e.g. `PILImageResampling.BILINEAR`.
    do_pad (`bool`, *optional*, defaults to `self.do_pad`):
        Whether to pad the image to `size`.
    padding_value (`float`, *optional*, defaults to `self.padding_value`):
        The value to pad the image with.
    padding_mode (`str`, *optional*, defaults to `self.padding_mode`):
        The padding mode to use when padding the image.
    do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
        Whether to normalize the image.
    image_mean (`float`, *optional*, defaults to `self.image_mean`):
        The mean to use when normalizing the image.
    image_std (`float`, *optional*, defaults to `self.image_std`):
        The standard deviation to use when normalizing the image.
    do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
        Whether to rescale the image.
    rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
        The factor to use when rescaling the image.
    patch_size (`Dict[str, int]`, *optional*, defaults to `self.patch_size`):
        Dictionary in the format `{"height": int, "width": int}` specifying the size of the patches.
    return_tensors (`str` or `TensorType`, *optional*):
        The type of tensors to return. Can be one of:
        - Unset: Return a list of `np.ndarray`.
        - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
        - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
        - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
        - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
    data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
        The channel dimension format of the output image. Can be one of:
        - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
        - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
    input_data_format (`ChannelDimension` or `str`, *optional*):
        The channel dimension format for the input image. If unset, the channel dimension format is inferred
        from the input image. Can be one of:
        - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
        - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.


def encode(self, input_ids: jnp.ndarray, attention_mask: Optional[jnp.ndarray]=None, position_ids: Optional[jnp.ndarray]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, train: bool=False, params: Optional[dict]=None, dropout_rng: PRNGKey=None):
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    return_dict = return_dict if return_dict is not None else self.config.return_dict
    if attention_mask is None:
        attention_mask = jnp.ones_like(input_ids)
    if position_ids is None:
        batch_size, sequence_length = input_ids.shape
        position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))
    rngs = {}
    if dropout_rng is not None:
        rngs['dropout'] = dropout_rng

    def _encoder_forward(module, input_ids, attention_mask, position_ids, **kwargs):
        encode_module = module._get_encoder_module()
        return encode_module(input_ids, attention_mask, position_ids, **kwargs)
    return self.module.apply({'params': params or self.params}, input_ids=jnp.array(input_ids, dtype='i4'), attention_mask=jnp.array(attention_mask, dtype='i4'), position_ids=jnp.array(position_ids, dtype='i4'), output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, deterministic=not train, rngs=rngs, method=_encoder_forward)

Returns:

Example:

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

>>> model = FlaxMBartForConditionalGeneration.from_pretrained("facebook/mbart-large-cc25")
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/mbart-large-cc25")

>>> text = "My friends are cool but they eat too many carbs."
>>> inputs = tokenizer(text, max_length=1024, return_tensors="jax")
>>> encoder_outputs = model.encode(**inputs)
```


def __init__(self, sess_creator):
    self._sess_creator = sess_creator
    _WrappedSession.__init__(self, self._create_session())

Create a new `_RecoverableSession`.

The value returned by calling `sess_creator.create_session()` will be the
session wrapped by this recoverable session.

Args:
  sess_creator: A 'SessionCreator' to be wrapped by recoverable.


def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor]=None, layer_head_mask: Optional[torch.Tensor]=None, past_key_value: Optional[Tuple[torch.Tensor]]=None, output_attentions: Optional[bool]=False, use_cache: Optional[bool]=True) -> torch.Tensor:
    residual = hidden_states
    hidden_states = self.self_attn_layer_norm(hidden_states)
    self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
    hidden_states, self_attn_weights, present_key_value = self.self_attn(hidden_states=hidden_states, past_key_value=self_attn_past_key_value, attention_mask=attention_mask, layer_head_mask=layer_head_mask, output_attentions=output_attentions)
    hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
    hidden_states = residual + hidden_states
    residual = hidden_states
    hidden_states = self.final_layer_norm(hidden_states)
    hidden_states = self.activation_fn(self.fc1(hidden_states))
    hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
    hidden_states = self.fc2(hidden_states)
    hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
    hidden_states = residual + hidden_states
    outputs = (hidden_states,)
    if output_attentions:
        outputs += (self_attn_weights,)
    if use_cache:
        outputs += (present_key_value,)
    return outputs

Args:
    hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
    attention_mask (`torch.FloatTensor`): attention mask of size
        `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
    layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size `(attention_heads,)`.
    past_key_value (`Tuple(torch.FloatTensor)`): cached past key and value projection states
    output_attentions (`bool`, *optional*):
        Whether or not to return the attentions tensors of all attention layers. See `attentions` under
        returned tensors for more detail.


def __init__(self, session_root, watch_fn=None, thread_name_filter=None):
    self._session_root = session_root
    self._watch_fn = watch_fn
    self._thread_name_filter = thread_name_filter
    self._session_wrapper = None

Create a local debugger command-line interface (CLI) hook.

Args:
  session_root: See doc of
    `dumping_wrapper.DumpingDebugWrapperSession.__init__`.
  watch_fn: See doc of
    `dumping_wrapper.DumpingDebugWrapperSession.__init__`.
  thread_name_filter: Regular-expression white list for threads on which the
    wrapper session will be active. See doc of `BaseDebugWrapperSession` for
    more details.


def _create_scalar_select(lhs_result: _sql_data_types.StandardSqlExpression, rhs_result: _sql_data_types.StandardSqlExpression, scalar_check_op: str, sql_data_type: _sql_data_types.StandardSqlDataType, sql_alias: str):
    return _sql_data_types.Select(select_part=_sql_data_types.RawExpression(f'({lhs_result.as_operand()} {scalar_check_op} {rhs_result.as_operand()})', _sql_data_type=sql_data_type, _sql_alias=sql_alias), from_part=None, sql_dialect=_sql_data_types.SqlDialect.SPARK)

Construct a Spark SQL select statement for scalar values.

Args:
    lhs_result: The result of the left-hand side expression.
    rhs_result: The result of the right-hand side expression.
    scalar_check_op: The scalar operation to be applied ('=' or '!=').
    sql_data_type: The SQL data type for the result.
    sql_alias: The SQL alias for the result.

Returns:
    A compiled Spark SQL select statement.


def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
    dtype = q.dtype
    q = q.float()
    k = k.float()
    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)
    q_embed = q * cos + rotate_half(q) * sin
    k_embed = k * cos + rotate_half(k) * sin
    return (q_embed.to(dtype=dtype), k_embed.to(dtype=dtype))

Applies Rotary Position Embedding to the query and key tensors.

Args:
    q (`torch.Tensor`): The query tensor.
    k (`torch.Tensor`): The key tensor.
    cos (`torch.Tensor`): The cosine part of the rotary embedding.
    sin (`torch.Tensor`): The sine part of the rotary embedding.
    position_ids (`torch.Tensor`, *optional*):
        Deprecated and unused.
    unsqueeze_dim (`int`, *optional*, defaults to 1):
        The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
        sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
        that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
        k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
        cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
        the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
Returns:
    `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.


def getargspec(obj):
    if isinstance(obj, functools.partial):
        return _get_argspec_for_partial(obj)
    decorators, target = tf_decorator.unwrap(obj)
    spec = next((d.decorator_argspec for d in decorators if d.decorator_argspec is not None), None)
    if spec:
        return spec
    try:
        return _getargspec(target)
    except TypeError:
        pass
    if isinstance(target, type):
        try:
            return _getargspec(target.__init__)
        except TypeError:
            pass
        try:
            return _getargspec(target.__new__)
        except TypeError:
            pass
    return _getargspec(type(target).__call__)

TFDecorator-aware replacement for `inspect.getargspec`.

Note: `getfullargspec` is recommended as the python 2/3 compatible
replacement for this function.

Args:
  obj: A function, partial function, or callable object, possibly decorated.

Returns:
  The `ArgSpec` that describes the signature of the outermost decorator that
  changes the callable's signature, or the `ArgSpec` that describes
  the object if not decorated.

Raises:
  ValueError: When callable's signature can not be expressed with
    ArgSpec.
  TypeError: For objects of unsupported types.


def max_epochs():
    return _MAX_EPOCHS

Get the maximum number of epochs for any call to fit.

Retrieves the limit on the number of epochs set by
`keras.config.set_max_epochs` or the `KERAS_MAX_EPOCHS` environment
variable.

Returns:
    The integer limit on the number of epochs or `None`, if no limit has
    been set.


def _resize_output_size_rescale_to_max_len(height: int, width: int, min_len: Optional[int]=1, max_len: Optional[int]=None) -> Tuple[int, int]:
    max_len = max(height, width) if max_len is None else max_len
    aspect_ratio = width / height
    if width >= height:
        width = max_len
        height = int(width / aspect_ratio)
        if height % 2 != 0:
            height += 1
    elif height > width:
        height = max_len
        width = int(height * aspect_ratio)
        if width % 2 != 0:
            width += 1
    height = max(height, min_len)
    width = max(width, min_len)
    return (height, width)

Get the output size of the image after resizing given a dictionary specifying the max and min sizes.
Args:
    height (`int`):
        Height of the input image.
    width (`int`):
        Width of the input image.
    min_len (`int`, *optional*, defaults to 1):
        Minimum size of the output image.
    max_len (`int`, *optional*, defaults to the maximum size of the image):
        Maximum size of the output image.
Returns:
    The output size of the image after resizing.


def __init__(self, value_set_codes_table: Optional[bigquery.TableReference]=None, value_set_codes_definitions: Optional[fhir_package.FhirPackageManager]=None) -> None:
    self._value_set_codes_table = value_set_codes_table
    self._value_set_codes_definitions = value_set_codes_definitions
    self._use_resource_alias = None

Creates a BigQuerySqlInterpreter.

Args:
  value_set_codes_table: The name of the database table containing value set
    code definitions. Used when building SQL for memberOf expressions. If
    given, value set definitions needed for memberOf expressions will be
    retrieved from this table if they can not be found in
    `value_set_codes_definitions`. If neither this nor
    `value_set_codes_definitions` is given, no memberOf SQL will be
    generated.
  value_set_codes_definitions: A package manager containing value set
    definitions which can be used to build SQL for memberOf expressions.
    These value set definitions can be consulted in favor of using an
    external `value_set_codes_table`. If neither this nor
    `value_set_codes_definitions` is given, no memberOf SQL will be
    generated.


def _get_token_budget(self, question_tokens, max_length=None):
    return (max_length if max_length is not None else self.model_max_length) - self._question_encoding_cost(question_tokens)

Computes the number of tokens left for the table after tokenizing a question, taking into account the max
sequence length of the model.

Args:
    question_tokens (`List[String]`):
        List of question tokens. Returns: `int`: the number of tokens left for the table, given the model max
        length.


def __init__(self, boundaries, values, name=None):
    super(PiecewiseConstantDecay, self).__init__()
    if len(boundaries) != len(values) - 1:
        raise ValueError('The length of boundaries should be 1 less than the length of values')
    self.boundaries = boundaries
    self.values = values
    self.name = name

Piecewise constant from boundaries and interval values.

Args:
  boundaries: A list of `Tensor`s or `int`s or `float`s with strictly
    increasing entries, and with all elements having the same type as the
    optimizer step.
  values: A list of `Tensor`s or `float`s or `int`s that specifies the
    values for the intervals defined by `boundaries`. It should have one
    more element than `boundaries`, and all elements should have the same
    type.
  name: A string. Optional name of the operation. Defaults to
    'PiecewiseConstant'.

Raises:
  ValueError: if the number of elements in the lists do not match.


def build_inputs_with_special_tokens(self, token_ids_0: List[int], token_ids_1: Optional[List[int]]=None) -> List[int]:
    if token_ids_1 is None:
        return self.prefix_tokens + token_ids_0 + self.suffix_tokens
    return self.prefix_tokens + token_ids_0 + token_ids_1 + self.suffix_tokens

Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
adding special tokens. An NLLB sequence has the following format, where `X` represents the sequence:

- `input_ids` (for encoder) `X [eos, src_lang_code]`
- `decoder_input_ids`: (for decoder) `X [eos, tgt_lang_code]`

BOS is never used. Pairs of sequences are not the expected use case, but they will be handled without a
separator.

Args:
    token_ids_0 (`List[int]`):
        List of IDs to which the special tokens will be added.
    token_ids_1 (`List[int]`, *optional*):
        Optional second list of IDs for sequence pairs.

Returns:
    `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.


def _maybe_add_call_warning(self, node, full_name, name):
    warned = False
    if isinstance(node.func, ast.Attribute):
        warned = self._maybe_add_warning(node, '*.' + name)
    arg_warnings = self._get_applicable_dict('function_arg_warnings', full_name, name)
    variadic_args = uses_star_args_or_kwargs_in_call(node)
    for (kwarg, arg), (level, warning) in sorted(arg_warnings.items()):
        present, _ = get_arg_value(node, kwarg, arg) or variadic_args
        if present:
            warned = True
            warning_message = warning.replace('<function name>', full_name or name)
            template = '%s called with %s argument, requires manual check: %s'
            if variadic_args:
                template = '%s called with *args or **kwargs that may include %s, requires manual check: %s'
            self.add_log(level, node.lineno, node.col_offset, template % (full_name or name, kwarg, warning_message))
    return warned

Print a warning when specific functions are called with selected args.

The function _print_warning_for_function matches the full name of the called
function, e.g., tf.foo.bar(). This function matches the function name that
is called, as long as the function is an attribute. For example,
`tf.foo.bar()` and `foo.bar()` are matched, but not `bar()`.

Args:
  node: ast.Call object
  full_name: The precomputed full name of the callable, if one exists, None
    otherwise.
  name: The precomputed name of the callable, if one exists, None otherwise.

Returns:
  Whether an error was recorded.


class TFSeq2SeqSequenceClassifierOutput(ModelOutput):
    loss: tf.Tensor | None = None
    logits: Optional[tf.Tensor] = None
    past_key_values: List[tf.Tensor] | None = None
    decoder_hidden_states: Tuple[tf.Tensor] | None = None
    decoder_attentions: Tuple[tf.Tensor] | None = None
    cross_attentions: Tuple[tf.Tensor] | None = None
    encoder_last_hidden_state: tf.Tensor | None = None
    encoder_hidden_states: Tuple[tf.Tensor] | None = None
    encoder_attentions: Tuple[tf.Tensor] | None = None

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

Args:
    loss (`tf.Tensor` of shape `(1,)`, *optional*, returned when `label` is provided):
        Classification (or regression if config.num_labels==1) loss.
    logits (`tf.Tensor` of shape `(batch_size, config.num_labels)`):
        Classification (or regression if config.num_labels==1) scores (before SoftMax).
    past_key_values (`List[tf.Tensor]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
        List of `tf.Tensor` 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(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `tf.Tensor` (one for the output of the embeddings + 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(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
        Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, 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(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
        Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
        sequence_length)`
    encoder_last_hidden_state (`tf.Tensor` 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(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `tf.Tensor` (one for the output of the embeddings + 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(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
        Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
        sequence_length)`.

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


def ParseNSSwitchConf(nsswitch_filename):
    with open(nsswitch_filename, 'r') as nsswitch_file:
        nsswitch = {}
        map_re = re.compile('^([a-z]+): *(.*)$')
        for line in nsswitch_file:
            match = map_re.match(line)
            if match:
                sources = match.group(2).split()
                nsswitch[match.group(1)] = sources
    return nsswitch

Parse /etc/nsswitch.conf and return the sources for each map.

Args:
  nsswitch_filename: Full path to an nsswitch.conf to parse.  See manpage
    nsswitch.conf(5) for full details on the format expected.

Returns:
  a dictionary keyed by map names and containing a list of sources
  for each map.


def equity_leg_cashflows(forward_prices, spots, notional, dividends=None, dtype=None, name=None):
    name = name or 'equity_leg_cashflows'
    with tf.name_scope(name):
        forward_prices = tf.convert_to_tensor(forward_prices, dtype=dtype, name='forward_prices')
        dtype = dtype or forward_prices.dtype
        spots = tf.convert_to_tensor(spots, dtype=dtype, name='spots')
        notional = tf.convert_to_tensor(notional, dtype=dtype, name='notional')
        dividends = 0 if dividends is None else dividends
        dividends = tf.convert_to_tensor(dividends, dtype=dtype, name='dividends')
        spots_expand = tf.expand_dims(spots, axis=-1)
        forward_prices = tf.concat([spots_expand, forward_prices], axis=-1)
        return tf.math.divide_no_nan(notional * (forward_prices[..., 1:] - forward_prices[..., :-1]) + dividends, forward_prices[..., :-1])

Computes cashflows for a batch of equity legs.

Equity cashflows are defined as a total equity return between pay dates, say,
`T_1, ..., T_n`. Let `S_i` represent the value of the equity at time `T_i` and
`d_i` be a discrete dividend paid at this time. Then the the payment at time
`T_i` is defined as `(S_i - S_{i - 1}) / S_{i-1} + d_i`. The value of
the cashflow is then the discounted sum of the paments. See, e.g., [1] for the
reference.

#### Example
```python
notional = 10000
forward_prices = [[110, 120, 140], [210, 220, 240]]
spots = [100, 200]
dividends = [[1, 1, 1], [2, 2, 2]]
equity_leg_cashflows(forward_prices, spots, notional, dividends,
                     dtype=tf.float64)
# Expected:
#  [[1000.01, 909.1, 1666.675],
#   [ 500.01, 476.2, 909.1]]
```

Args:
  forward_prices: A real `Tensor` of shape `batch_shape + [num_cashflows]`,
    where `num_cashflows` is the number of cashflows for each batch element.
    Equity forward prices at leg reset times.
  spots:  A `Tensor` of the same `dtype` as `forward_prices` and of
    shape compatible with `batch_shape`. Spot prices for each batch element
  notional: A `Tensor` of the same `dtype` as `forward_prices` and of
    compatible shape. Notional amount for each cashflow.
  dividends:  A `Tensor` of the same `dtype` as `forward_prices` and of
    compatible shape. Discrete dividends paid at the leg reset times.
    Default value: None which maps to zero dividend.
  dtype: `tf.Dtype`. If supplied the dtype for the input and output `Tensor`s.
    Default value: None which maps to the default dtype inferred from
    `forward_prices`.
  name: Python str. The name to give to the ops created by this function.
    Default value: None which maps to 'equity_leg_cashflows'.

Returns:
  A `Tensor` of the same `dtype` as `forward_prices` and of shape
  `batch_shape + [num_cashflows]`.

#### References
[1] Don M. Chance and Don R Rich,
  The Pricing of Equity Swaps and Swaptions, 1998
  https://jod.pm-research.com/content/5/4/19


def _dequeue_return_value(self, tensors):
    if self._names:
        return {n: tensors[i] for i, n in enumerate(self._names)}
    elif len(tensors) == 1:
        return tensors[0]
    else:
        return tensors

Return the value to return from a dequeue op.

If the queue has names, return a dictionary with the
names as keys.  Otherwise return either a single tensor
or a list of tensors depending on the length of `tensors`.

Args:
  tensors: List of tensors from the dequeue op.

Returns:
  A single tensor, a list of tensors, or a dictionary
  of tensors.


def __getitem__(self, index):
    rank = self.rank
    if isinstance(index, slice):
        if index.step is not None and index.step != 1:
            raise IndexError('Cannot stride through a shape')
        start = index.start
        stop = index.stop
        if start is None:
            start = 0
        start = _fix_start_index(start, rank, self.num_row_partitions)
        stop = _fix_stop_index(stop, rank)
        return self._slice_shape(start, stop)
    elif isinstance(index, int):
        if index < 0:
            if rank is None:
                raise ValueError('Rank must be known to use __getitem__ with a negative index.')
            return self._dimension(rank + index)
        return self._dimension(index)
    else:
        raise TypeError('Argument is not an int or a slice')

Returns a dimension or a slice of the shape.

Ragged shapes can have ragged dimensions that depend upon other dimensions.
Therefore, if you ask for a dimension that is ragged, this function returns
a ValueError. For similar reasons, if a slice is selected that includes
a ragged dimension without including the zero dimension, then this fails.

Any slice that does not start at zero will return a shape
with num_row_partitions == 0.

Args:
  index: the index: can be an int or a slice.

Raises:
  IndexError: if the index is not in range.
  ValueError: if the rank is unknown, or a ragged rank is requested
  incorrectly.


def save(self, output_saved_model_dir, save_gpu_specific_engines=True, options=None):
    assert self._converted
    if trt_utils.is_experimental_feature_activated('remove_native_segments'):
        logging.info("'remove_native_segments' experimental feature is enabled during saving of converted SavedModel.")
        self._converted_func = _remove_native_segments(self._converted_func)
        self._converted_graph_def = self._converted_func.graph.as_graph_def()
    if self._need_calibration and (not self._calibrated):
        raise RuntimeError('A model that requires INT8 calibration has to be built before saving it. Call build() to build and calibrate the TensorRT engines.')
    engine_asset_dir = tempfile.mkdtemp()
    resource_map = {}

    def _serialize_and_track_engine(node):
        """Serialize TRT engines in the cache and track them."""
        canonical_engine_name = _get_canonical_engine_name(node.name)
        if canonical_engine_name in resource_map:
            return
        filename = os.path.join(engine_asset_dir, 'trt-serialized-engine.' + canonical_engine_name)
        try:
            gen_trt_ops.serialize_trt_resource(resource_name=canonical_engine_name, filename=filename, delete_resource=True, save_gpu_specific_engines=save_gpu_specific_engines)
        except errors.NotFoundError:
            logging.info('Could not find %s in TF-TRT cache. This can happen if build() is not called, which means TensorRT engines will be built and cached at runtime.', canonical_engine_name)
            return
        resource_map[canonical_engine_name] = _TRTEngineResource(canonical_engine_name, filename, self._conversion_params.maximum_cached_engines)
    self._for_each_trt_node(self._converted_graph_def, _serialize_and_track_engine)
    trackable = autotrackable.AutoTrackable() if self.freeze else self._saved_model
    trackable.trt_engine_resources = resource_map
    if not self._conversion_params.allow_build_at_runtime:

        def _reset_allow_build_at_runtime(node):
            node.attr['_allow_build_at_runtime'].b = False
        self._for_each_trt_node(self._converted_graph_def, _reset_allow_build_at_runtime)
        reset_converted_func = wrap_function.function_from_graph_def(self._converted_graph_def, [tensor.name for tensor in self._converted_func.inputs], [tensor.name for tensor in self._converted_func.outputs])
        reset_converted_func.graph.structured_outputs = nest.pack_sequence_as(self._converted_func.graph.structured_outputs, reset_converted_func.graph.structured_outputs)
        reset_converted_func.graph.structured_input_signature = self._converted_func.structured_input_signature
        self._converted_func = reset_converted_func
    signatures = {self._input_saved_model_signature_key: self._converted_func}
    save.save(trackable, output_saved_model_dir, signatures, options=options)

Save the converted SavedModel.

Args:
  output_saved_model_dir: directory to saved the converted SavedModel.
  save_gpu_specific_engines: whether to save TRT engines that have been
    built. When True, all engines are saved and when False, the engines
    are not saved and will be rebuilt at inference time. By using
    save_gpu_specific_engines=False after doing INT8 calibration, inference
    can be done on different GPUs than the GPU that the model was calibrated
    and saved on.
  options: `tf.saved_model.SaveOptions` object for configuring save options.
Raises:
  RuntimeError: if the needed calibration hasn't been done.


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.LongTensor]=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, BaseModelOutput]]=None, past_key_values: Optional[List[torch.FloatTensor]]=None, inputs_embeds: Optional[torch.Tensor]=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) -> Union[Tuple[torch.FloatTensor], 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)

    BlenderbotSmall uses the `bos_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, BlenderbotSmallModel

>>> model = BlenderbotSmallModel.from_pretrained("facebook/blenderbot_small-90M")
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/blenderbot_small-90M")

>>> inputs = tokenizer("Studies have been shown that owning a dog is good for you", return_tensors="pt")
>>> decoder_inputs = tokenizer("Studies show that", return_tensors="pt")  # Batch size 1
>>> 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, 3, 512]
```


def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: Optional[torch.LongTensor]=None):
    pixel_values = pixel_values.type(self.visual.dtype)
    image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
    return image_embeds

Encodes images into continuous embeddings that can be forwarded to the language model.

Args:
    pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
        The tensors corresponding to the input images.
    image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
        The temporal, height and width of feature shape of each image in LLM.


def forward(self, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.Tensor]=None, position_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[List[torch.FloatTensor]]=None, inputs_embeds: Optional[torch.FloatTensor]=None, pixel_values: Optional[torch.FloatTensor]=None, pixel_attention_mask: Optional[torch.BoolTensor]=None, image_hidden_states: 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, cache_position: Optional[torch.LongTensor]=None, return_dict: Optional[bool]=None, logits_to_keep: Union[int, torch.Tensor]=0, **kwargs: Unpack[KwargsForCausalLM]) -> Union[Tuple, Idefics3CausalLMOutputWithPast]:
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, pixel_values=pixel_values, pixel_attention_mask=pixel_attention_mask, image_hidden_states=image_hidden_states, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, cache_position=cache_position, return_dict=True, **kwargs)
    hidden_states = outputs[0]
    slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
    logits = self.lm_head(hidden_states[:, slice_indices, :])
    loss = None
    if labels is not None:
        loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs)
    return Idefics3CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, image_hidden_states=outputs.image_hidden_states)

pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*):
    Mask to avoid performing attention on padding pixel indices.
image_hidden_states (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
    The hidden states of the image encoder after modality projection.
labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
    Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
    config.vocab_size]` or `model.image_token_id` (where `model` is your instance of `Idefics3ForConditionalGeneration`).
    Tokens with indices set to `model.image_token_id` are ignored (masked), the loss is only
    computed for the tokens with labels in `[0, ..., config.vocab_size]`.

Example:

```python
>>> import requests
>>> import torch
>>> from PIL import Image
>>> from io import BytesIO

>>> from transformers import AutoProcessor, AutoModelForVision2Seq
>>> from transformers.image_utils import load_image

>>> # Note that passing the image urls (instead of the actual pil images) to the processor is also possible
>>> image1 = load_image("https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg")
>>> image2 = load_image("https://cdn.britannica.com/59/94459-050-DBA42467/Skyline-Chicago.jpg")
>>> image3 = load_image("https://cdn.britannica.com/68/170868-050-8DDE8263/Golden-Gate-Bridge-San-Francisco.jpg")

>>> processor = AutoProcessor.from_pretrained("HuggingFaceM4/Idefics3-8B-Llama3")
>>> model = AutoModelForVision2Seq.from_pretrained("HuggingFaceM4/Idefics3-8B-Llama3", torch_dtype=torch.bfloat16, device_map="auto")

>>> # Create inputs
>>> messages = [
...     {
...         "role": "user",
...         "content": [
...             {"type": "image"},
...             {"type": "text", "text": "In this image, we can see the city of New York, and more specifically the Statue of Liberty."},
...             {"type": "image"},
...             {"type": "text", "text": "What can we see in this image?"},
...         ]
...     },
...     {
...         "role": "user",
...         "content": [
...             {"type": "image"},
...             {"type": "text", "text": "In which city is that bridge located?"},
...         ]
...     }
... ]

>>> prompts = [processor.apply_chat_template([message], add_generation_prompt=True) for message in messages]
>>> images = [[image1, image2], [image3]]
>>> inputs = processor(text=prompts, images=images, padding=True, return_tensors="pt").to(model.device)

>>> # Generate
>>> generated_ids = model.generate(**inputs, max_new_tokens=256)
>>> generated_texts = processor.batch_decode(generated_ids, skip_special_tokens=True)

>>> print(generated_texts[0])
Assistant: There are buildings, trees, lights, and water visible in this image.

>>> print(generated_texts[1])
Assistant: The bridge is in San Francisco.
```


def build_inputs_with_special_tokens(self, token_ids_0: List[int], token_ids_1: Optional[List[int]]=None) -> List[int]:
    sep = [self.sep_token_id]
    cls = [self.cls_token_id]
    if token_ids_1 is None:
        return cls + token_ids_0 + sep
    return cls + token_ids_0 + sep + token_ids_1 + sep

Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
adding special tokens. An FNet sequence has the following format:

- single sequence: `[CLS] X [SEP]`
- pair of sequences: `[CLS] A [SEP] B [SEP]`

Args:
    token_ids_0 (`List[int]`):
        List of IDs to which the special tokens will be added
    token_ids_1 (`List[int]`, *optional*):
        Optional second list of IDs for sequence pairs.

Returns:
    `List[int]`: list of [input IDs](../glossary#input-ids) with the appropriate special tokens.


def stateless_random_jpeg_quality(image, min_jpeg_quality, max_jpeg_quality, seed):
    if min_jpeg_quality < 0 or max_jpeg_quality < 0 or min_jpeg_quality > 100 or (max_jpeg_quality > 100):
        raise ValueError('jpeg encoding range must be between 0 and 100.')
    if min_jpeg_quality >= max_jpeg_quality:
        raise ValueError('`min_jpeg_quality` must be less than `max_jpeg_quality`.')
    jpeg_quality = stateless_random_ops.stateless_random_uniform(shape=[], minval=min_jpeg_quality, maxval=max_jpeg_quality, seed=seed, dtype=dtypes.int32)
    return adjust_jpeg_quality(image, jpeg_quality)

Deterministically radomize jpeg encoding quality for inducing jpeg noise.

Guarantees the same results given the same `seed` independent of how many
times the function is called, and independent of global seed settings (e.g.
`tf.random.set_seed`).

`min_jpeg_quality` must be in the interval `[0, 100]` and less than
`max_jpeg_quality`.
`max_jpeg_quality` must be in the interval `[0, 100]`.

Usage Example:

>>> x = tf.constant([[[1, 2, 3],
...                   [4, 5, 6]],
...                  [[7, 8, 9],
...                   [10, 11, 12]]], dtype=tf.uint8)
>>> seed = (1, 2)
>>> tf.image.stateless_random_jpeg_quality(x, 75, 95, seed)
<tf.Tensor: shape=(2, 2, 3), dtype=uint8, numpy=
array([[[ 0,  4,  5],
        [ 1,  5,  6]],
       [[ 5,  9, 10],
        [ 5,  9, 10]]], dtype=uint8)>

Args:
  image: 3D image. Size of the last dimension must be 1 or 3.
  min_jpeg_quality: Minimum jpeg encoding quality to use.
  max_jpeg_quality: Maximum jpeg encoding quality to use.
  seed: A shape [2] Tensor, the seed to the random number generator. Must have
    dtype `int32` or `int64`. (When using XLA, only `int32` is allowed.)

Returns:
  Adjusted image(s), same shape and DType as `image`.

Raises:
  ValueError: if `min_jpeg_quality` or `max_jpeg_quality` is invalid.


def _VerifyExplicitPaddings(self, tensor_in_sizes, filter_in_sizes, strides, padding, data_format, dtype, use_gpu, op_name, dilations=(1, 1), test_grappler_layout_optimizer=False, tol=1e-05):
    input_tensor = self._CreateNumpyTensor(tensor_in_sizes)
    filter_tensor = self._CreateNumpyTensor(filter_in_sizes)
    input_tensor = array_ops.pad(input_tensor, [(0, 0)] + padding + [(0, 0)])
    dilations = list(dilations)
    conv2d_result = nn_ops.conv2d(input_tensor, filter_tensor, [1] + list(strides) + [1], 'VALID', dilations=[1] + dilations + [1])
    expected = list(self.evaluate(array_ops.reshape(conv2d_result, [-1])))
    self._VerifyValuesParameters(tensor_in_sizes, filter_in_sizes, strides, padding, expected, data_format, dtype, use_gpu, op_name, dilations, test_grappler_layout_optimizer=test_grappler_layout_optimizer, tol=tol)

Verifies Conv2D with explicit padding generates correct values.

It does this by comparing with Conv2D without explicit padding. This
function assumes Conv2D without explicit padding works correctly.

Args:
  tensor_in_sizes: Input tensor dimensions in [batch, input_rows,
    input_cols, input_depth].
  filter_in_sizes: Filter tensor dimensions in [kernel_rows, kernel_cols,
    input_depth, output_depth].
  strides: [row_stride, col_stride] for the convolution;
  padding: Explicit padding amounts.
  data_format: "NCHW" or "NHWC"
  dtype: data type to perform test
  use_gpu: True if testing on the GPU
  op_name: "Conv" or "Conv2D"
  dilations: Dilation values
  test_grappler_layout_optimizer: If True, allow the Grappler layout
    optimizer to run, which turns NHWC Conv2Ds on the GPU to NCHW Conv2Ds.
  tol: The absolute and relative tolerance.


def _compute_new_attention_mask(hidden_states: torch.Tensor, seq_lens: torch.Tensor):
    batch_size, mask_seq_len = hidden_states.shape[:2]
    indices = torch.arange(mask_seq_len, device=seq_lens.device).expand(batch_size, -1)
    bool_mask = indices >= seq_lens.unsqueeze(1).expand(-1, mask_seq_len)
    mask = hidden_states.new_ones((batch_size, mask_seq_len))
    mask = mask.masked_fill(bool_mask, 0)
    return mask

Computes an attention mask of the form `(batch, seq_len)` with an attention for each element in the batch that
stops at the corresponding element in `seq_lens`.
Args:
    hidden_states (`torch.FloatTensor` of shape `(batch, seq_len, *)`):
        The sequences to mask, where `*` is any number of sequence-specific dimensions including none.
    seq_lens (`torch.Tensor` of shape `(batch)`:
        Each element represents the length of the sequence at the same index in `hidden_states`
Returns:
    `torch.FloatTensor`: The float attention mask of shape `(batch, seq_len)`


def get_module_file(self, namespace, module, version):
    module_parts = module.split('.')
    module_path = path_utils.join(*module_parts)
    paths = []
    if namespace == 'stdlib':
        path = path_utils.join(namespace, module_path)
        if self._is_module_in_typeshed(module_parts, version) or path in self.missing:
            paths.append(path)
    elif namespace == 'third_party':
        for package in sorted(self._third_party_packages[module_parts[0]]):
            paths.append(path_utils.join('stubs', package, module_path))
    for path_rel in paths:
        if path_rel in self.missing:
            relpath = path_utils.join('nonexistent', path_rel + '.pyi')
            return (relpath, builtin_stubs.DEFAULT_SRC)
        for path in [path_utils.join(path_rel, '__init__.pyi'), path_rel + '.pyi']:
            try:
                name, src = self._store.load_file(path)
                return (name, src)
            except OSError:
                pass
    raise OSError(f"Couldn't find {module}")

Get the contents of a typeshed .pyi file.

Arguments:
  namespace: selects a top-level directory within typeshed/ Allowed values
    are "stdlib" and "third_party". "third_party" corresponds to the the
    typeshed/stubs/ directory.
  module: module name (e.g., "sys" or "__builtins__"). Can contain dots, if
    it's a submodule. Package names should omit the "__init__" suffix (e.g.,
    pass in "os", not "os.__init__").
  version: The Python version. (major, minor)

Returns:
  A tuple with the filename and contents of the file
Raises:
  IOError: if file not found


def from_dict(cls, feature_extractor_dict: dict[str, Any], **kwargs) -> PreTrainedFeatureExtractor:
    return_unused_kwargs = kwargs.pop('return_unused_kwargs', False)
    to_remove = []
    for key, value in kwargs.items():
        if key in feature_extractor_dict:
            feature_extractor_dict[key] = value
            to_remove.append(key)
    for key in to_remove:
        kwargs.pop(key, None)
    feature_extractor = cls(**feature_extractor_dict)
    logger.info(f'Feature extractor {feature_extractor}')
    if return_unused_kwargs:
        return (feature_extractor, kwargs)
    else:
        return feature_extractor

Instantiates a type of [`~feature_extraction_utils.FeatureExtractionMixin`] from a Python dictionary of
parameters.

Args:
    feature_extractor_dict (`Dict[str, Any]`):
        Dictionary that will be used to instantiate the feature extractor object. Such a dictionary can be
        retrieved from a pretrained checkpoint by leveraging the
        [`~feature_extraction_utils.FeatureExtractionMixin.to_dict`] method.
    kwargs (`Dict[str, Any]`):
        Additional parameters from which to initialize the feature extractor object.

Returns:
    [`~feature_extraction_utils.FeatureExtractionMixin`]: The feature extractor object instantiated from those
    parameters.


def __init__(self, functions, inference_args, input_tangents, tape_watching):
    self._functions = functions
    self._inference_args = inference_args
    self._input_tangents = input_tangents
    self._tape_watching = tape_watching

Collects information about the function call.

Args:
  functions: An object which produces forward and backward functions, either
    a _DelayedRewriteGradientFunctions or a _TapeGradientFunctions object.
  inference_args: A flat list of Tensors, arguments to the inference
    function.
  input_tangents: A flat list of Tensors, jvps associated with
    `inference_args`.
  tape_watching: Boolean, with True indicating that recording is necessary.


def forward(self, input, padding_mask: Optional[torch.Tensor]=None, causal_mask: Optional[torch.Tensor]=None, past_key_values: Optional[Tuple[torch.Tensor]]=None, output_attentions=False, use_cache=False):
    seq_len, bsz, embed_dim = input.size()
    if embed_dim != self.config.hidden_size:
        raise ValueError(f'Input embedding dimension should be {self.config.hidden_size}; received {embed_dim}')
    residual = input
    if self.config.normalize_before_mega:
        input = self.norm(input)
    value = self.activation(self.v_proj(input))
    if self.config.is_decoder and past_key_values is not None:
        if seq_len > 1:
            raise ValueError(f'Incremental decoding only supports self sequence length of 1; received {seq_len}')
        prev_self_key, prev_self_value, prev_ema_state = past_key_values[0:3]
    else:
        prev_self_key = prev_self_value = prev_ema_state = None
    ema_out, updated_ema_state = self.ema_gate(input, attention_mask=padding_mask, prev_state=prev_ema_state, use_cache=use_cache)
    ema_out = self.dropout(ema_out)
    base = self.mx_proj(ema_out)
    residual_weight, query_key_gates, intermediate_state = torch.split(base, [self.config.hidden_size, self.config.shared_representation_size + self.config.intermediate_size, self.config.hidden_size], dim=-1)
    residual_weight = torch.sigmoid(residual_weight)
    query_key_gates = F.silu(query_key_gates)
    query_key, attention_gate = torch.split(query_key_gates, [self.config.shared_representation_size, self.config.intermediate_size], dim=-1)
    query_key = query_key.unsqueeze(2) * self.qk_weight + self.qk_bias
    query, key = torch.unbind(query_key, dim=2)
    query = query.transpose(0, 1)
    key = key.transpose(0, 1)
    value = value.transpose(0, 1)
    if self.config.is_decoder:
        if prev_self_key is not None:
            key = torch.cat([prev_self_key, key], dim=1)
        if prev_self_value is not None:
            value = torch.cat([prev_self_value, value], dim=1)
        if not self.config.use_chunking:
            updated_self_key = key
            updated_self_value = value
        else:
            curr_len = key.size(1) % self.config.chunk_size
            if curr_len == 0:
                updated_self_key = None
                updated_self_value = None
            else:
                updated_self_key = key
                updated_self_value = value
    ctx_len = key.size(1)
    if not self.config.use_chunking:
        query = query.unsqueeze(1)
        key = key.unsqueeze(1)
        value = value.unsqueeze(1)
        if padding_mask is not None:
            padding_mask = padding_mask.unsqueeze(1)
    else:
        if seq_len < self.config.chunk_size:
            query = query.unsqueeze(1)
        else:
            n_chunks = seq_len // self.config.chunk_size
            query = query.reshape(bsz, n_chunks, self.config.chunk_size, self.config.shared_representation_size)
        if ctx_len < self.config.chunk_size:
            key = key.unsqueeze(1)
            value = value.unsqueeze(1)
            if padding_mask is not None:
                padding_mask = padding_mask.unsqueeze(1)
        else:
            n_chunks = ctx_len // self.config.chunk_size
            key = key.reshape(bsz, n_chunks, self.config.chunk_size, self.config.shared_representation_size)
            value = value.reshape(bsz, n_chunks, self.config.chunk_size, self.config.intermediate_size)
            if padding_mask is not None:
                padding_mask = padding_mask.view(bsz, n_chunks, self.config.chunk_size)
    if padding_mask is not None and padding_mask.dim() == 0:
        padding_mask = None
    attn_weights = self.attention_function(query, key, padding_mask=padding_mask, causal_mask=causal_mask)
    value = self.hidden_dropout(value, batch_first=True)
    kernel = self.attention_dropout(attn_weights)
    weighted_self_output = torch.matmul(kernel, value).view(bsz, seq_len, self.config.intermediate_size).transpose(0, 1)
    weighted_self_output = self.activation(intermediate_state + self.h_proj(weighted_self_output * attention_gate))
    weighted_self_output = self.dropout(weighted_self_output)
    out = torch.addcmul(residual, residual_weight, weighted_self_output - residual)
    if not self.config.normalize_before_mega:
        out = self.norm(out)
    return_values = (out, attn_weights) if output_attentions else (out,)
    if self.config.is_decoder:
        return_values = return_values + (updated_self_key, updated_self_value, updated_ema_state)
    return return_values

Mega's self-attention block, which combines multi-headed EMA with traditional self-attention

Args:
    input (`torch.Tensor` of shape `(sequence_length, batch_size, hidden_size)`):
        Hidden states to be updated by Mega's self-attention
    padding_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
        Indicates which inputs are to be ignored due to padding, where elements are either 1 for *not masked*
        or 0 for *masked*
    causal_mask (`torch.LongTensor` of shape `(sequence_length, sequence_length)`, *optional*):
        Indicates which inputs are to be ignored due to causal attention, where elements are either 1 for *not
        masked* or 0 for *masked*
    past_key_values (`tuple(torch.Tensor)`, *optional*):
        The hidden states returned from the previous timestep during incremental decoding; expects that
        self-attention key, value, and EMA states are the first 3 entries in the tuple
    output_attentions (`bool`, default `False`):
        Whether to return self-attention weights
    use_cache (`bool`, default `False`):
        Whether to perform incremental decoding; uses `past_key_values` as prior state, and returns the updated
        states for use in the next step

Returns:
    `tuple(torch.FloatTensor)` containing various elements depending on configuration ([`MegaConfig`]) and
    inputs:
    - **hidden_states** (`torch.FloatTensor` of shape `(sequence_length, batch_size, hidden_size)`) -- Hidden
      states from target sequence updated by Mega's self-attention
    - **attn_weights** (*optional*, returned when `output_attentions=True`) `torch.FloatTensor` of shape
      `(batch_size, 1, sequence_length, sequence_length)` -- The self-attention weights corresponding to how
      each token in the input sequence attends to every other token
    - **self_key** (*optional*, returned when `use_cache=True`) `torch.FloatTensor` of shape `(batch_size,
      sequence_length, config.shared_representation_size)` -- The self-attention key state for use in the next
      step of incremental decoding
    - **self_value** (*optional*, returned when `use_cache=True`) `torch.FloatTensor` of shape `(batch_size,
      sequence_length, config.hidden_size)` -- The self-attention value state for use in the next step of
      incremental decoding
    - **self_ema_state** (*optional*, returned when `use_cache=True`) `torch.FloatTensor` of shape
      `(batch_size, config.ndim)` The incremental EMA state for use in the next step of incremental decoding.


class Softmax(Layer):

    def __init__(self, axis=-1, **kwargs):
        super(Softmax, self).__init__(**kwargs)
        self.supports_masking = True
        self.axis = axis

    def call(self, inputs, mask=None):
        if mask is not None:
            adder = (1.0 - math_ops.cast(mask, inputs.dtype)) * _large_compatible_negative(inputs.dtype)
            inputs += adder
        if isinstance(self.axis, (tuple, list)):
            if len(self.axis) > 1:
                return math_ops.exp(inputs - math_ops.reduce_logsumexp(inputs, axis=self.axis, keepdims=True))
            else:
                return backend.softmax(inputs, axis=self.axis[0])
        return backend.softmax(inputs, axis=self.axis)

    def get_config(self):
        config = {'axis': self.axis}
        base_config = super(Softmax, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

    @tf_utils.shape_type_conversion
    def compute_output_shape(self, input_shape):
        return input_shape

Softmax activation function.

Example without mask:

>>> inp = np.asarray([1., 2., 1.])
>>> layer = tf.keras.layers.Softmax()
>>> layer(inp).numpy()
array([0.21194157, 0.5761169 , 0.21194157], dtype=float32)
>>> mask = np.asarray([True, False, True], dtype=bool)
>>> layer(inp, mask).numpy()
array([0.5, 0. , 0.5], dtype=float32)

Input shape:
  Arbitrary. Use the keyword argument `input_shape`
  (tuple of integers, does not include the samples axis)
  when using this layer as the first layer in a model.

Output shape:
  Same shape as the input.

Args:
  axis: Integer, or list of Integers, axis along which the softmax
    normalization is applied.
Call arguments:
  inputs: The inputs, or logits to the softmax layer.
  mask: A boolean mask of the same shape as `inputs`. Defaults to `None`. The
    mask specifies 1 to keep and 0 to mask.

Returns:
  softmaxed output with the same shape as `inputs`.


def batch_reduce_and_verify(self, inputs, expect, options):

    def replica_fn():
        CollectiveReplicaLauncher._prefer_unique_instance_key = options.prefer_unique_instance_key
        collective, devices, pid = self.make_collective(options.num_processes, options.gpus_per_process)

        def batch_reduce_fn():
            batch_size = len(inputs[0])
            value_dst_pairs = []
            for i in range(batch_size):

                def value_fn(device_idx, idx=i):
                    return inputs[pid * len(devices) + device_idx][idx]
                per_replica_value = make_per_replica_value(value_fn, devices)
                value_dst_pairs.append((per_replica_value, per_replica_value))
            reduced_values = collective.batch_reduce(options.reduce_op, value_dst_pairs, options.communication_options)
            if options.gpus_per_process > 1:
                for v in reduced_values:
                    self.assertIsInstance(v, value_lib.Mirrored)
            reduced_values = [self.as_list(v) for v in reduced_values]
            for v in reduced_values:
                self.assertAllEqual(devices, [t.device for t in v])
            return nest.map_structure(ops.convert_to_tensor, reduced_values)
        per_replica_expect = nest.map_structure(lambda x: [ops.convert_to_tensor(x)] * len(devices), expect)
        if 'eager' in options.mode:
            got = batch_reduce_fn()
            self.assertAllClose(got, per_replica_expect)
        if 'func_graph' in options.mode:
            got = def_function.function(batch_reduce_fn)()
            self.assertAllClose(got, per_replica_expect)
    get_global_mpr(options.num_processes).run(replica_fn)

Batch reduce the given `inputs` and verify the output matches `expect`.

Args:
  inputs: a 2-level nested list of `Tensor` or `IndexedSlices`, where i-th
    value will be fed to i-th replica.
  expect: a list of `Tensor` or `IndexedSlices`. This should be the expected
    value for one replica.
  options: a `RunOpotions` instance.


class DetrModelOutput(Seq2SeqModelOutput):
    intermediate_hidden_states: Optional[torch.FloatTensor] = None

Base class for outputs of the DETR encoder-decoder model. This class adds one attribute to Seq2SeqModelOutput,
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 decoder of the model.
    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.
    intermediate_hidden_states (`torch.FloatTensor` of shape `(config.decoder_layers, batch_size, sequence_length, 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.


def _tag_and_add_meta_graph(self, meta_graph_def, tags, signature_def_map):
    for tag in tags:
        meta_graph_def.meta_info_def.tags.append(tag)
    if signature_def_map is not None:
        for key in signature_def_map:
            meta_graph_def.signature_def[key].CopyFrom(signature_def_map[key])
    proto_meta_graph_def = self._saved_model.meta_graphs.add()
    proto_meta_graph_def.CopyFrom(meta_graph_def)

Tags the meta graph def and adds it to the SavedModel.

Tags the meta graph def with the supplied tags, adds signature defs to it if
provided and appends the meta graph def to the SavedModel proto.

Args:
  meta_graph_def: The meta graph def to add to the SavedModel.
  tags: The set of tags to annotate the meta graph def with.
  signature_def_map: The map of signature defs to be added to the meta graph
    def.


Called at the end of session.

The `session` argument can be used in case the hook wants to run final ops,
such as saving a last checkpoint.

If `session.run()` raises exception other than OutOfRangeError or
StopIteration then `end()` is not called.
Note the difference between `end()` and `after_run()` behavior when
`session.run()` raises OutOfRangeError or StopIteration. In that case
`end()` is called but `after_run()` is not called.

Args:
  session: A TensorFlow Session that will be soon closed.


def input(self):
    if not self._inbound_nodes:
        raise AttributeError('Layer ' + self.name + ' is not connected, no input to return.')
    return self._get_node_attribute_at_index(0, 'input_tensors', 'input')

Retrieves the input tensor(s) of a layer.

Only applicable if the layer has exactly one input,
i.e. if it is connected to one incoming layer.

Returns:
    Input tensor or list of input tensors.

Raises:
  RuntimeError: If called in Eager mode.
  AttributeError: If no inbound nodes are found.


def connect(self, uid=UNKNOWN_UID, cmd=JsonRpcCommand.INIT):
    self._counter = self._id_counter()
    try:
        self._conn = socket.create_connection(('localhost', self.host_port), _SOCKET_CONNECTION_TIMEOUT)
    except ConnectionRefusedError as err:
        self.log.debug('Failed to connect to localhost, trying 127.0.0.1: {}'.format(str(err)))
        self._conn = socket.create_connection(('127.0.0.1', self.host_port), _SOCKET_CONNECTION_TIMEOUT)
    self._conn.settimeout(_SOCKET_READ_TIMEOUT)
    self._client = self._conn.makefile(mode='brw')
    resp = self._cmd(cmd, uid)
    if not resp:
        raise ProtocolError(self._ad, ProtocolError.NO_RESPONSE_FROM_HANDSHAKE)
    result = json.loads(str(resp, encoding='utf8'))
    if result['status']:
        self.uid = result['uid']
    else:
        self.uid = UNKNOWN_UID

Opens a connection to a JSON RPC server.

Opens a connection to a remote client. The connection attempt will time
out if it takes longer than _SOCKET_CONNECTION_TIMEOUT seconds. Each
subsequent operation over this socket will time out after
_SOCKET_READ_TIMEOUT seconds as well.

Args:
  uid: int, The uid of the session to join, or UNKNOWN_UID to start a
    new session.
  cmd: JsonRpcCommand, The command to use for creating the connection.

Raises:
  IOError: Raised when the socket times out from io error
  socket.timeout: Raised when the socket waits to long for connection.
  ProtocolError: Raised when there is an error in the protocol.


def call(self, context: tf.Tensor, latents: tf.Tensor) -> tf.Tensor:
    context = self.context_layer_norm(context)
    latents = self.latents_layer_norm(latents)
    batch_size, seq_length, embed_dim = shape_list(context)
    q = self.q_proj(latents)
    k = self.k_proj(tf.concat([context, latents], axis=-2))
    v = self.v_proj(tf.concat([context, latents], axis=-2))
    q, k, v = [tf.transpose(tf.reshape(x, (batch_size, x.shape[1], self.n_heads, self.head_dim)), perm=[0, 2, 1, 3]) for x in (q, k, v)]
    if self.qk_layer_norms:
        q = self.q_layer_norm(q)
        k = self.k_layer_norm(k)
    scores = tf.einsum('... i d, ... j d -> ... i j', q * self.qk_scale, k)
    stabilized_scores = scores - tf.reduce_max(scores, axis=-1, keepdims=True)
    attn = tf.nn.softmax(stabilized_scores, axis=-1)
    resampled = tf.einsum('... i j, ... j d -> ... i d', attn, v)
    return self.output_proj(tf.reshape(tf.transpose(resampled, perm=[0, 2, 1, 3]), (batch_size, -1, self.n_heads * self.head_dim)))

Runs Perceiver Self-Attention, with special (context, latents) appended along the `seq` dimension!

Args:
    context (`tf.Tensor`):
        Tensor of shape `[bsz, seq, embed_dim]` representing long-form context to resample.
    latents (`tf.Tensor`):
        Tensor of shape `[bsz, n_latents, embed_dim]` representing fixed length latents to compress to.

Returns:
    `tf.Tensor`: Tensor of shape `[bsz, n_latents, embed_dim]` representing attention over latents w/ cross
    from context.


def set_record_attn(self, record_attn):

    def _should_record_attn(layer_idx):
        if isinstance(record_attn, bool):
            return record_attn
        return layer_idx in record_attn
    for i, layer in enumerate(self._attn_mods):
        layer.attn.record_attn = _should_record_attn(i)
    if not record_attn:
        self.saved_attn_weights = []

Makes forward prop dump self-attention softmaxes to self.saved_attn_weights.

Args:
    record_attn (`Union[bool,set]`):
        Either a set of layer indices indicating which layers to store, or a boolean value indicating Whether
        to dump all.


def set(self, value):
    pywrap_tfe.TFE_MonitoringIntGaugeCellSet(self._cell, value)

Atomically set the value.

Args:
  value: integer value.


def _PushParameterListState(self, newline):
    current = self.next_token
    previous = current.previous_token
    if _IsFunctionDefinition(previous):
        first_param_column = previous.total_length + self.stack[-2].indent
        self.param_list_stack.append(object_state.ParameterListState(previous, newline, first_param_column))

Push a new parameter list state for a function definition.

Args:
  newline: Whether the current token is to be added on a newline.


def checksum(self, path):
    raise NotImplementedError

Fetch checksum metadata of a file on the
:class:`~apache_beam.io.filesystem.FileSystem`.

This operation returns checksum metadata as stored in the underlying
FileSystem. It should not need to read file data to obtain this value.
Checksum type and format are FileSystem dependent and are not compatible
between FileSystems.
FileSystem implementations may return file size if a checksum isn't
available.

Args:
  path: string path of a file.

Returns: string containing checksum

Raises:
  ``BeamIOError``: if path isn't a file or doesn't exist.


def register_for_auto_class(cls, auto_class='AutoConfig'):
    if not isinstance(auto_class, str):
        auto_class = auto_class.__name__
    import transformers.models.auto as auto_module
    if not hasattr(auto_module, auto_class):
        raise ValueError(f'{auto_class} is not a valid auto class.')
    cls._auto_class = auto_class

Register this class with a given auto class. This should only be used for custom configurations as the ones in
the library are already mapped with `AutoConfig`.

Args:
    auto_class (`str` or `type`, *optional*, defaults to `"AutoConfig"`):
        The auto class to register this new configuration with.


def convert_format(boxes, source, target, height=None, width=None, dtype='float32'):
    box_utils = BoundingBox()
    if backend_utils.in_tf_graph():
        box_utils.backend.set_backend('tensorflow')
    boxes = box_utils.convert_format(boxes=boxes, source=source, target=target, height=height, width=width, dtype=dtype)
    box_utils.backend.reset()
    return boxes

Converts bounding boxes between formats.

Supported formats (case-insensitive):
`"xyxy"`: [left, top, right, bottom]
`"yxyx"`: [top, left, bottom, right]
`"xywh"`: [left, top, width, height]
`"center_xywh"`: [center_x, center_y, width, height]
`"center_yxhw"`: [center_y, center_x, height, width]
`"rel_xyxy"`, `"rel_yxyx"`, `"rel_xywh"`, `"rel_center_xywh"`:  Relative
    versions of the above formats, where coordinates are normalized
    to the range [0, 1] based on the image `height` and `width`.

Args:
    boxes: Bounding boxes tensor/array or dictionary of `boxes` and
        `labels`.
    source: Source format string.
    target: Target format string.
    height: Image height (required for relative target format).
    width: Image width (required for relative target format).
    dtype: Data type for conversion (optional).

Returns:
    Converted boxes.

Raises:
    ValueError: For invalid formats, shapes, or missing dimensions.

Example:
```python
boxes = np.array([[10, 20, 30, 40], [50, 60, 70, 80]])
# Convert from 'xyxy' to 'xywh' format
boxes_xywh = keras.utils.bounding_boxes.convert_format(
    boxes, source='xyxy', target='xywh'
)  # Output: [[10. 20. 20. 20.], [50. 60. 20. 20.]]

# Convert to relative 'rel_xyxy' format
boxes_rel_xyxy = keras.utils.bounding_boxes.convert_format(
    boxes, source='xyxy', target='rel_xyxy', height=200, width=300
) # Output: [[0.03333334 0.1        0.1        0.2       ],
           #[0.16666667 0.3        0.23333333 0.4       ]]
```